1 #ifndef SPATIAL_SOFTMAX_WITH_LOSS_OP_H_ 2 #define SPATIAL_SOFTMAX_WITH_LOSS_OP_H_ 4 #include "caffe2/core/context.h" 5 #include "caffe2/core/logging.h" 6 #include "caffe2/core/operator.h" 7 #include "caffe2/utils/math.h" 11 template <
typename T,
class Context>
16 scale_(OperatorBase::GetSingleArgument<float>(
"scale", 1.)),
17 order_(StringToStorageOrder(
18 OperatorBase::GetSingleArgument<string>(
"order",
"NCHW"))) {
19 CAFFE_ENFORCE(scale_ >= 0);
21 order_, StorageOrder::NCHW,
"Only NCHW order is supported right now.");
23 USE_OPERATOR_CONTEXT_FUNCTIONS;
25 bool RunOnDevice()
override;
39 template <
typename T,
class Context>
44 scale_(OperatorBase::GetSingleArgument<float>(
"scale", 1.)),
45 order_(StringToStorageOrder(
46 OperatorBase::GetSingleArgument<string>(
"order",
"NCHW"))),
47 only_loss_(OperatorBase::GetSingleArgument<bool>(
"only_loss",
false)) {
48 CAFFE_ENFORCE(scale_ >= 0);
50 order_, StorageOrder::NCHW,
"Only NCHW order is supported right now.");
52 USE_OPERATOR_CONTEXT_FUNCTIONS;
54 bool RunOnDevice()
override;
68 #endif // SOFTMAX_WITH_LOSS_OP_H_
Tensor is the basic class in Caffe2 that stores a contiguous memory with its shape information...
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...