1 #ifndef SOFTMAX_WITH_LOSS_OP_H_ 2 #define 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 label_prob_mode_(OperatorBase::GetSingleArgument<int>(
"label_prob", 0)),
18 order_(StringToStorageOrder(
19 OperatorBase::GetSingleArgument<string>(
"order",
"NCHW"))),
20 axis_(OperatorBase::GetSingleArgument<int>(
"axis", 1)) {
21 CAFFE_ENFORCE(scale_ >= 0);
23 order_, StorageOrder::NCHW,
"Only NCHW order is supported right now.");
25 USE_OPERATOR_CONTEXT_FUNCTIONS;
27 bool RunOnDevice()
override;
43 template <
typename T,
class Context>
48 scale_(OperatorBase::GetSingleArgument<float>(
"scale", 1.)),
49 label_prob_mode_(OperatorBase::GetSingleArgument<int>(
"label_prob", 0)),
50 order_(StringToStorageOrder(
51 OperatorBase::GetSingleArgument<string>(
"order",
"NCHW"))),
52 only_loss_(OperatorBase::GetSingleArgument<bool>(
"only_loss",
false)),
53 axis_(OperatorBase::GetSingleArgument<int>(
"axis", 1)) {
54 CAFFE_ENFORCE(scale_ >= 0);
56 order_, StorageOrder::NCHW,
"Only NCHW order is supported right now.");
58 USE_OPERATOR_CONTEXT_FUNCTIONS;
60 bool RunOnDevice()
override;
76 #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 ...