/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators/scale_blobs_op.h (1458B)
#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_