#ifndef CAFFE2_OPERATORS_SCALE_BLOBS_OP_H_ #define CAFFE2_OPERATORS_SCALE_BLOBS_OP_H_ #include "caffe2/core/context.h" #include "caffe2/core/operator.h" #include "caffe2/utils/math.h" namespace caffe2 { template class ScaleBlobsOp final : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; template explicit ScaleBlobsOp(Args&&... args) : Operator(std::forward(args)...), OP_SINGLE_ARG(float, "scale", scale_, 1.0f) {} template bool DoRunWithType() { int batchSize = InputSize(); for (int i = 0; i < batchSize; ++i) { const auto& X = Input(i); auto* Y = Output(i, X.sizes(), at::dtype()); math::Scale( X.numel(), scale_, X.template data(), Y->template mutable_data(), &context_); } return true; } bool RunOnDevice() override { for (int i = 0; i < InputSize(); ++i) { auto& input = this->template Input(i, CPU); auto* output = this->template Output(i, CPU); output->ResizeLike(input); } return DispatchHelper>::call(this, Input(0)); } private: const float scale_; Tensor blobSizes_; Tensor inputs_; Tensor outputs_; Tensor hostBlobSizes_; Tensor hostInputs_; Tensor hostOutputs_; }; } // namespace caffe2 #endif // CAFFE2_OPERATORS_SCALE_BLOBS_OP_H_