#ifndef CAFFE2_CORE_OPERATOR_GRADIENT_H_
|
#define CAFFE2_CORE_OPERATOR_GRADIENT_H_
|
|
#include "c10/util/Registry.h"
|
#include "caffe2/core/operator_schema.h"
|
#include "caffe2/proto/caffe2_pb.h"
|
#include "caffe2/utils/proto_utils.h"
|
|
namespace caffe2 {
|
|
/* @brief A struct that abstracts on top of dense and sparse blobs.
|
*
|
* For a dense blob, its gradient name should be written into dense_, and for
|
* a sparse blob, its gradient name should be written into indice_ for
|
* the sparse indices and value_ for the values.
|
*/
|
struct CAFFE2_API GradientWrapper {
|
string dense_;
|
string indices_;
|
string values_;
|
|
inline bool IsDense() const {
|
return (dense_.size() != 0);
|
}
|
inline bool IsSparse() const {
|
return (indices_.size() != 0 || values_.size() != 0);
|
}
|
inline bool IsEmpty() const {
|
return (!IsDense() && !IsSparse());
|
}
|
};
|
|
/**
|
* A struct that holds the gradient operators and related gradient maps.
|
*/
|
struct CAFFE2_API GradientOpsMeta {
|
vector<OperatorDef> ops_;
|
vector<GradientWrapper> g_input_;
|
|
GradientOpsMeta() {}
|
GradientOpsMeta(
|
const vector<OperatorDef>& ops,
|
const vector<GradientWrapper>& v)
|
: ops_(ops), g_input_(v) {}
|
};
|
|
class CAFFE2_API GradientMakerBase {
|
public:
|
GradientMakerBase(
|
const OperatorDef& def,
|
const vector<GradientWrapper>& g_output)
|
: def_(def), g_output_(g_output), g_input_(def.input_size()){};
|
virtual ~GradientMakerBase() {}
|
virtual bool CopyDeviceOption() const {
|
return true;
|
}
|
virtual bool CopyEngine() const {
|
return true;
|
}
|
virtual bool CopyArguments() const {
|
return true;
|
}
|
|
virtual void VerifyOp() const {
|
auto* schema = OpSchemaRegistry::Schema(def_.type());
|
if (schema) {
|
CAFFE_ENFORCE(
|
schema->Verify(def_),
|
"(GradientMaker) Operator def did not pass schema checking: ",
|
ProtoDebugString(def_));
|
}
|
}
|
|
/**
|
* @brief Returns the gradient ops meta.
|
*
|
* If your gradient op generator only use standard input and output
|
* manipulations, you can simply implement GetGradientDefs() that
|
* returns vector<OperatorDef>. In that, you can call GI, GI_V and GI_I
|
* that will automatically create the gradient registration for you.
|
*
|
* If you need to do custom gradient name registration, overload this
|
* function directly.
|
*/
|
virtual GradientOpsMeta Get() {
|
VerifyOp();
|
vector<OperatorDef> new_defs = GetGradientDefs();
|
for (auto& opdef : new_defs) {
|
opdef.set_is_gradient_op(true);
|
}
|
return GradientOpsMeta(new_defs, g_input_);
|
};
|
|
const OperatorDef& Def() const {
|
return def_;
|
}
|
|
protected:
|
virtual vector<OperatorDef> GetGradientDefs() {
|
CAFFE_NOT_IMPLEMENTED;
|
}
|
|
// Helper functions to return names for the gradient computation.
|
// I(idx), O(idx): return the input and output names.
|
// GO(idx): return the name of the gradient for output idx.
|
// GI(idx), GI_I(idx), GI_V(idx): return the name of the gradient for
|
// input idx, and also registers that name into the gradient
|
// registry to be returned.
|
string I(const int i) {
|
CAFFE_ENFORCE((i >= 0) && (i < def_.input().size()));
|
return def_.input(i);
|
}
|
string O(const int i) {
|
CAFFE_ENFORCE((i >= 0) && (i < def_.output().size()));
|
return def_.output(i);
|
}
|
string GI(const int i) {
|
CAFFE_ENFORCE(
|
!g_input_.at(i).IsSparse(),
|
"Input ",
|
def_.input(i),
|
" already set to sparse.");
|
g_input_.at(i).dense_ = GradientName(def_.input(i));
|
return GradientName(def_.input(i));
|
}
|
string GI_I(const int i) {
|
CAFFE_ENFORCE(
|
!g_input_.at(i).IsDense(),
|
"Input ",
|
def_.input(i),
|
" already set to dense.");
|
g_input_.at(i).indices_ = GradientSliceIndices(def_.input(i));
|
return GradientSliceIndices(def_.input(i));
|
}
|
string GI_V(const int i) {
|
CAFFE_ENFORCE(
|
!g_input_.at(i).IsDense(),
|
"Input ",
|
def_.input(i),
|
" already set to dense.");
|
g_input_.at(i).values_ = GradientSliceValues(def_.input(i));
|
return GradientSliceValues(def_.input(i));
|
}
|
string GO(const int i) {
|
CAFFE_ENFORCE(
|
g_output_.at(i).IsDense(),
|
"Gradient of output ",
|
def_.output(i),
|
(g_output_.at(i).IsSparse() ? " is sparse (expected dense)."
|
: " is not provided!"));
|
return g_output_.at(i).dense_;
|
}
|
string GO_I(const int i) {
|
CAFFE_ENFORCE(
|
g_output_.at(i).IsSparse(),
|
"Gradient of output ",
|
def_.output(i),
|
(g_output_.at(i).IsDense() ? " is dense (expected sparse)."
|
: " is not provided!"));
|
return g_output_.at(i).indices_;
|
}
|
string GO_V(const int i) {
|
CAFFE_ENFORCE(
|
g_output_.at(i).IsSparse(),
|
"Gradient of output ",
|
def_.output(i),
|
(g_output_.at(i).IsDense() ? " is dense (expected sparse)."
|
: " is not provided!"));
|
return g_output_.at(i).values_;
|
}
|
const GradientWrapper& GradOut(int i) {
|
return g_output_.at(i);
|
}
|
|
// Function to add a gradient pair to map.
|
void SetDense(const int i, const string& name) {
|
CAFFE_ENFORCE(
|
!g_input_.at(i).IsSparse(),
|
"Input ",
|
def_.input(i),
|
" already set to sparse.");
|
g_input_.at(i).dense_ = name;
|
}
|
void SetSparse(const int i, const string& indices, const string& values) {
|
CAFFE_ENFORCE(
|
!g_input_.at(i).IsDense(),
|
"Input ",
|
def_.input(i),
|
" already set to dense.");
|
g_input_.at(i).indices_ = indices;
|
g_input_.at(i).values_ = values;
|
}
|
|
/**
|
* @brief a helper function to allow one to create one single operator
|
* def, which is usually the case for many simple operators.
|
*/
|
template <class... Args>
|
inline static vector<OperatorDef> SingleGradientDef(const Args&... args) {
|
return vector<OperatorDef>{CreateOperatorDef(args...)};
|
}
|
|
public:
|
/**
|
* Returns map that returns the parameters that the gradients are for.
|
*/
|
static CaffeMap<string, string> MatchGradsToParams(const OperatorDef& op) {
|
// NOTE: how to go beyond string-matching?
|
CaffeMap<string, string> m;
|
for (auto& out : op.output()) {
|
if (IsGradientBlob(out)) {
|
m[out] = out.substr(0, out.length() - 5);
|
}
|
}
|
return m;
|
}
|
|
private:
|
// Utility functions for gradient name computation. We don't expose them
|
// in order to discourage the use of such names explicitly.
|
static string GradientName(const string& name) {
|
return name + "_grad";
|
}
|
|
static bool IsGradientBlob(const string& name) {
|
return name.length() > 5 && name.find("_grad") == name.length() - 5;
|
}
|
|
static string GradientNameToParam(const string& name) {
|
CHECK(IsGradientBlob(name));
|
return name.substr(0, name.length() - 5);
|
}
|
|
static string GradientSliceIndices(const string& name) {
|
return name + "_grad_indices";
|
}
|
|
static string GradientSliceValues(const string& name) {
|
return name + "_grad_values";
|
}
|
|
protected:
|
// We make the member variables protected in case someone wants to write
|
// a fully custom Get() function.
|
const OperatorDef& def_;
|
const vector<GradientWrapper>& g_output_;
|
vector<GradientWrapper> g_input_;
|
};
|
|
/**
|
* @brief A helper class to indicate that the operator does not need gradient
|
* computation.
|
*
|
* Use the macro NO_GRADIENT to register operators that do not have gradients.
|
* Note that this is different fron SHOULD_NOT_DO_GRADIENT: the latter means
|
* that the gradient computation should not flow through it at all, and throws
|
* an error if it is called.
|
*/
|
class CAFFE2_API NoGradient : public GradientMakerBase {
|
using GradientMakerBase::GradientMakerBase;
|
vector<OperatorDef> GetGradientDefs() override {
|
return vector<OperatorDef>();
|
}
|
};
|
|
/**
|
* @brief A helper class to indicate that the operator should have no gradient.
|
*
|
* This is used when the operator definition is designed to not have a gradient.
|
* Calling a gradient on this operator def will cause Caffe2 to quit.
|
*/
|
struct ThrowInTheTowelIfGradientIsCalled : public GradientMakerBase {
|
using GradientMakerBase::GradientMakerBase;
|
GradientOpsMeta Get() override {
|
CAFFE_THROW("One should not call gradient for operator ", def_.type(), ".");
|
}
|
};
|
|
/**
|
* @brief A helper class to indicate that the gradient mechanism is not ready.
|
*
|
* This should only be used sparsely when the gradient does exist, but we have
|
* not implemented it yet and are using this as a lazy excuse. Eventually, a
|
* gradient operator should be implemented.
|
*/
|
struct GradientNotImplementedYet : public GradientMakerBase {
|
using GradientMakerBase::GradientMakerBase;
|
GradientOpsMeta Get() override {
|
CAFFE_THROW(
|
"Operator ",
|
def_.type(),
|
" should have a gradient but is not implemented yet.");
|
}
|
};
|
|
C10_DECLARE_REGISTRY(
|
GradientRegistry,
|
GradientMakerBase,
|
const OperatorDef&,
|
const vector<GradientWrapper>&);
|
|
#ifdef CAFFE2_NO_GRADIENT_OPS
|
|
#define REGISTER_GRADIENT(name, ...) /* No gradients. */
|
#define REGISTER_GRADIENT_STR(str_name, ...) /* No gradients. */
|
|
#else
|
|
#define REGISTER_GRADIENT(name, ...) \
|
C10_REGISTER_CLASS(GradientRegistry, name, __VA_ARGS__)
|
#define REGISTER_GRADIENT_STR(str_name, ...) \
|
C10_REGISTER_TYPED_CLASS(GradientRegistry, str_name, __VA_ARGS__)
|
|
#endif
|
|
// NO_GRADIENT means that the operator does not need any gradient computation.
|
#define NO_GRADIENT(name) REGISTER_GRADIENT(name, NoGradient)
|
|
// SHOULD_NOT_DO_GRADIENT means that the operator is not designed to have
|
// gradient operators. If you attempt to call the gradient, a log fatal will
|
// occur.
|
#define SHOULD_NOT_DO_GRADIENT(name) \
|
REGISTER_GRADIENT(name, ThrowInTheTowelIfGradientIsCalled)
|
|
#define GRADIENT_NOT_IMPLEMENTED_YET(name) \
|
REGISTER_GRADIENT(name, GradientNotImplementedYet)
|
|
/**
|
* @brief Gets the GradientOpsMeta for the given operator def.
|
*/
|
CAFFE2_API GradientOpsMeta GetGradientForOp(
|
const OperatorDef& def,
|
const vector<GradientWrapper>& g_output);
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_CORE_OPERATOR_GRADIENT_H_
|
