1 #include "caffe2/core/operator.h" 5 #include "caffe2/core/logging.h" 6 #include "caffe2/core/net.h" 7 #include "caffe2/core/operator_gradient.h" 8 #include "caffe2/core/tensor.h" 9 #include "caffe2/core/workspace.h" 11 #include "caffe2/proto/caffe2.pb.h" 12 #include "caffe2/utils/proto_utils.h" 13 #include "caffe2/utils/string_utils.h" 17 OperatorBase::OperatorBase(
const OperatorDef& operator_def, Workspace* ws)
18 : operator_def_(operator_def), arg_helper_(operator_def_) {
19 for (
const string& input_str : operator_def_.input()) {
20 auto* blob = ws->GetBlob(input_str);
25 ": Encountered a non-existing input blob: ",
27 inputs_.push_back(blob);
30 GetOperatorLogger()(operator_def_);
32 for (
const string& output_str : operator_def_.output()) {
33 outputs_.push_back(CHECK_NOTNULL(ws->CreateBlob(output_str)));
38 unique_ptr<OperatorBase> TryCreateOperator(
39 const string& key,
const OperatorDef& operator_def, Workspace* ws) {
40 auto type = operator_def.device_option().device_type();
42 gDeviceTypeRegistry()->count(type),
46 OperatorRegistry* registry = gDeviceTypeRegistry()->at(type);
47 VLOG(1) <<
"Creating operator with device type " << type;
49 return registry->Create(key, operator_def, ws);
50 }
catch (
const UnsupportedOperatorFeature& err) {
51 VLOG(1) <<
"Operator " << operator_def.type()
52 <<
" does not support the requested feature. Msg: " << err.what()
53 <<
". Proto is: " << ProtoDebugString(operator_def);
59 unique_ptr<OperatorBase> CreateOperator(
60 const OperatorDef& operator_def, Workspace* ws) {
62 auto* schema = OpSchemaRegistry::Schema(operator_def.type());
65 schema->Verify(operator_def),
66 "Operator def did not pass schema checking: ",
67 ProtoDebugString(operator_def));
72 LOG(ERROR) <<
"Cannot find operator schema for " 73 << operator_def.type()
74 <<
". Will skip schema checking.";
78 if (operator_def.engine().size()) {
79 vector<string> engine_choices = split(
',', operator_def.engine());
80 for (
const string& engine : engine_choices) {
81 string key = operator_def.type() +
"_ENGINE_" + engine;
82 VLOG(1) <<
"Trying to create operator " << operator_def.type()
83 <<
" with engine " << engine;
84 auto op = TryCreateOperator(key, operator_def, ws);
90 VLOG(1) <<
"Operator with engine " << engine
91 <<
" is not available. Using default implementation.";
97 auto op = TryCreateOperator(operator_def.type(), operator_def, ws);
100 "Cannot create operator of type '",
103 DeviceTypeName(operator_def.device_option().device_type()),
104 "'. Verify that implementation for the corresponding device exist. It " 105 "might also happen if the binary is not linked with the operator " 106 "implementation code. If Python frontend is used it might happen if " 107 "dyndep.InitOpsLibrary call is missing. Operator def: ",
108 ProtoDebugString(operator_def));
112 std::map<int32_t, OperatorRegistry*>* gDeviceTypeRegistry() {
113 static std::map<int32_t, OperatorRegistry*> g_device_type_registry;
114 return &g_device_type_registry;
117 CAFFE_DEFINE_REGISTRY(
122 CAFFE_REGISTER_DEVICE_TYPE(DeviceType::CPU, CPUOperatorRegistry);
124 CAFFE_DEFINE_REGISTRY(
125 CUDAOperatorRegistry,
129 CAFFE_REGISTER_DEVICE_TYPE(DeviceType::CUDA, CUDAOperatorRegistry);
131 CAFFE_DEFINE_REGISTRY(
134 const OperatorDef&,
const vector<GradientWrapper>&);
137 const OperatorDef& def,
const vector<GradientWrapper>& g_output) {
138 std::unique_ptr<GradientMakerBase> maker(
139 GradientRegistry()->Create(def.type(), def, g_output));
141 "Gradient maker for operator ", def.type(),
" not implemented.");
144 if (maker->CopyDeviceOption() && def.has_device_option()) {
145 for (OperatorDef& grad_def : meta.ops_) {
146 grad_def.mutable_device_option()->CopyFrom(def.device_option());
150 if (maker->CopyEngine() && def.has_engine()) {
151 for (OperatorDef& grad_def : meta.ops_) {
152 grad_def.set_engine(def.engine());
156 if (maker->CopyArguments() && def.arg_size()) {
157 for (OperatorDef& grad_def : meta.ops_) {
158 for (
auto& arg : def.arg()) {
159 grad_def.add_arg()->CopyFrom(arg);
164 for (
const OperatorDef& grad_def : meta.ops_) {
165 VLOG(1) <<
"Gradient ops: " << ProtoDebugString(grad_def);
169 CAFFE_ENFORCE_EQ(meta.g_input_.size(), def.input_size());
170 VLOG(1) <<
"Gradients:";
174 if (!grad.IsDense() && !grad.IsSparse()) {
175 VLOG(1) <<
"\t [no gradient]";
176 }
else if (grad.IsDense()) {
177 VLOG(1) <<
"\t [dense]" << grad.dense_;
180 grad.indices_.size() && grad.values_.size(),
181 "For sparse gradient, one should set both indices and values. " 182 "Currently we have: (" +
183 grad.indices_ +
", " + grad.values_ +
").");
184 VLOG(1) <<
"\t [sparse] " << grad.indices_ <<
", " << grad.values_;
190 static TensorShapes InferBlobShapesAndTypes(
191 CaffeMap<string, TensorShape>& blob_desc,
192 const vector<std::unique_ptr<NetDef>>& nets) {
193 for (
auto& defptr : nets) {
194 for (
const OperatorDef& op : defptr.get()->op()) {
195 vector<TensorShape> input_desc;
196 for (
const string& in : op.input()) {
197 auto inp_desc = blob_desc.find(in);
198 if (inp_desc == blob_desc.end()) {
200 "Shape and type inference failed, could not find shape for ", in);
202 input_desc.push_back(inp_desc->second);
204 auto op_schema = OpSchemaRegistry::Schema(op.type());
205 if (op_schema ==
nullptr) {
206 CAFFE_THROW(
"Shape inference failed, since no schema for: ", op.type());
209 std::vector<TensorShape> out = op_schema->InferTensor(op, input_desc);
210 if (op.is_gradient_op() && out.size()) {
216 CaffeMap<string, string> grads_to_params =
219 for (
int i = 0; i < out.size(); i++) {
220 if (out[i].unknown_shape()) {
221 std::string gradout = op.output(i);
223 if (grads_to_params.find(gradout) != grads_to_params.end()) {
224 std::string var = grads_to_params[gradout];
225 if (blob_desc.find(var) != blob_desc.end()) {
226 out[i] = blob_desc[var];
233 if (out.size() != op.output_size()) {
235 "Invalid shape inference for operator ",
239 " outputs, but got ",
242 for (
int i = 0; i < out.size(); i++) {
243 blob_desc[op.output(i)] = out[i];
249 for (
auto kv : blob_desc) {
250 TensorShape& tp = kv.second;
251 TensorShape* tpnew = tps.add_shapes();
253 tpnew->set_name(kv.first);
258 TensorShapes InferBlobShapesAndTypesFromWorkspace(
260 const vector<std::unique_ptr<NetDef>>& nets) {
261 CaffeMap<string, TensorShape> blob_desc;
263 const std::vector<string>& ws_blobs = ws->
Blobs();
264 for (
const auto& s : ws_blobs) {
266 TypeCall type_fun = GetTypeCallFunction(b->
meta().
id());
267 ShapeCall shape_fun = GetShapeCallFunction(b->
meta().
id());
271 tp.set_data_type(TypeMetaToDataType(type_fun(b->GetRaw())));
276 auto shape = shape_fun(b->GetRaw(), _shares_data, _capacity);
277 for (
auto d : shape) {
281 tp.set_unknown_shape(
true);
285 return InferBlobShapesAndTypes(blob_desc, nets);
288 TensorShapes InferBlobShapesAndTypesFromMap(
289 const CaffeMap<std::string, std::vector<TIndex>>& blob_dimensions,
290 const vector<std::unique_ptr<NetDef>>& nets) {
291 CaffeMap<string, TensorShape> blob_desc;
293 for (
const auto& blob : blob_dimensions) {
295 for (
auto d : blob.second) {
296 CAFFE_ENFORCE_GT(d, 0);
299 blob_desc[blob.first] = tp;
301 return InferBlobShapesAndTypes(blob_desc, nets);
const Blob * GetBlob(const string &name) const
Gets the blob with the given name as a const pointer.
static CaffeMap< string, string > MatchGradsToParams(const OperatorDef &op)
Returns map that returns the parameters that the gradients are for.
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
vector< string > Blobs() const
Return a list of blob names.
Simple registry implementation in Caffe2 that uses static variables to register object creators durin...
Blob is a general container that hosts a typed pointer.
GradientOpsMeta GetGradientForOp(const OperatorDef &def, const vector< GradientWrapper > &g_output)
Gets the GradientOpsMeta for the given operator def.
const TypeMeta & meta() const
Returns the meta info of the blob.