summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--caffe2/operators/conv_transpose_op_cudnn.cc261
-rw-r--r--caffe2/operators/conv_transpose_op_impl.h889
-rw-r--r--caffe2/operators/conv_transpose_unpool_op_base.h38
-rw-r--r--caffe2/python/operator_test/conv_transpose_test.py63
4 files changed, 765 insertions, 486 deletions
diff --git a/caffe2/operators/conv_transpose_op_cudnn.cc b/caffe2/operators/conv_transpose_op_cudnn.cc
index 8f8c9a2db8..459ccd7661 100644
--- a/caffe2/operators/conv_transpose_op_cudnn.cc
+++ b/caffe2/operators/conv_transpose_op_cudnn.cc
@@ -1,7 +1,10 @@
+#include "caffe2/operators/conv_transpose_op.h"
+
+#include <vector>
+
#include "caffe2/core/context_gpu.h"
#include "caffe2/core/cudnn_wrappers.h"
#include "caffe2/operators/conv_op_cache_cudnn.h"
-#include "caffe2/operators/conv_transpose_op.h"
#include "caffe2/operators/op_utils_cudnn.h"
namespace caffe2 {
@@ -49,6 +52,7 @@ class CudnnConvTransposeOpBase : public ConvTransposeUnpoolBase<CUDAContext> {
CUDNN_ENFORCE(cudnnCreateFilterDescriptor(&filter_desc_));
if (InputSize() == 3) {
CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&bias_desc_));
+ CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&top_desc_for_bias_));
}
CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&top_desc_));
CUDNN_ENFORCE(cudnnCreateConvolutionDescriptor(&conv_desc_));
@@ -59,27 +63,59 @@ class CudnnConvTransposeOpBase : public ConvTransposeUnpoolBase<CUDAContext> {
CUDNN_ENFORCE(cudnnDestroyFilterDescriptor(filter_desc_));
if (InputSize() == 3) {
CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(bias_desc_));
+ CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(top_desc_for_bias_));
}
CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(top_desc_));
CUDNN_ENFORCE(cudnnDestroyConvolutionDescriptor(conv_desc_));
}
protected:
- vector<int64_t> cudnn_input_dims_;
- vector<int64_t> cudnn_filter_dims_;
+ void SetTensor4DDescriptorWithGroup(
+ const cudnnDataType_t data_type,
+ const int N,
+ const int C,
+ const int H,
+ const int W,
+ cudnnTensorDescriptor_t* desc) const {
+#if CUDNN_VERSION_MIN(7, 0, 0)
+ const int CC = C;
+#else
+ const int CC = C / group_;
+#endif
+ switch (order_) {
+ case StorageOrder::NCHW: {
+ CUDNN_ENFORCE(cudnnSetTensor4dDescriptorEx(
+ *desc, data_type, N, CC, H, W, C * H * W, H * W, W, 1));
+ break;
+ }
+ case StorageOrder::NHWC: {
+ CUDNN_ENFORCE(cudnnSetTensor4dDescriptorEx(
+ *desc, data_type, N, CC, H, W, H * W * C, 1, W * C, C));
+ break;
+ }
+ default: {
+ LOG(FATAL) << "Unknown storage order: " << order_;
+ }
+ }
+ }
+
+ std::vector<std::int64_t> cudnn_input_dims_;
+ std::vector<std::int64_t> cudnn_filter_dims_;
CuDNNWrapper cudnn_wrapper_;
cudnnTensorDescriptor_t bottom_desc_;
cudnnFilterDescriptor_t filter_desc_;
cudnnTensorDescriptor_t bias_desc_;
cudnnTensorDescriptor_t top_desc_;
+ cudnnTensorDescriptor_t top_desc_for_bias_;
cudnnConvolutionDescriptor_t conv_desc_;
+
const size_t cudnn_ws_nbytes_limit_;
size_t cudnn_ws_nbytes_;
bool exhaustive_search_;
bool deterministic_;
size_t cudnn_state_;
- vector<int> force_algo_; // stored as FWD, dFILTER, dDATA
+ std::vector<int> force_algo_; // stored as FWD, dFILTER, dDATA
bool enable_tensor_core_;
};
@@ -141,10 +177,10 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
int C = 0;
switch (order_) {
case StorageOrder::NHWC:
- C = filter.dim32(3);
+ C = filter.dim32(3) * group_;
break;
case StorageOrder::NCHW:
- C = filter.dim32(1);
+ C = filter.dim32(1) * group_;
break;
default:
LOG(FATAL) << "Unknown storage order: " << order_;
@@ -162,9 +198,8 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
H_out = Y->dim32(1);
W_out = Y->dim32(2);
CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
- CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_w());
- CAFFE_ENFORCE_EQ(filter.dim32(3), C);
+ CAFFE_ENFORCE_EQ(filter.dim32(3), C / group_);
break;
case StorageOrder::NCHW:
N = X.dim32(0);
@@ -173,13 +208,14 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
W = X.dim32(3);
H_out = Y->dim32(2);
W_out = Y->dim32(3);
- CAFFE_ENFORCE_EQ(filter.dim32(1), C);
+ CAFFE_ENFORCE_EQ(filter.dim32(1), C / group_);
CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_h());
CAFFE_ENFORCE_EQ(filter.dim32(3), kernel_w());
break;
default:
LOG(FATAL) << "Unknown storage order: " << order_;
}
+ CAFFE_ENFORCE_EQ(M % group_, 0);
if (InputSize() == 3) {
auto& bias = Input(BIAS);
@@ -188,30 +224,29 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
}
// Set up the cudnn algorithms & workspace if necessary
- bool input_changed = (X.sizes() != cudnn_input_dims_);
- bool filter_changed = (filter.sizes() != cudnn_filter_dims_);
+ const bool input_changed = (X.sizes() != cudnn_input_dims_);
+ const bool filter_changed = (filter.sizes() != cudnn_filter_dims_);
if (input_changed || filter_changed) {
VLOG(1) << "Changing the cudnn descriptor configurations.";
if (input_changed) {
cudnn_input_dims_ = X.sizes().vec();
- CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
- bottom_desc_,
- GetCudnnTensorFormat(order_),
- cudnnTypeWrapper<T>::type,
- N,
- M,
- H,
- W));
+ SetTensor4DDescriptorWithGroup(
+ cudnnTypeWrapper<T>::type, N, M, H, W, &bottom_desc_);
}
if (filter_changed) {
cudnn_filter_dims_ = filter.sizes().vec();
+#if CUDNN_VERSION_MIN(7, 0, 0)
+ const int MM = M;
+#else
+ const int MM = M / group_;
+#endif
CUDNN_ENFORCE(cudnnSetFilter4dDescriptor(
filter_desc_,
cudnnTypeWrapper<T>::type,
GetCudnnTensorFormat(order_),
- M,
- C,
+ MM,
+ C / group_,
kernel_h(),
kernel_w()));
if (InputSize() == 3) {
@@ -226,14 +261,19 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
}
}
// Set the output
- CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
- top_desc_,
- GetCudnnTensorFormat(order_),
- cudnnTypeWrapper<T>::type,
- N,
- C,
- H_out,
- W_out));
+ SetTensor4DDescriptorWithGroup(
+ cudnnTypeWrapper<T>::type, N, C, H_out, W_out, &top_desc_);
+ if (InputSize() == 3) {
+ CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
+ top_desc_for_bias_,
+ GetCudnnTensorFormat(order_),
+ cudnnTypeWrapper<T>::type,
+ N,
+ C,
+ H_out,
+ W_out));
+ }
+
// Set the convolution descriptor
CAFFE_ENFORCE_EQ(
pad_t(),
@@ -246,7 +286,7 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
"The current padding scheme leads to unequal padding on the left "
"and right, which is not supported by cudnn.");
// Set the convolution descriptor
-#if CUDNN_VERSION_MIN(6,0,0)
+#if CUDNN_VERSION_MIN(6, 0, 0)
CUDNN_ENFORCE(cudnnSetConvolution2dDescriptor(
conv_desc_,
pad_t(),
@@ -268,6 +308,7 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
1,
CUDNN_CROSS_CORRELATION));
#endif
+
#if CUDNN_VERSION_MIN(7, 0, 0)
// enable TensorCore math if desired
enable_tensor_core_ &= TensorCoreAvailable();
@@ -275,7 +316,10 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
CUDNN_ENFORCE(
cudnnSetConvolutionMathType(conv_desc_, CUDNN_TENSOR_OP_MATH));
}
+ // set cuDNN groups if appropriate
+ CUDNN_ENFORCE(cudnnSetConvolutionGroupCount(conv_desc_, group_));
#endif
+
if (force_algo_[ALGO_DGRAD] >= 0) {
bwd_data_algo_ = (cudnnConvolutionBwdDataAlgo_t)force_algo_[ALGO_DGRAD];
} else if (deterministic_) {
@@ -331,24 +375,56 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
VLOG(1) << "CuDNN workspace size: " << bwd_data_ws_size;
}
+ const T* X_data = X.template data<T>();
+ const T* filter_data = filter.template data<T>();
+ T* Y_data = Y->template mutable_data<T>();
+
// Now, actually run the computation.
// Filter
+#if CUDNN_VERSION_MIN(7, 0, 0)
cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
CUDNN_ENFORCE(cudnnConvolutionBackwardData(
state->cudnn_handle(),
cudnnTypeWrapper<T>::kOne(),
filter_desc_,
- filter.template data<T>(),
+ filter_data,
bottom_desc_,
- X.template data<T>(),
+ X_data,
conv_desc_,
bwd_data_algo_,
state->workspace().get(cudnn_ws_nbytes_),
cudnn_ws_nbytes_,
cudnnTypeWrapper<T>::kZero(),
top_desc_,
- Y->template mutable_data<T>()));
+ Y_data));
});
+#else
+ const int X_HxW = H * W;
+ const int Y_HxW = H_out * W_out;
+ const int group_offset_X =
+ order_ == StorageOrder::NCHW ? M / group_ * X_HxW : M / group_;
+ const int group_offset_Y =
+ order_ == StorageOrder::NCHW ? C / group_ * Y_HxW : C / group_;
+ const int group_offset_filter = filter.numel() / group_;
+ for (int i = 0; i < group_; ++i) {
+ cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
+ CUDNN_ENFORCE(
+ cudnnConvolutionBackwardData(state->cudnn_handle(),
+ cudnnTypeWrapper<T>::kOne(),
+ filter_desc_,
+ filter_data + i * group_offset_filter,
+ bottom_desc_,
+ X_data + i * group_offset_X;
+ conv_desc_,
+ bwd_data_algo_,
+ state->workspace().get(cudnn_ws_nbytes_),
+ cudnn_ws_nbytes_,
+ cudnnTypeWrapper<T_DX>::kZero(),
+ top_desc_,
+ Y_data + i * group_offset_Y));
+ });
+ }
+#endif
// Bias
if (InputSize() == 3) {
CUDNN_ENFORCE(cudnnAddTensor(
@@ -357,7 +433,7 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
bias_desc_,
Input(BIAS).template data<T>(),
cudnnTypeWrapper<T>::kOne(),
- top_desc_,
+ top_desc_for_bias_,
Y->template mutable_data<T>()));
}
// Done.
@@ -368,19 +444,19 @@ bool CudnnConvTransposeOp<T>::RunOnDevice() {
// consolidating them.
template <typename T>
bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
- auto& X = Input(INPUT);
- auto& filter = Input(FILTER);
- auto& dY = Input(OUTPUT_GRAD);
+ const auto& X = Input(INPUT);
+ const auto& filter = Input(FILTER);
+ const auto& dY = Input(OUTPUT_GRAD);
CAFFE_ENFORCE_EQ(X.dim(), 4);
CAFFE_ENFORCE_EQ(filter.dim(), 4);
int C = 0;
switch (order_) {
case StorageOrder::NHWC:
- C = filter.dim32(3);
+ C = filter.dim32(3) * group_;
break;
case StorageOrder::NCHW:
- C = filter.dim32(1);
+ C = filter.dim32(1) * group_;
break;
default:
LOG(FATAL) << "Unknown storage order: " << order_;
@@ -398,7 +474,7 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
CAFFE_ENFORCE_EQ(filter.dim32(1), kernel_h());
CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_w());
- CAFFE_ENFORCE_EQ(filter.dim32(3), C);
+ CAFFE_ENFORCE_EQ(filter.dim32(3), C / group_);
break;
case StorageOrder::NCHW:
N = X.dim32(0);
@@ -407,41 +483,42 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
W = X.dim32(3);
H_out = dY.dim32(2);
W_out = dY.dim32(3);
- CAFFE_ENFORCE_EQ(filter.dim32(1), C);
+ CAFFE_ENFORCE_EQ(filter.dim32(1), C / group_);
CAFFE_ENFORCE_EQ(filter.dim32(2), kernel_h());
CAFFE_ENFORCE_EQ(filter.dim32(3), kernel_w());
break;
default:
LOG(FATAL) << "Unknown storage order: " << order_;
}
+ CAFFE_ENFORCE_EQ(M % group_, 0);
+
// Since we only handle LegacyPadding::NOTSET, we don't need to
// compute padding.
auto* dfilter = Output(FILTER_GRAD, filter.sizes(), at::dtype<T>());
// Set up the cudnn algorithms & workspace if necessary
- bool input_changed = (X.sizes() != cudnn_input_dims_);
- bool filter_changed = (filter.sizes() != cudnn_filter_dims_);
+ const bool input_changed = (X.sizes() != cudnn_input_dims_);
+ const bool filter_changed = (filter.sizes() != cudnn_filter_dims_);
if (input_changed || filter_changed) {
VLOG(1) << "Changing the cudnn descriptor configurations.";
if (input_changed) {
cudnn_input_dims_ = X.sizes().vec();
- CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
- bottom_desc_,
- GetCudnnTensorFormat(order_),
- cudnnTypeWrapper<T>::type,
- N,
- M,
- H,
- W));
+ SetTensor4DDescriptorWithGroup(
+ cudnnTypeWrapper<T>::type, N, M, H, W, &bottom_desc_);
}
if (filter_changed) {
cudnn_filter_dims_ = filter.sizes().vec();
+#if CUDNN_VERSION_MIN(7, 0, 0)
+ const int MM = M;
+#else
+ const int MM = M / group_;
+#endif
CUDNN_ENFORCE(cudnnSetFilter4dDescriptor(
filter_desc_,
cudnnTypeWrapper<T>::type,
GetCudnnTensorFormat(order_),
- M,
- C,
+ MM,
+ C / group_,
kernel_h(),
kernel_w()));
if (!no_bias_) {
@@ -456,14 +533,19 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
}
}
// Set the output
- CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
- top_desc_,
- GetCudnnTensorFormat(order_),
- cudnnTypeWrapper<T>::type,
- N,
- C,
- H_out,
- W_out));
+ SetTensor4DDescriptorWithGroup(
+ cudnnTypeWrapper<T>::type, N, C, H_out, W_out, &top_desc_);
+ if (!no_bias_) {
+ CUDNN_ENFORCE(cudnnSetTensor4dDescriptor(
+ top_desc_for_bias_,
+ GetCudnnTensorFormat(order_),
+ cudnnTypeWrapper<T>::type,
+ N,
+ C,
+ H_out,
+ W_out));
+ }
+
// Set the convolution descriptor
CAFFE_ENFORCE_EQ(
pad_t(),
@@ -475,7 +557,7 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
pad_r(),
"The current padding scheme leads to unequal padding on the left "
"and right, which is not supported by cudnn.");
-#if CUDNN_VERSION_MIN(6,0,0)
+#if CUDNN_VERSION_MIN(6, 0, 0)
CUDNN_ENFORCE(cudnnSetConvolution2dDescriptor(
conv_desc_,
pad_t(),
@@ -504,6 +586,8 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
CUDNN_ENFORCE(
cudnnSetConvolutionMathType(conv_desc_, CUDNN_TENSOR_OP_MATH));
}
+ // set cuDNN groups if appropriate
+ CUDNN_CHECK(cudnnSetConvolutionGroupCount(conv_desc_, group_));
#endif
if (force_algo_[ALGO_WGRAD] >= 0) {
bwd_filter_algo_ =
@@ -622,13 +706,14 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
CUDNN_ENFORCE(cudnnConvolutionBackwardBias(
cudnn_wrapper_.inline_cudnn_handle(),
cudnnTypeWrapper<T>::kOne(),
- top_desc_,
+ top_desc_for_bias_,
dY.template data<T>(),
cudnnTypeWrapper<T>::kZero(),
bias_desc_,
dbias->template mutable_data<T>()));
}
+#if CUDNN_VERSION_MIN(7, 0, 0)
cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
CUDNN_ENFORCE(cudnnConvolutionBackwardFilter(
state->cudnn_handle(),
@@ -647,7 +732,6 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
if (OutputSize() == 3 || (no_bias_ && (OutputSize() == 2))) {
// Compute the gradient w.r.t. the input.
-
auto* dX = Output(
no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD,
X.sizes(),
@@ -668,6 +752,55 @@ bool CudnnConvTransposeGradientOp<T>::RunOnDevice() {
dX->template mutable_data<T>()));
}
});
+#else
+ const int X_HxW = H * W;
+ const int Y_HxW = H_out * W_out;
+ const int group_offset_X =
+ order_ == StorageOrder::NCHW ? M / group_ * X_HxW : M / group_;
+ const int group_offset_Y =
+ order_ == StorageOrder::NCHW ? C / group_ * Y_HxW : C / group_;
+ const int group_offset_filter = filter.numel() / group_;
+ for (int i = 0; i < group_; ++i) {
+ cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
+ CUDNN_ENFORCE(cudnnConvolutionBackwardFilter(
+ state->cudnn_handle(),
+ cudnnTypeWrapper<T>::kOne(),
+ top_desc_,
+ dY.template data<T>() + i * group_offset_Y,
+ bottom_desc_,
+ X.template data<T>() + i * group_offset_X,
+ conv_desc_,
+ bwd_filter_algo_,
+ state->workspace().get(cudnn_ws_nbytes_),
+ cudnn_ws_nbytes_,
+ cudnnTypeWrapper<T>::kZero(),
+ filter_desc_,
+ dfilter->template mutable_data<T>() + i * group_offset_filter));
+ if (OutputSize() == 3 || (no_bias_ && (OutputSize() == 2))) {
+ // Compute the gradient w.r.t. the input.
+ auto* dX = Output(
+ no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD,
+ X.sizes(),
+ at::dtype<T>());
+ cudnn_wrapper_.with_cudnn_state(cudnn_state_, [&](CuDNNState* state) {
+ CUDNN_ENFORCE(cudnnConvolutionForward(
+ state->cudnn_handle(),
+ cudnnTypeWrapper<T>::kOne(),
+ top_desc_,
+ dY.template data<T>() + i * group_offset_Y,
+ filter_desc_,
+ filter.template data<T>() + i * group_offset_filter,
+ conv_desc_,
+ algo_,
+ state->workspace().get(cudnn_ws_nbytes_),
+ cudnn_ws_nbytes_,
+ cudnnTypeWrapper<T>::kZero(),
+ bottom_desc_,
+ dX->template mutable_data<T>() + i * group_offset_X));
+ });
+ }
+ }
+#endif
return true;
}
diff --git a/caffe2/operators/conv_transpose_op_impl.h b/caffe2/operators/conv_transpose_op_impl.h
index 41af81c2c6..333f782920 100644
--- a/caffe2/operators/conv_transpose_op_impl.h
+++ b/caffe2/operators/conv_transpose_op_impl.h
@@ -3,11 +3,15 @@
#ifndef CAFFE2_OPERATORS_CONV_TRANSPOSE_OP_IMPL_H_
#define CAFFE2_OPERATORS_CONV_TRANSPOSE_OP_IMPL_H_
+#include "caffe2/operators/conv_transpose_op.h"
+
+#include <array>
+#include <vector>
+
#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/operators/conv_op_shared.h"
-#include "caffe2/operators/conv_transpose_op.h"
#include "caffe2/operators/conv_transpose_unpool_op_base.h"
#include "caffe2/utils/math.h"
@@ -17,551 +21,618 @@ namespace caffe2 {
template <typename T, class Context>
bool ConvTransposeOp<T, Context>::RunOnDeviceWithOrderNCHW() {
- const Tensor& X = Input(INPUT);
- auto& filter = Input(FILTER);
- const int N = X.dim32(0), M = X.dim32(1), H = X.dim32(2), W = X.dim32(3);
- CAFFE_ENFORCE(filter.dim() == 4, "filter must be 4D tensor");
- CAFFE_ENFORCE(
- filter.dim32(0) == M,
- "filter number must be equal to input channel number");
- const int C = filter.dim32(1);
- CAFFE_ENFORCE(
- filter.dim32(2) == this->kernel_h(),
+ const auto& X = Input(INPUT);
+ const auto& filter = Input(FILTER);
+ CAFFE_ENFORCE_EQ(X.dim(), 4, "Input must be 4D tensor");
+ CAFFE_ENFORCE_EQ(filter.dim(), 4, "filter must be 4D tensor");
+ const int N = X.dim32(0);
+ const int M = X.dim32(1);
+ const int H = X.dim32(2);
+ const int W = X.dim32(3);
+ const int G = group_;
+ CAFFE_ENFORCE_EQ(M, filter.dim32(0));
+ CAFFE_ENFORCE_EQ(
+ M % G, 0, "The number of input channels is not divisible by group.");
+ const int C = filter.dim32(1) * G;
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(2),
+ kernel_h(),
"filter height must be equal to kernel height");
- CAFFE_ENFORCE(
- filter.dim32(3) == this->kernel_w(),
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(3),
+ this->kernel_w(),
"filter width must be equal to kernel width");
- auto sizes = ConvTransposeUnpoolBase<Context>::GetOutputSize(X, C);
- Tensor* Y = Output(0, sizes, at::dtype<T>());
+ const std::vector<std::int64_t> Y_dims =
+ ConvTransposeUnpoolBase<Context>::GetOutputSize(X, C);
+ auto* Y = Output(0, Y_dims, at::dtype<T>());
- const int kernel_dim = C * this->kernel_h() * this->kernel_w();
- const int input_image_size = H * W;
- const int output_image_size = Y->dim32(2) * Y->dim32(3);
+ if (N == 0) {
+ return true;
+ }
+ const int K_HxW = kernel_h() * kernel_w();
+ const int kernel_dim = C / G * K_HxW;
+ const int X_HxW = H * W;
+ const int Y_HxW = Y->dim32(2) * Y->dim32(3);
+
+ const T* X_data = X.template data<T>();
+ const T* filter_data = filter.template data<T>();
+ const T* bias_data = nullptr;
if (InputSize() == 3) {
auto& bias = Input(BIAS);
- CAFFE_ENFORCE(bias.dim() == 1, "bias must be 1D tensor");
- CAFFE_ENFORCE(
- bias.dim32(0) == C,
+ CAFFE_ENFORCE_EQ(bias.dim(), 1, "bias must be 1D tensor");
+ CAFFE_ENFORCE_EQ(
+ bias.dim32(0),
+ C,
"bias dimension must be equal to output channel number");
- ReinitializeTensor(
- &bias_multiplier_,
- {1, output_image_size},
- at::dtype<T>().device(Context::GetDeviceType()));
- T* bm_data = bias_multiplier_.template mutable_data<T>();
- math::Set<T, Context>(
- output_image_size,
- static_cast<T>(1),
- bm_data,
- &context_);
+ bias_data = bias.template data<T>();
}
+ T* Y_data = Y->template mutable_data<T>();
- const T* Xdata = X.template data<T>();
- const T* filter_data = filter.template data<T>();
- T* Ydata = Y->template mutable_data<T>();
+ const std::vector<std::int64_t> buffer_shape = {
+ C, kernel_h(), kernel_w(), H, W};
- auto f = [&](Tensor* col_buffer) {
- ReinitializeTensor(col_buffer, vector<int64_t>{C, this->kernel_h(), this->kernel_w(), H, W}, at::dtype<T>().device(Context::GetDeviceType()));
+ const auto func = [&](Tensor* col_buffer) {
+ ReinitializeTensor(
+ col_buffer,
+ buffer_shape,
+ at::dtype<T>().device(Context::GetDeviceType()));
T* col_buffer_data = col_buffer->template mutable_data<T>();
- for (auto image_id = 0; image_id < N; ++image_id) {
+ for (int image_id = 0; image_id < N; ++image_id) {
// Weight term
- math::Gemm<T, Context>(
- CblasTrans,
- CblasNoTrans,
- kernel_dim,
- input_image_size,
- M,
- 1,
- filter_data,
- Xdata,
- 0,
- col_buffer_data,
- &context_);
+ if (G == 1) {
+ math::Gemm<T, Context>(
+ CblasTrans,
+ CblasNoTrans,
+ kernel_dim,
+ X_HxW,
+ M,
+ 1.0f,
+ filter_data,
+ X_data + image_id * M * X_HxW,
+ 0.0f,
+ col_buffer_data,
+ &context_);
+ } else {
+ math::GemmStridedBatched<T, Context>(
+ CblasTrans,
+ CblasNoTrans,
+ G,
+ kernel_dim,
+ X_HxW,
+ M / G,
+ 1.0f,
+ filter_data,
+ M / G * kernel_dim,
+ X_data + image_id * M * X_HxW,
+ M / G * X_HxW,
+ 0.0f,
+ col_buffer_data,
+ col_buffer->numel() / G,
+ &context_);
+ }
// Col2Im
math::Col2Im<T, Context, StorageOrder::NCHW>(
C,
Y->dim32(2),
Y->dim32(3),
- this->kernel_h(),
- this->kernel_w(),
+ kernel_h(),
+ kernel_w(),
1,
1,
- this->pad_t(),
- this->pad_l(),
- this->pad_b(),
- this->pad_r(),
- this->stride_h(),
- this->stride_w(),
+ pad_t(),
+ pad_l(),
+ pad_b(),
+ pad_r(),
+ stride_h(),
+ stride_w(),
col_buffer_data,
- Ydata,
+ Y_data + image_id * C * Y_HxW,
&context_);
- // Bias term
- if (InputSize() == 3) {
- const T* bias_data = Input(BIAS).template data<T>();
- const T* bm_data = bias_multiplier_.template data<T>();
-#if !defined(__ARM_NEON__) && !defined(__ARM_NEON)
- math::Gemm<T, Context>(
- CblasNoTrans,
- CblasNoTrans,
- C,
- output_image_size,
- 1,
- 1,
- bias_data,
- bm_data,
- 1,
- Ydata,
- &context_);
-#else
+ if (bias_data != nullptr) {
+ // Bias term
+#if defined(__ARM_NEON__) || defined(__ARM_NEON)
math::BiasCHW<T, Context>(
bias_data,
- bm_data,
+ nullptr,
C,
- output_image_size,
- Ydata,
+ Y_HxW,
+ Y_data + image_id * C * Y_HxW,
&context_);
#endif // !defined(__ARM_NEON__) && !defined(__ARM_NEON)
}
-
- Xdata += M * H * W;
- Ydata += Y->numel() / Y->dim32(0);
+ }
+ if (bias_data != nullptr) {
+#if !defined(__ARM_NEON__) && !defined(__ARM_NEON)
+ // Bias term
+ const std::array<int, 3> Y_dims = {N, C, Y_HxW};
+ const std::array<int, 3> b_dims = {1, C, 1};
+ math::Add<T, Context>(
+ 3,
+ Y_dims.data(),
+ 3,
+ b_dims.data(),
+ Y_data,
+ bias_data,
+ Y_data,
+ &context_);
+#endif
}
};
+
if (FLAGS_caffe2_force_shared_col_buffer || shared_buffer_) {
- runWithSharedBuffer<Context>(ws_, f);
+ runWithSharedBuffer<Context>(ws_, func);
} else {
- f(&col_buffer_);
+ func(&col_buffer_);
}
return true;
}
template <typename T, class Context>
bool ConvTransposeOp<T, Context>::RunOnDeviceWithOrderNHWC() {
- const Tensor& X = Input(INPUT);
+ const auto& X = Input(INPUT);
auto& filter = Input(FILTER);
- const auto N = X.dim32(0), H = X.dim32(1), W = X.dim32(2), M = X.dim32(3);
- CAFFE_ENFORCE(filter.dim() == 4, "filter must be 4D tensor");
- CAFFE_ENFORCE(
- filter.dim32(0) == M,
+ CAFFE_ENFORCE_EQ(filter.dim(), 4, "filter must be 4D tensor");
+ const int N = X.dim32(0);
+ const int H = X.dim32(1);
+ const int W = X.dim32(2);
+ const int M = X.dim32(3);
+ const int G = group_;
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(0),
+ M,
"filter number must be equal to input channel number");
- CAFFE_ENFORCE(
- filter.dim32(1) == this->kernel_h(),
+ CAFFE_ENFORCE_EQ(
+ M % G, 0, "The number of input channels is not divisible by group.");
+ const int C = filter.dim32(3) * G;
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(1),
+ kernel_h(),
"filter height must be equal to kernel height");
- CAFFE_ENFORCE(
- filter.dim32(2) == this->kernel_w(),
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(2),
+ kernel_w(),
"filter width must be equal to kernel width");
- const int C = filter.dim32(3);
- auto sizes = ConvTransposeUnpoolBase<Context>::GetOutputSize(X, C);
- Tensor* Y = Output(0, sizes, at::dtype<T>());
- const auto kernel_dim = C * this->kernel_h() * this->kernel_w();
- const auto input_image_size = H * W;
- const auto output_image_size = Y->dim32(1) * Y->dim32(2);
+ const std::vector<std::int64_t> Y_dims =
+ ConvTransposeUnpoolBase<Context>::GetOutputSize(X, C);
+ auto* Y = Output(0, Y_dims, at::dtype<T>());
+ if (N == 0) {
+ return true;
+ }
+
+ const int K_HxW = kernel_h() * kernel_w();
+ const int kernel_dim = C / G * K_HxW;
+ const int X_HxW = H * W;
+ const int Y_HxW = Y->dim32(1) * Y->dim32(2);
+ const T* X_data = X.template data<T>();
+ const T* filter_data = filter.template data<T>();
+ const T* bias_data = nullptr;
if (InputSize() == 3) {
auto& bias = Input(BIAS);
- CAFFE_ENFORCE(bias.dim() == 1, "bias must be 1D tensor");
- CAFFE_ENFORCE(
- bias.dim32(0) == C,
+ CAFFE_ENFORCE_EQ(bias.dim(), 1, "bias must be 1D tensor");
+ CAFFE_ENFORCE_EQ(
+ bias.dim32(0),
+ C,
"bias dimension must be equal to output channel number");
- // TODO(jerryzh): is it OK to remove the check of whether numel is output_image_size
- ReinitializeTensor(
- &bias_multiplier_,
- {1, output_image_size},
- at::dtype<T>().device(Context::GetDeviceType()));
- T* bm_data = bias_multiplier_.template mutable_data<T>();
- math::Set<T, Context>(
- output_image_size,
- static_cast<T>(1),
- bm_data,
- &context_);
+ bias_data = bias.template data<T>();
}
- const T* Xdata = X.template data<T>();
- const T* filter_data = filter.template data<T>();
- T* Ydata = Y->template mutable_data<T>();
+ T* Y_data = Y->template mutable_data<T>();
- auto f = [&](Tensor* /*col_buffer*/) {
+ const std::vector<std::int64_t> buffer_shape = {
+ G, H, W, kernel_h(), kernel_w(), C / G};
+ const auto func = [&](Tensor* /*col_buffer*/) {
ReinitializeTensor(
&col_buffer_,
- vector<int64_t>{H, W, this->kernel_h(), this->kernel_w(), C},
+ buffer_shape,
at::dtype<T>().device(Context::GetDeviceType()));
T* col_buffer_data = col_buffer_.template mutable_data<T>();
- for (auto image_id = 0; image_id < N; ++image_id) {
+ for (int image_id = 0; image_id < N; ++image_id) {
// Weight term
- math::Gemm<T, Context>(
- CblasNoTrans,
- CblasNoTrans,
- input_image_size,
- kernel_dim,
- M,
- 1,
- Xdata,
- filter_data,
- 0,
- col_buffer_data,
- &context_);
+ if (G == 1) {
+ math::Gemm<T, Context>(
+ CblasNoTrans,
+ CblasNoTrans,
+ X_HxW,
+ kernel_dim,
+ M,
+ 1.0f,
+ X_data + image_id * M * X_HxW,
+ filter_data,
+ 0.0f,
+ col_buffer_data,
+ &context_);
+ } else {
+ for (int group_id = 0; group_id < G; ++group_id) {
+ math::GemmEx<T, Context>(
+ CblasNoTrans,
+ CblasNoTrans,
+ X_HxW,
+ kernel_dim,
+ M / G,
+ 1.0f,
+ X_data + image_id * M * X_HxW + group_id * M / G,
+ M,
+ filter_data + group_id * M / G * kernel_dim,
+ kernel_dim,
+ 0.0f,
+ col_buffer_data + group_id * kernel_dim,
+ G * kernel_dim,
+ &context_);
+ }
+ }
// Col2Im
math::Col2Im<T, Context, StorageOrder::NHWC>(
C,
Y->dim32(1),
Y->dim32(2),
- this->kernel_h(),
- this->kernel_w(),
+ kernel_h(),
+ kernel_w(),
1,
1,
- this->pad_t(),
- this->pad_l(),
- this->pad_b(),
- this->pad_r(),
- this->stride_h(),
- this->stride_w(),
+ pad_t(),
+ pad_l(),
+ pad_b(),
+ pad_r(),
+ stride_h(),
+ stride_w(),
col_buffer_data,
- Ydata,
- &context_);
+ Y_data + image_id * C * Y_HxW,
+ &context_,
+ G);
+ }
+ if (bias_data != nullptr) {
// Bias term
- if (InputSize() == 3) {
- const T* bm_data = bias_multiplier_.template data<T>();
- const T* bias_data = Input(BIAS).template data<T>();
- math::Gemm<T, Context>(
- CblasNoTrans,
- CblasNoTrans,
- output_image_size,
- C,
- 1,
- 1,
- bm_data,
- bias_data,
- 1,
- Ydata,
- &context_);
- }
- Xdata += M * H * W;
- Ydata += Y->numel() / Y->dim32(0);
+ const std::array<int, 2> Y_dims = {N * Y_HxW, C};
+ const std::array<int, 2> b_dims = {1, C};
+ math::Add<T, Context>(
+ 2,
+ Y_dims.data(),
+ 2,
+ b_dims.data(),
+ Y_data,
+ bias_data,
+ Y_data,
+ &context_);
}
};
+
if (FLAGS_caffe2_force_shared_col_buffer || shared_buffer_) {
- runWithSharedBuffer<Context>(ws_, f);
+ runWithSharedBuffer<Context>(ws_, func);
} else {
- f(&col_buffer_);
+ func(&col_buffer_);
}
return true;
}
template <typename T, class Context>
bool ConvTransposeGradientOp<T, Context>::RunOnDeviceWithOrderNCHW() {
- auto& X = Input(INPUT);
- auto& filter = Input(FILTER);
- auto& dY = Input(OUTPUT_GRAD);
-
- const int N = X.dim32(0), M = X.dim32(1), H = X.dim32(2), W = X.dim32(3);
- // We only handle LegacyPadding::NOTSET case and ignore cases of
- // LegacyPadding::VALID and LegacyPadding::SAME
- // Thus, we don't need to manually compute padding values
- // We simply use the values from the user
- CAFFE_ENFORCE(filter.dim() == 4);
- const int C = filter.dim32(1);
- CAFFE_ENFORCE(
- filter.dim32(2) == this->kernel_h(),
+ const auto& X = Input(INPUT);
+ const auto& filter = Input(FILTER);
+ const auto& dY = Input(OUTPUT_GRAD);
+ CAFFE_ENFORCE_EQ(filter.dim(), 4);
+ const int N = X.dim32(0);
+ const int M = X.dim32(1);
+ const int H = X.dim32(2);
+ const int W = X.dim32(3);
+ const int G = group_;
+ CAFFE_ENFORCE_EQ(M, filter.dim32(0));
+ CAFFE_ENFORCE_EQ(
+ M % G, 0, "The number of input channels is not divisible by group.");
+ const int C = filter.dim32(1) * G;
+ CAFFE_ENFORCE_EQ(C, dY.dim32(1));
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(2),
+ kernel_h(),
"filter height must be equal to kernel height");
- CAFFE_ENFORCE(
- filter.dim32(3) == this->kernel_w(),
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(3),
+ this->kernel_w(),
"filter width must be equal to kernel width");
+
+ const int K_HxW = kernel_h() * kernel_w();
+ const int kernel_dim = C / G * K_HxW;
+ const int X_HxW = H * W;
+ const int Y_HxW = dY.dim32(2) * dY.dim32(3);
auto* dfilter = Output(FILTER_GRAD, filter.sizes(), at::dtype<T>());
- const int kernel_dim = C * this->kernel_h() * this->kernel_w();
- const int output_image_size = dY.dim32(2) * dY.dim32(3);
- // The col buffer is stored in CHW order as well
- ReinitializeTensor(
- &col_buffer_,
- vector<int64_t>{C, this->kernel_h(), this->kernel_w(), H, W},
- at::dtype<T>().device(Context::GetDeviceType()));
- if (!no_bias_) {
- auto* dbias = Output(BIAS_OR_INPUT_GRAD);
- dbias->Resize(C);
- // TODO(jerryzh): is it OK to remove the check of whether numel is output_image_size
- ReinitializeTensor(
- &bias_multiplier_,
- {1, output_image_size},
- at::dtype<T>().device(Context::GetDeviceType()));
- T* bm_data = bias_multiplier_.template mutable_data<T>();
- math::Set<T, Context>(
- output_image_size,
- static_cast<T>(1),
- bm_data,
- &context_);
- }
- T* col_buffer_data = col_buffer_.template mutable_data<T>();
- const T* Xdata = X.template data<T>();
+ const T* X_data = X.template data<T>();
const T* filter_data = filter.template data<T>();
- const T* dYdata = dY.template data<T>();
+ const T* dY_data = dY.template data<T>();
T* dfilter_data = dfilter->template mutable_data<T>();
- // Pre-setting the gradients to zero
- math::Set<T, Context>(dfilter->numel(), 0, dfilter_data, &context_);
+ T* dbias_data = nullptr;
+ T* dX_data = nullptr;
if (!no_bias_) {
- auto* dbias = Output(BIAS_OR_INPUT_GRAD);
- T* dbias_data = dbias->template mutable_data<T>();
- math::Set<T, Context>(dbias->numel(), 0, dbias_data, &context_);
+ auto* dbias = Output(BIAS_OR_INPUT_GRAD, {C}, at::dtype<T>());
+ dbias_data = dbias->template mutable_data<T>();
}
- for (auto image_id = 0; image_id < N; ++image_id) {
+ const bool compute_dX =
+ (OutputSize() == 3) || (no_bias_ && (OutputSize() == 2));
+ if (compute_dX) {
+ auto* dX = Output(
+ no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD, X.sizes(), at::dtype<T>());
+ dX_data = dX->template mutable_data<T>();
+ }
+ math::Set<T, Context>(filter.numel(), T(0), dfilter_data, &context_);
+
+ if (N == 0) {
+ math::Set<T, Context>(C, T(0), dbias_data, &context_);
+ return true;
+ }
+
+ ReinitializeTensor(
+ &col_buffer_,
+ std::vector<std::int64_t>{C, kernel_h(), kernel_w(), H, W},
+ at::dtype<T>().device(Context::GetDeviceType()));
+ T* col_buffer_data = col_buffer_.template mutable_data<T>();
+
+ for (int image_id = 0; image_id < N; ++image_id) {
// gradient w.r.t. filters. Im2Col followed by Gemm
// Im2Col.
math::Im2Col<T, Context, StorageOrder::NCHW>(
C,
dY.dim32(2),
dY.dim32(3),
- this->kernel_h(),
- this->kernel_w(),
+ kernel_h(),
+ kernel_w(),
1,
1,
- this->pad_t(),
- this->pad_l(),
- this->pad_b(),
- this->pad_r(),
- this->stride_h(),
- this->stride_w(),
- dYdata,
+ pad_t(),
+ pad_l(),
+ pad_b(),
+ pad_r(),
+ stride_h(),
+ stride_w(),
+ dY_data + image_id * C * Y_HxW,
col_buffer_data,
&context_);
// Gemm
- math::Gemm<T, Context>(
- CblasNoTrans,
- CblasTrans,
- M,
- kernel_dim,
- H * W,
- 1,
- Xdata,
- col_buffer_data,
- 1,
- dfilter_data,
- &context_);
- // gradient w.r.t. bias
- if (!no_bias_) {
- const T* bm_data = bias_multiplier_.template data<T>();
- T* input_grad_data = Output(BIAS_OR_INPUT_GRAD)->template mutable_data<T>();
+ if (G == 1) {
math::Gemm<T, Context>(
CblasNoTrans,
- CblasNoTrans,
- C,
- 1,
- output_image_size,
- 1,
- dYdata,
- bm_data,
- 1,
- input_grad_data,
- &context_);
- }
- dYdata += dY.numel() / dY.dim32(0);
- Xdata += X.numel() / X.dim32(0);
- }
- if (OutputSize() == 3 || (no_bias_ && (OutputSize() == 2))) {
- // Compute gradients w.r.t. the input
- // Since we have changed dYdata in the above loop, we will need to reset.
- dYdata = dY.template data<T>();
-
- auto* dX = Output(
- no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD, X.sizes(), at::dtype<T>());
- T* dXdata = dX->template mutable_data<T>();
- for (auto image_id = 0; image_id < N; ++image_id) {
- // Im2Col.
- // TODO(zyan3): Probably duplicate work as in gradient computation
- // w.r.t filters
- math::Im2Col<T, Context, StorageOrder::NCHW>(
- C,
- dY.dim32(2),
- dY.dim32(3),
- this->kernel_h(),
- this->kernel_w(),
- 1,
- 1,
- this->pad_t(),
- this->pad_l(),
- this->pad_b(),
- this->pad_r(),
- this->stride_h(),
- this->stride_w(),
- dYdata,
+ CblasTrans,
+ M,
+ kernel_dim,
+ X_HxW,
+ 1.0f,
+ X_data + image_id * M * X_HxW,
col_buffer_data,
+ 1.0f,
+ dfilter_data,
&context_);
- // Gemm
- math::Gemm<T, Context>(
+ } else {
+ math::GemmStridedBatched<T, Context>(
CblasNoTrans,
- CblasNoTrans,
- M,
- H * W,
+ CblasTrans,
+ G,
+ M / G,
kernel_dim,
- 1,
- filter_data,
+ X_HxW,
+ 1.0f,
+ X_data + image_id * M * X_HxW,
+ M / G * X_HxW,
col_buffer_data,
- 0,
- dXdata,
+ col_buffer_.numel() / G,
+ 1.0f,
+ dfilter_data,
+ M / G * kernel_dim,
&context_);
- dYdata += dY.numel() / dY.dim32(0);
- dXdata += X.numel() / X.dim32(0);
+ }
+
+ if (dX_data != nullptr) {
+ // Compute gradients w.r.t. the input
+ if (G == 1) {
+ math::Gemm<T, Context>(
+ CblasNoTrans,
+ CblasNoTrans,
+ M,
+ X_HxW,
+ kernel_dim,
+ 1.0f,
+ filter_data,
+ col_buffer_data,
+ 0.0f,
+ dX_data + image_id * M * X_HxW,
+ &context_);
+ } else {
+ math::GemmStridedBatched<T, Context>(
+ CblasNoTrans,
+ CblasNoTrans,
+ G,
+ M / G,
+ X_HxW,
+ kernel_dim,
+ 1.0f,
+ filter_data,
+ M / G * kernel_dim,
+ col_buffer_data,
+ col_buffer_.numel() / G,
+ 0.0f,
+ dX_data + image_id * M * X_HxW,
+ M / G * X_HxW,
+ &context_);
+ }
}
}
+
+ if (dbias_data != nullptr) {
+ // gradient w.r.t. bias
+ const std::array<int, 3> Y_dims = {N, C, Y_HxW};
+ const std::array<int, 3> b_dims = {1, C, 1};
+ math::ReduceSum<T, Context>(
+ 3, Y_dims.data(), b_dims.data(), T(1), dY_data, dbias_data, &context_);
+ }
+
return true;
}
template <typename T, class Context>
bool ConvTransposeGradientOp<T, Context>::RunOnDeviceWithOrderNHWC() {
- auto& X = Input(INPUT);
- auto& filter = Input(FILTER);
- auto& dY = Input(OUTPUT_GRAD);
-
- const int N = X.dim32(0), H = X.dim32(1), W = X.dim32(2), M = X.dim32(3);
- // We only handle LegacyPadding::NOTSET case and ignore cases of
- // LegacyPadding::VALID and LegacyPadding::SAME
- // Thus, we don't need to manually compute padding values
- // We simply use the values from the user
- CAFFE_ENFORCE(filter.dim() == 4, "filter must be 4D tensor");
- CAFFE_ENFORCE(
- filter.dim32(1) == this->kernel_h(),
+ const auto& X = Input(INPUT);
+ const auto& filter = Input(FILTER);
+ const auto& dY = Input(OUTPUT_GRAD);
+ CAFFE_ENFORCE_EQ(filter.dim(), 4);
+ const int N = X.dim32(0);
+ const int H = X.dim32(1);
+ const int W = X.dim32(2);
+ const int M = X.dim32(3);
+ const int G = group_;
+ CAFFE_ENFORCE_EQ(M, filter.dim32(0));
+ CAFFE_ENFORCE_EQ(
+ M % G, 0, "The number of input channels is not divisible by group.");
+ const int C = filter.dim32(3) * G;
+ CAFFE_ENFORCE_EQ(C, dY.dim32(3));
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(1),
+ kernel_h(),
"filter height must be equal to kernel height");
- CAFFE_ENFORCE(
- filter.dim32(2) == this->kernel_w(),
+ CAFFE_ENFORCE_EQ(
+ filter.dim32(2),
+ this->kernel_w(),
"filter width must be equal to kernel width");
- const int C = filter.dim32(3);
+ CAFFE_ENFORCE_EQ(dY.dim32(3), C);
+
+ const int K_HxW = kernel_h() * kernel_w();
+ const int kernel_dim = C / G * K_HxW;
+ const int X_HxW = H * W;
+ const int Y_HxW = dY.dim32(1) * dY.dim32(2);
auto* dfilter = Output(FILTER_GRAD, filter.sizes(), at::dtype<T>());
- const int kernel_dim = C * this->kernel_h() * this->kernel_w();
- const int output_image_size = dY.dim32(1) * dY.dim32(2);
- // The col buffer is stored in HWC order as well
- ReinitializeTensor(
- &col_buffer_,
- vector<int64_t>{H, W, this->kernel_h(), this->kernel_w(), C},
- at::dtype<T>().device(Context::GetDeviceType()));
- if (!no_bias_) {
- auto* dbias = Output(BIAS_OR_INPUT_GRAD);
- dbias->Resize(C);
- // TODO(jerryzh): is it OK to remove the check of whether numel is output_image_size
- ReinitializeTensor(
- &bias_multiplier_,
- {1, output_image_size},
- at::dtype<T>().device(Context::GetDeviceType()));
- T* bm_data = bias_multiplier_.template mutable_data<T>();
- math::Set<T, Context>(
- output_image_size,
- static_cast<T>(1),
- bm_data,
- &context_);
- }
- T* col_buffer_data = col_buffer_.template mutable_data<T>();
- const T* Xdata = X.template data<T>();
+ const T* X_data = X.template data<T>();
const T* filter_data = filter.template data<T>();
- const T* dYdata = dY.template data<T>();
+ const T* dY_data = dY.template data<T>();
T* dfilter_data = dfilter->template mutable_data<T>();
- // Pre-setting the gradients to zero
- math::Set<T, Context>(dfilter->numel(), 0, dfilter_data, &context_);
+ T* dbias_data = nullptr;
+ T* dX_data = nullptr;
if (!no_bias_) {
- auto* dbias = Output(BIAS_OR_INPUT_GRAD);
- T* dbias_data = dbias->template mutable_data<T>();
- math::Set<T, Context>(dbias->numel(), 0, dbias_data, &context_);
+ auto* dbias = Output(BIAS_OR_INPUT_GRAD, {C}, at::dtype<T>());
+ dbias_data = dbias->template mutable_data<T>();
+ }
+ const bool compute_dX =
+ (OutputSize() == 3) || (no_bias_ && (OutputSize() == 2));
+ if (compute_dX) {
+ auto* dX = Output(
+ no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD, X.sizes(), at::dtype<T>());
+ dX_data = dX->template mutable_data<T>();
}
- for (auto image_id = 0; image_id < N; ++image_id) {
+ math::Set<T, Context>(filter.numel(), T(0), dfilter_data, &context_);
+
+ if (N == 0) {
+ math::Set<T, Context>(C, T(0), dbias_data, &context_);
+ return true;
+ }
+
+ ReinitializeTensor(
+ &col_buffer_,
+ std::vector<std::int64_t>{C, kernel_h(), kernel_w(), H, W},
+ at::dtype<T>().device(Context::GetDeviceType()));
+ T* col_buffer_data = col_buffer_.template mutable_data<T>();
+
+ for (int image_id = 0; image_id < N; ++image_id) {
// gradient w.r.t. filters. Im2Col followed by Gemm
// Im2Col.
math::Im2Col<T, Context, StorageOrder::NHWC>(
C,
dY.dim32(1),
dY.dim32(2),
- this->kernel_h(),
- this->kernel_w(),
+ kernel_h(),
+ kernel_w(),
1,
1,
- this->pad_t(),
- this->pad_l(),
- this->pad_b(),
- this->pad_r(),
- this->stride_h(),
- this->stride_w(),
- dYdata,
+ pad_t(),
+ pad_l(),
+ pad_b(),
+ pad_r(),
+ stride_h(),
+ stride_w(),
+ dY_data + image_id * C * Y_HxW,
col_buffer_data,
- &context_);
+ &context_,
+ G);
// Gemm
- math::Gemm<T, Context>(
- CblasTrans,
- CblasNoTrans,
- M,
- kernel_dim,
- H * W,
- 1,
- Xdata,
- col_buffer_data,
- 1,
- dfilter_data,
- &context_);
- // gradients w.r.t. bias
- if (!no_bias_) {
- const T* bm_data = bias_multiplier_.template data<T>();
- T* input_grad_data = Output(BIAS_OR_INPUT_GRAD)->template mutable_data<T>();
+ if (G == 1) {
math::Gemm<T, Context>(
CblasTrans,
CblasNoTrans,
- C,
- 1,
- output_image_size,
- 1,
- dYdata,
- bm_data,
- 1,
- input_grad_data,
- &context_);
- }
- dYdata += dY.numel() / dY.dim32(0);
- Xdata += X.numel() / X.dim32(0);
- }
- if (OutputSize() == 3 || (no_bias_ && (OutputSize() == 2))) {
- // Compute gradients w.r.t. the input
- // Since we have changed dYdata in the above loop, we will need to reset.
- dYdata = dY.template data<T>();
-
- auto* dX = Output(
- no_bias_ ? BIAS_OR_INPUT_GRAD : INPUT_GRAD, X.sizes(), at::dtype<T>());
- T* dXdata = dX->template mutable_data<T>();
- for (auto image_id = 0; image_id < N; ++image_id) {
- // Im2Col.
- // TODO(zyan3): Probably duplicate work as in gradient computation
- // w.r.t filters
- math::Im2Col<T, Context, StorageOrder::NHWC>(
- C,
- dY.dim32(1),
- dY.dim32(2),
- this->kernel_h(),
- this->kernel_w(),
- 1,
- 1,
- this->pad_t(),
- this->pad_l(),
- this->pad_b(),
- this->pad_r(),
- this->stride_h(),
- this->stride_w(),
- dYdata,
- col_buffer_data,
- &context_);
- // Gemm
- math::Gemm<T, Context>(
- CblasNoTrans,
- CblasTrans,
- H * W,
M,
kernel_dim,
- 1,
+ X_HxW,
+ 1.0f,
+ X_data + image_id * M * X_HxW,
col_buffer_data,
- filter_data,
- 0,
- dXdata,
+ 1.0f,
+ dfilter_data,
&context_);
- dYdata += dY.numel() / dY.dim32(0);
- dXdata += X.numel() / X.dim32(0);
+ } else {
+ for (int group_id = 0; group_id < G; ++group_id) {
+ math::GemmEx<T, Context>(
+ CblasTrans,
+ CblasNoTrans,
+ M / G,
+ kernel_dim,
+ X_HxW,
+ 1.0f,
+ X_data + image_id * M * X_HxW + group_id * M / G,
+ M,
+ col_buffer_data + group_id * kernel_dim,
+ G * kernel_dim,
+ 1.0f,
+ dfilter_data + group_id * M / G * kernel_dim,
+ kernel_dim,
+ &context_);
+ }
+ }
+
+ if (dX_data != nullptr) {
+ // Compute gradients w.r.t. the input
+ if (G == 1) {
+ math::Gemm<T, Context>(
+ CblasNoTrans,
+ CblasTrans,
+ X_HxW,
+ M,
+ kernel_dim,
+ 1.0f,
+ col_buffer_data,
+ filter_data,
+ 0.0f,
+ dX_data + image_id * M * X_HxW,
+ &context_);
+ } else {
+ for (int group_id = 0; group_id < G; ++group_id) {
+ math::GemmEx<T, Context>(
+ CblasNoTrans,
+ CblasTrans,
+ X_HxW,
+ M / G,
+ kernel_dim,
+ 1.0f,
+ col_buffer_data + group_id * kernel_dim,
+ G * kernel_dim,
+ filter_data + group_id * M / G * kernel_dim,
+ kernel_dim,
+ 0.0f,
+ dX_data + image_id * M * X_HxW + group_id * M / G,
+ M,
+ &context_);
+ }
+ }
}
}
+
+ if (dbias_data != nullptr) {
+ const std::array<int, 2> Y_dims = {N * Y_HxW, C};
+ const std::array<int, 2> b_dims = {1, C};
+ math::ReduceSum<T, Context>(
+ 2, Y_dims.data(), b_dims.data(), T(1), dY_data, dbias_data, &context_);
+ }
+
return true;
}
} // namespace caffe2
+
#endif // CAFFE2_OPERATORS_CONV_TRANSPOSE_OP_IMPL_H_
diff --git a/caffe2/operators/conv_transpose_unpool_op_base.h b/caffe2/operators/conv_transpose_unpool_op_base.h
index 7ebfda7e09..c98b3bae21 100644
--- a/caffe2/operators/conv_transpose_unpool_op_base.h
+++ b/caffe2/operators/conv_transpose_unpool_op_base.h
@@ -17,7 +17,9 @@ template <class Context>
class ConvTransposeUnpoolBase : public Operator<Context> {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
- explicit ConvTransposeUnpoolBase(const OperatorDef& operator_def, Workspace* ws)
+ explicit ConvTransposeUnpoolBase(
+ const OperatorDef& operator_def,
+ Workspace* ws)
: Operator<Context>(operator_def, ws),
legacy_pad_(
static_cast<LegacyPadding>(this->template GetSingleArgument<int>(
@@ -27,6 +29,7 @@ class ConvTransposeUnpoolBase : public Operator<Context> {
stride_(this->template GetRepeatedArgument<int>("strides")),
pads_(this->template GetRepeatedArgument<int>("pads")),
adj_(this->template GetRepeatedArgument<int>("adjs")),
+ group_(this->template GetSingleArgument<int>("group", 1)),
order_(StringToStorageOrder(
this->template GetSingleArgument<string>("order", "NCHW"))),
shared_buffer_(
@@ -206,19 +209,7 @@ class ConvTransposeUnpoolBase : public Operator<Context> {
virtual ~ConvTransposeUnpoolBase() {}
- private:
- LegacyPadding legacy_pad_;
- int pad_;
-
protected:
- vector<int> kernel_;
- vector<int> stride_;
- vector<int> pads_;
- vector<int> adj_;
- StorageOrder order_;
- bool shared_buffer_;
- Workspace* ws_;
-
// Accessors for 2D conv params.
inline int pad_t() const {
@@ -289,14 +280,35 @@ class ConvTransposeUnpoolBase : public Operator<Context> {
break;
}
}
+
+ LegacyPadding legacy_pad_;
+ int pad_;
+
+ std::vector<int> kernel_;
+ std::vector<int> stride_;
+ std::vector<int> pads_;
+ std::vector<int> adj_;
+ int group_;
+ StorageOrder order_;
+ bool shared_buffer_;
+ Workspace* ws_;
};
#define USE_CONV_TRANSPOSE_UNPOOL_BASE_FUNCTIONS(Context) \
USE_OPERATOR_FUNCTIONS(Context); \
using ConvTransposeUnpoolBase<Context>::kernel_; \
+ using ConvTransposeUnpoolBase<Context>::kernel_h; \
+ using ConvTransposeUnpoolBase<Context>::kernel_w; \
using ConvTransposeUnpoolBase<Context>::stride_; \
+ using ConvTransposeUnpoolBase<Context>::stride_h; \
+ using ConvTransposeUnpoolBase<Context>::stride_w; \
using ConvTransposeUnpoolBase<Context>::pads_; \
+ using ConvTransposeUnpoolBase<Context>::pad_t; \
+ using ConvTransposeUnpoolBase<Context>::pad_l; \
+ using ConvTransposeUnpoolBase<Context>::pad_b; \
+ using ConvTransposeUnpoolBase<Context>::pad_r; \
using ConvTransposeUnpoolBase<Context>::adj_; \
+ using ConvTransposeUnpoolBase<Context>::group_; \
using ConvTransposeUnpoolBase<Context>::order_; \
using ConvTransposeUnpoolBase<Context>::shared_buffer_; \
using ConvTransposeUnpoolBase<Context>::ws_
diff --git a/caffe2/python/operator_test/conv_transpose_test.py b/caffe2/python/operator_test/conv_transpose_test.py
index 2a32f9acee..272ac3a3e7 100644
--- a/caffe2/python/operator_test/conv_transpose_test.py
+++ b/caffe2/python/operator_test/conv_transpose_test.py
@@ -6,6 +6,7 @@ import numpy as np
from hypothesis import assume, given, settings
import hypothesis.strategies as st
+from caffe2.proto import caffe2_pb2
from caffe2.python import core, utils
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.hip_test_util as hiputl
@@ -360,6 +361,68 @@ class TestConvolutionTranspose(hu.HypothesisTestCase):
for i in outputs_to_check:
self.assertGradientChecks(gc, op, inputs, i, [0])
+ @given(stride=st.integers(1, 3),
+ pad=st.integers(0, 3),
+ kernel=st.integers(1, 3),
+ adj=st.integers(0, 2),
+ size=st.integers(7, 10),
+ input_channels=st.integers(1, 8),
+ output_channels=st.integers(1, 8),
+ batch_size=st.integers(1, 4),
+ group=st.integers(1, 4),
+ order=st.sampled_from(["NCHW", "NHWC"]),
+ engine=st.sampled_from(["", "CUDNN", "BLOCK"]),
+ shared_buffer=st.booleans(),
+ use_bias=st.booleans(),
+ **hu.gcs)
+ def test_convolution_transpose_with_group(
+ self, stride, pad, kernel, adj, size, input_channels,
+ output_channels, batch_size, group, order, engine, shared_buffer,
+ use_bias, gc, dc):
+ assume(adj < stride)
+ # TODO: Group conv_transpose in NHWC not implemented for GPU yet.
+ assume(group == 1 or order == "NCHW" or
+ gc.device_type == caffe2_pb2.CPU)
+ if group != 1 and order == "NHWC":
+ dc = [d for d in dc if d.device_type == caffe2_pb2.CPU]
+
+ if hiputl.run_in_hip(gc, dc) and order == "NHWC":
+ engine = ""
+
+ op = core.CreateOperator(
+ "ConvTranspose",
+ ["X", "w", "b"] if use_bias else ["X", "w"],
+ ["Y"],
+ stride=stride,
+ kernel=kernel,
+ pad=pad,
+ adj=adj,
+ group=group,
+ order=order,
+ engine=engine,
+ shared_buffer=int(shared_buffer),
+ device_option=gc,
+ )
+
+ input_channels *= group
+ output_channels *= group
+
+ X = np.random.rand(
+ batch_size, size, size, input_channels).astype(np.float32) - 0.5
+ w = np.random.rand(
+ input_channels, kernel, kernel, int(output_channels / group)) \
+ .astype(np.float32) - 0.5
+ b = np.random.rand(output_channels).astype(np.float32) - 0.5
+ if order == "NCHW":
+ X = utils.NHWC2NCHW(X)
+ w = utils.NHWC2NCHW(w)
+
+ inputs = [X, w, b] if use_bias else [X, w]
+ self.assertDeviceChecks(dc, op, inputs, [0])
+ for i in range(len(inputs)):
+ self.assertGradientChecks(gc, op, inputs, i, [0])
+
+
if __name__ == "__main__":
import unittest
unittest.main()
generated by cgit v1.2.3 (git 2.25.1) at 2024-11-26 05:36:03 +0000