| #ifndef CAFFE2_OPERATORS_TILE_OP_H_ |
| #define CAFFE2_OPERATORS_TILE_OP_H_ |
| |
| #include <array> |
| #include <string> |
| #include <type_traits> |
| #include <vector> |
| |
| #include "caffe2/core/common_omp.h" |
| #include "caffe2/core/context.h" |
| #include "caffe2/core/logging.h" |
| #include "caffe2/core/operator.h" |
| #include "caffe2/utils/eigen_utils.h" |
| #include "caffe2/utils/math.h" |
| |
| namespace caffe2 { |
| |
| // Copy a Blob n times along a specified axis. |
| template <class Context> |
| class TileOp final : public Operator<Context> { |
| public: |
| USE_OPERATOR_CONTEXT_FUNCTIONS; |
| |
| template <class... Args> |
| explicit TileOp(Args&&... args) |
| : Operator<Context>(std::forward<Args>(args)...), |
| OP_SINGLE_ARG(std::int32_t, "tiles", tiles_, 1), |
| OP_SINGLE_ARG(std::int32_t, "axis", axis_, 0) {} |
| |
| bool RunOnDevice() override { |
| return DispatchHelper< |
| TensorTypes<std::int32_t, std::int64_t, float, double>>:: |
| call(this, Input(0)); |
| } |
| |
| template <typename T> |
| bool DoRunWithType() { |
| if (InputSize() > 1) { |
| // We potentially have tiles and/or axis specified as inputs |
| // as well. We will check for them in that order. In other words: |
| // InputSize() == 2: tiles is specified |
| // InputSize() == 3: tiles is specified and axis. |
| // Anything specified as input will override the arguments |
| CAFFE_ENFORCE( |
| Input(1).dim() == 1 && Input(1).numel() == 1, |
| "Input `tiles` should be a vector of size 1."); |
| tiles_ = GetArgFromTensor(Input(1)); |
| |
| // Because of a bug in original code, temporarily adds this part to keep |
| // backward compatibility. |
| // TODO(yangxm): Remove this part when prod runtime upgraded with fixed |
| // model config. |
| if (Input(1).template IsType<std::int64_t>()) { |
| axis_ = 0; |
| } |
| |
| if (InputSize() > 2) { |
| CAFFE_ENFORCE( |
| Input(2).dim() == 1 && Input(2).numel() == 1, |
| "Input `axis` should be a vector of size 1."); |
| axis_ = GetArgFromTensor(Input(2)); |
| } else { |
| CAFFE_ENFORCE( |
| OperatorBase::HasArgument("axis"), |
| "Argument `axis` is missing and was not specified as input."); |
| } |
| } else { |
| CAFFE_ENFORCE( |
| OperatorBase::HasArgument("tiles"), |
| "Argument `tiles` is missing and was not specified as input."); |
| CAFFE_ENFORCE( |
| OperatorBase::HasArgument("axis"), |
| "Argument `axis` is missing and was not specified as input."); |
| } |
| |
| const auto& X = Input(0); |
| auto* Y = Output(0); |
| const int axis = X.canonical_axis_index(axis_); |
| |
| // reshape output to be input tiled along the axis |
| std::vector<std::int64_t> Y_dims = X.sizes().vec(); |
| Y_dims[axis] *= tiles_; |
| Y->Resize(Y_dims); |
| |
| // size up to (and not including) axis |
| const int outer_size = X.size_to_dim(axis); |
| // size from axis up |
| const int inner_size = X.size_from_dim(axis); |
| |
| const T* X_data = X.template data<T>(); |
| T* Y_data = Y->template mutable_data<T>(); |
| return DoTile<T>(outer_size, inner_size, X_data, Y_data); |
| } |
| |
| private: |
| std::int32_t GetArgFromTensor(const Tensor& tensor) { |
| CAFFE_ENFORCE( |
| tensor.IsType<std::int32_t>() || tensor.IsType<std::int64_t>()); |
| std::int32_t val = -1; |
| if (tensor.IsType<std::int32_t>()) { |
| context_.template CopyToCPU<std::int32_t>( |
| 1, tensor.data<std::int32_t>(), &val); |
| } else if (tensor.IsType<std::int64_t>()) { |
| std::int64_t val_int64; |
| context_.template CopyToCPU<std::int64_t>( |
| 1, tensor.data<std::int64_t>(), &val_int64); |
| val = static_cast<std::int32_t>(val_int64); |
| } |
| return val; |
| } |
| |
| template <typename T> |
| bool DoTile(const int outer_size, const int inner_size, const T* X, T* Y) { |
| if (inner_size == 1) { |
| EigenArrayMap<T> Y_arr(Y, tiles_, outer_size); |
| for (const auto i : c10::irange(outer_size)) { |
| Y_arr.col(i) = X[i]; |
| } |
| } else { |
| ConstEigenArrayMap<T> X_arr(X, inner_size, outer_size); |
| for (const auto i : c10::irange(outer_size)) { |
| EigenArrayMap<T>(Y + i * tiles_ * inner_size, inner_size, tiles_) |
| .colwise() = X_arr.col(i); |
| } |
| } |
| return true; |
| } |
| |
| std::int32_t tiles_; |
| std::int32_t axis_; |
| }; |
| |
| template <class Context> |
| class TileGradientOp final : public Operator<Context> { |
| public: |
| USE_OPERATOR_CONTEXT_FUNCTIONS; |
| |
| template <class... Args> |
| explicit TileGradientOp(Args&&... args) |
| : Operator<Context>(std::forward<Args>(args)...), |
| OP_SINGLE_ARG(std::int32_t, "tiles", tiles_, 1), |
| OP_SINGLE_ARG(std::int32_t, "axis", axis_, 0) {} |
| |
| bool RunOnDevice() override { |
| return DispatchHelper< |
| TensorTypes<std::int32_t, std::int64_t, float, double>>:: |
| call(this, Input(0)); |
| } |
| |
| template <typename T> |
| bool DoRunWithType() { |
| if (InputSize() > 1) { |
| // We potentially have tiles and/or axis specified as inputs |
| // as well. We will check for them in that order. In other words: |
| // InputSize() == 2: tiles is specified |
| // InputSize() == 3: tiles is specified and axis. |
| // Anything specified as input will override the arguments |
| CAFFE_ENFORCE( |
| Input(1).dim() == 1 && Input(1).numel() == 1, |
| "Input `tiles` should be a vector of size 1."); |
| tiles_ = GetArgFromTensor(Input(1)); |
| if (InputSize() > 2) { |
| CAFFE_ENFORCE( |
| Input(2).dim() == 1 && Input(2).numel() == 1, |
| "Input `axis` should be a vector of size 1."); |
| axis_ = GetArgFromTensor(Input(2)); |
| } else { |
| CAFFE_ENFORCE( |
| OperatorBase::HasArgument("axis"), |
| "Argument `axis` is missing and was not specified as input."); |
| } |
| } else { |
| CAFFE_ENFORCE( |
| OperatorBase::HasArgument("tiles"), |
| "Argument `tiles` is missing and was not specified as input."); |
| CAFFE_ENFORCE( |
| OperatorBase::HasArgument("axis"), |
| "Argument `axis` is missing and was not specified as input."); |
| } |
| |
| const auto& dY = Input(0); |
| auto* dX = Output(0); |
| const int axis = dY.canonical_axis_index(axis_); |
| |
| // reshape output to be input "untiled" along the axis |
| std::vector<std::int64_t> X_dims = dY.sizes().vec(); |
| CAFFE_ENFORCE_EQ(X_dims[axis] % tiles_, 0); |
| X_dims[axis] /= tiles_; |
| dX->Resize(X_dims); |
| |
| // size up to (and not including) axis |
| const int outer_size = dX->size_to_dim(axis); |
| // size from axis up |
| const int inner_size = dX->size_from_dim(axis); |
| |
| /** |
| * How this works: |
| * Imagine a 2D tensor (matrix) of size 3x10, tiled 2 times along axis 1 |
| * (column). |
| * This is equivalent to multiplying by a vector of 1s transposed. |
| * The gradient of this is all 1s in the shape of the input matrix |
| * (call it X). |
| * So the output gradient should be the matrix multiplication result |
| * of input gradient (gradient of tiled tensor output) and X. |
| */ |
| const T* dY_data = dY.template data<T>(); |
| T* dX_data = dX->template mutable_data<T>(); |
| return DoTileGradient<T>(outer_size, inner_size, dY_data, dX_data); |
| } |
| |
| private: |
| std::int32_t GetArgFromTensor(const Tensor& tensor) { |
| CAFFE_ENFORCE( |
| tensor.IsType<std::int32_t>() || tensor.IsType<std::int64_t>()); |
| std::int32_t val = -1; |
| if (tensor.IsType<std::int32_t>()) { |
| context_.template CopyToCPU<std::int32_t>( |
| 1, tensor.data<std::int32_t>(), &val); |
| } else if (tensor.IsType<std::int64_t>()) { |
| std::int64_t val_int64; |
| context_.template CopyToCPU<std::int64_t>( |
| 1, tensor.data<std::int64_t>(), &val_int64); |
| val = static_cast<std::int32_t>(val_int64); |
| } |
| return val; |
| } |
| |
| template <typename T> |
| bool DoTileGradient( |
| const int outer_size, |
| const int inner_size, |
| const T* dY, |
| T* dX) { |
| if (inner_size == 1) { |
| const std::array<int, 2> dY_dims = {outer_size, tiles_}; |
| const std::array<int, 2> dX_dims = {outer_size, 1}; |
| math::ReduceSum<T, Context>( |
| 2, dY_dims.data(), dX_dims.data(), T(1), dY, dX, &context_); |
| } else { |
| math::CopyMatrix<T, Context>( |
| outer_size, |
| inner_size, |
| dY, |
| inner_size * tiles_, |
| dX, |
| inner_size, |
| &context_); |
| for (const auto i : c10::irange(outer_size)) { |
| const T* dY_ptr = dY + i * tiles_ * inner_size; |
| T* dX_ptr = dX + i * inner_size; |
| for (const auto j : c10::irange(1, tiles_)) { |
| math::Add<T, Context>( |
| inner_size, dX_ptr, dY_ptr + j * inner_size, dX_ptr, &context_); |
| } |
| } |
| } |
| return true; |
| } |
| |
| std::int32_t tiles_; |
| std::int32_t axis_; |
| |
| Tensor ones_; |
| }; |
| |
| } // namespace caffe2 |
| |
| #endif // CAFFE2_OPERATORS_TILE_OP_H_ |
