Caffe2 - C++ API
A deep learning, cross platform ML framework
module.h
1 #pragma once
2 #include <torch/csrc/autograd/variable.h>
3 #include <torch/csrc/autograd/generated/variable_factories.h>
4 #include <torch/csrc/jit/argument_spec.h>
5 #include <c10/util/Exception.h>
6 #include <torch/csrc/jit/graph_executor.h>
7 #include <torch/csrc/jit/ir.h>
8 #include <torch/csrc/jit/named_value.h>
9 #include <torch/csrc/jit/passes/shape_analysis.h>
10 #include <torch/csrc/jit/source_range.h>
11 
12 #include <torch/csrc/WindowsTorchApiMacro.h>
13 #include <torch/csrc/api/include/torch/ordered_dict.h>
14 #include <torch/csrc/utils/memory.h>
15 
16 #include <ATen/core/function_schema.h>
17 #include <c10/util/ArrayRef.h>
18 #include <c10/util/Optional.h>
19 
20 #include <functional>
21 #include <memory>
22 #include <mutex>
23 #include <ostream>
24 #include <string>
25 #include <unordered_map>
26 #include <vector>
27 
28 // This file contains classes which assist in desugaring Python style
29 // modules and their methods into flattened graphs which don't have any
30 // function calls.
31 
32 namespace torch {
33 namespace jit {
34 namespace script {
35 
36 using ::c10::Argument;
37 using ::c10::FunctionSchema;
38 // Map which stores filename to content.
39 using ExtraFilesMap = std::unordered_map<std::string, std::string>;
40 
41 // A method in a module, e.g. f in:
42 //
43 // class M(ScriptModule):
44 // @script_method
45 // def f(self, x):
46 // ...
47 // Note: because Method/Module are exposed to python these
48 // classes use python method naming conventions
49 
50 struct Module;
51 
52 using ModuleLookup = std::function<std::shared_ptr<Module>(
53  const std::vector<std::string>&)>;
54 
55 struct Method {
56  Method(
57  Module* owner,
58  std::string name,
59  bool optimize,
60  std::shared_ptr<Graph> graph,
61  std::vector<IValue*> initial_members,
62  std::function<void(Method&)> method_creator)
63  : owner_(owner),
64  name_(std::move(name)),
65  graph_(std::move(graph)),
66  optimize(optimize),
67  initial_ivalues_(std::move(initial_members)),
68  method_creator(std::move(method_creator)) {
69  AT_ASSERT(graph_->inputs().size() >= initial_ivalues_.size());
70  int i = graph_->inputs().size() - initial_ivalues_.size();
71  for (auto member : initial_ivalues_) {
72  initial_ivalue_index[member] = i++;
73  }
74  }
75 
76  void run(Stack& stack) {
77  for (auto input : initial_ivalues_) {
78  push(stack, *input);
79  }
80  get_executor().run(stack);
81  }
82 
83  void run(Stack&& stack) {
84  run(stack);
85  }
86 
87  IValue operator()(std::vector<IValue> stack) {
88  checkInputsAgainstSchema(stack);
89  run(stack);
90  return stack.front();
91  }
92 
93  std::shared_ptr<Graph> graph_for(Stack inputs) {
94  for (auto tp : initial_ivalues_) {
95  inputs.emplace_back(*tp);
96  }
97  return get_executor().graphFor(inputs);
98  }
99  TORCH_API std::shared_ptr<Graph> graph() const {
100  return graph_;
101  }
102 
103  TORCH_API const std::string& name() const {
104  return name_;
105  }
106  // emit a function call by inlining the callees Graph into this one
107  // adding any extra parameters necessary to do this call
108 
109  // defined here to keep details of member_input handling confined to this
110  // class
111  Value* emit_call_to(
112  const SourceRange& loc,
113  Method& callee,
115  ArrayRef<NamedValue> kwargs);
116 
117  // if this isn't yet defined, run its method_creator function
118  TORCH_API void ensure_defined();
119 
120  size_t num_inputs() const {
121  return graph()->inputs().size() - initial_ivalues_.size();
122  }
123  TORCH_API Value* get_or_add_parameter(IValue* slot) {
124  AT_ASSERT(slot->isTensor());
125  return get_or_add_attribute(TensorType::get(), slot);
126  }
127 
128  TORCH_API Value* get_or_add_attribute(TypePtr type, IValue* slot) {
129  auto it = initial_ivalue_index.find(slot);
130  if (it != initial_ivalue_index.end()) {
131  return graph()->inputs().at(it->second);
132  }
133  initial_ivalues_.push_back(slot);
134  initial_ivalue_index[slot] = graph()->inputs().size();
135  return graph()->addInput()->setType(type);
136  }
137 
138  std::shared_ptr<Graph> propagate_shapes(
139  std::vector<at::Tensor> inputs,
140  bool with_grad = false) {
141  auto retval = graph_->copy();
142  Stack stack;
143  stack.reserve(inputs.size() + initial_ivalues_.size());
144  for (at::Tensor& i : inputs) {
145  stack.emplace_back(std::move(i));
146  }
147  for (IValue* inp : initial_ivalues_) {
148  stack.push_back(*inp);
149  }
150  const auto size = stack.size();
151  setInputTypes(*retval, ArgumentSpec(with_grad, stack, size));
152  PropagateInputShapes(retval);
153  return retval;
154  }
155 
156  std::shared_ptr<Graph> propagate_and_assign_input_and_output_shapes(
157  std::vector<at::Tensor> inputs,
158  std::vector<at::Tensor> outputs,
159  bool with_grad = false,
160  bool propagate = true) {
161  auto retval = graph_->copy();
162  for (auto inp : initial_ivalues_) {
163  if (inp->isTensor()) {
164  inputs.push_back(inp->toTensor());
165  }
166  }
167  if (propagate) {
168  setInputTypes(
169  *retval,
170  ArgumentSpec(with_grad, fmap<IValue>(inputs), inputs.size()));
171  PropagateInputShapes(retval);
172  }
173  AT_ASSERT(retval->inputs().size() == inputs.size());
174  for (size_t i = 0; i < retval->inputs().size(); ++i) {
175  auto scalar_type = inputs[i].scalar_type();
176  auto sizes = inputs[i].sizes();
177  auto type =
178  torch::jit::CompleteTensorType::create(scalar_type, at::kCPU, sizes);
179  retval->inputs()[i]->setType(type);
180  }
181  at::ArrayRef<Value*> output_values = retval->outputs();
182  // patch this to still work if we are returning a tuple of multiple values
183  if (output_values.at(0)->type()->kind() == TupleType::Kind) {
184  AT_ASSERT(output_values.at(0)->node()->kind() == prim::TupleConstruct);
185  output_values = output_values.at(0)->node()->inputs();
186  }
187  AT_ASSERT(output_values.size() == outputs.size());
188  for (size_t i = 0; i < retval->outputs().size(); ++i) {
189  auto scalar_type = outputs[i].scalar_type();
190  auto sizes = outputs[i].sizes();
191  auto type =
192  torch::jit::CompleteTensorType::create(scalar_type, at::kCPU, sizes);
193  output_values[i]->setType(type);
194  }
195  return retval;
196  }
197 
198  const std::vector<IValue*>& initial_ivalues() const {
199  return initial_ivalues_;
200  }
201 
202  Method& setSchema(FunctionSchema schema_) {
203  schema = make_unique<FunctionSchema>(std::move(schema_));
204  return *this;
205  }
206 
207  TORCH_API const FunctionSchema& getSchema() const {
208  if (schema == nullptr) {
209  schema = make_unique<FunctionSchema>(defaultSchemaFor(*this));
210  }
211  return *schema;
212  }
213 
214  std::string pretty_print_schema() const {
215  AT_ASSERT(schema);
216  std::stringstream ss;
217  ss << *schema;
218  return ss.str();
219  }
220 
221  GraphExecutorState getDebugState() {
222  return get_executor().getDebugState();
223  }
224 
225  void debugDisableAutodiffSubgraphInlining() {
226  return get_executor().debugDisableAutodiffSubgraphInlining();
227  }
228 
229  bool is_optimized() const {
230  return optimize;
231  }
232 
233  // the module that contains this method.
234  Module& owner() const {
235  return *owner_;
236  }
237 
238  void check_single_output() {
239  AT_CHECK(
240  graph()->outputs().size() == 1,
241  "Method (but not graphs in general) require a single output. Use None/Tuple for 0 or 2+ outputs");
242  }
243 
244  private:
245  static FunctionSchema defaultSchemaFor(const Method& method) {
246  std::vector<Argument> args;
247  std::vector<Argument> returns;
248  Graph& g = *method.graph();
249  size_t num_inputs = method.num_inputs();
250  for (size_t i = 0; i < num_inputs; ++i) {
251  const Value* v = g.inputs().at(i);
252  std::string name = v->hasUniqueName() ? v->uniqueNameBase()
253  : ("argument_" + std::to_string(i));
254  args.emplace_back(std::move(name), unshapedType(g.inputs()[i]->type()));
255  }
256  for (size_t i = 0; i < g.outputs().size(); ++i) {
257  returns.emplace_back("", unshapedType(g.outputs()[i]->type()));
258  }
259  return {method.name(), "", std::move(args), std::move(returns)};
260  }
261 
262  GraphExecutor& get_executor() {
263  std::call_once(executor_init, [&] {
264  check_single_output();
265  executor = GraphExecutor(graph(), optimize);
266  });
267  return executor;
268  }
269 
270  void checkInputsAgainstSchema(std::vector<IValue>& inputs) {
271  const auto& schema = getSchema();
272  // Do we have more inputs than the schema accepts?
273  AT_CHECK(
274  inputs.size() <= schema.arguments().size(),
275  "Expected at most ",
276  schema.arguments().size(),
277  " argument(s) for operator '",
278  schema.name(),
279  "', but received ",
280  inputs.size(),
281  " argument(s). Declaration: ",
282  schema);
283 
284  for (size_t pos = 0; pos < schema.arguments().size(); ++pos) {
285  const auto& argument = schema.arguments()[pos];
286  if (pos < inputs.size()) {
287  if (!isSubvalueOf(inputs[pos], argument.type())) {
288  AT_ERROR(
289  "Expected value of type ",
290  *argument.type(),
291  " for argument '",
292  argument.name(),
293  "' in position ",
294  pos,
295  ", but instead got value of type ",
296  attemptToRecoverType(inputs[pos])->str(),
297  ". Declaration: ",
298  schema);
299  }
300  } else if (argument.default_value()) {
301  inputs.push_back(*argument.default_value());
302  } else {
303  AT_ERROR(
304  schema.name(),
305  "() is missing value for argument '",
306  argument.name(),
307  "'. Declaration: ",
308  schema);
309  }
310  }
311  }
312 
313  // Methods are uniqued onwed by a single module. This raw pointer allows
314  // looking up the module.
315  Module* owner_;
316 
317  std::string name_;
318  std::shared_ptr<Graph> graph_; // for debugging and for inlining
319  bool optimize;
320 
321  GraphExecutor executor; // for execution
322  // initial_ivalues are a list of additional arguments appended to graph
323  // that are inputs that come from the members of the Module or its submodules.
324  // each is a pointer to a slot in the module that owns this parameter
325  // parameters and submodules can only be _added_ to script Modules to ensure
326  // these pointers always stay valid
327  std::vector<IValue*> initial_ivalues_;
328 
329  // map from a IValue* in initial_ivalues to the offset it appears at
330  // in graph. used to accelerate get_or_add_parameter
331  std::unordered_map<IValue*, size_t> initial_ivalue_index;
332 
333  // TODO: support that case where we allow _writes_ to parameters from
334  // compiled functions.
335  // This requires more sophisticated tracking of ssa values in Graphs so that
336  // stores to all modules can be lifted to the end of a graph execution.
337  // It also adds more complexity to adding actual module invocations
338  // to the executor, so currently it is not done.
339  // std::vector<at::Tensor*> member_outputs;
340 
341  std::once_flag executor_init;
342 
343  // an optional function that actually creates the method when
344  // emit_call_to(this,...) is first called. this is used by the compiler so
345  // that it can construct methods out of order
346  std::function<void(Method&)> method_creator;
347 
348  // if absent, then we generate a default schema based on the graph
349  // mutable because getSchema caches the default schema if one is requested
350  // before a call to setSchema
351  mutable std::unique_ptr<FunctionSchema> schema;
352 };
353 
354 struct Module;
355 
356 struct NamedModule {
357  std::string name;
358  std::shared_ptr<Module> module;
359 };
360 
361 struct NamedIValue {
362  NamedIValue(std::string name, TypePtr type, IValue ivalue)
363  : name_(name),
364  type(type),
365  ivalue(torch::make_unique<IValue>(std::move(ivalue))) {}
366 
367  IValue* slot() const {
368  return ivalue.get();
369  }
370  const std::string name_;
371  const TypePtr type;
372  std::unique_ptr<IValue> ivalue;
373 };
374 
375 struct Module {
376  TH_DISALLOW_COPY_AND_ASSIGN(Module);
377  Module()
378  : modules("Module"),
379  parameters("Parameter"),
380  attributes("Attributes"),
381  methods("Method"),
382  optimize(true) {}
383 
384  // note this doesn't change the flags of existing methods just ones
385  // added afterward.
386  void set_optimized(bool o) {
387  optimize = o;
388  }
389 
390  bool is_optimized() const {
391  return optimize;
392  }
393 
394  IValue forward(std::vector<IValue> inputs) {
395  return get_method("forward")(std::move(inputs));
396  }
397 
398  void register_buffer(const std::string& name, autograd::Variable v) {
399  if (auto b = attributes.find(name)) {
400  AT_ASSERT(b->type->isSubtypeOf(TensorType::get()));
401  *b->slot() = v;
402  return;
403  }
404  attributes.insert(name, NamedIValue(name, TensorType::get(), std::move(v)));
405  }
406  void register_parameter(
407  const std::string& name,
409  bool is_buffer) {
410  if (is_buffer) {
411  register_buffer(name, std::move(v));
412  return;
413  }
414  if (auto p = parameters.find(name)) {
415  *p->slot() = v;
416  return;
417  }
418  parameters.insert(name, NamedIValue(name, TensorType::get(), std::move(v)));
419  }
420  void register_attribute(
421  const std::string& name,
422  const TypePtr type,
423  IValue ivalue) {
424  attributes.insert(name, NamedIValue(name, type, ivalue));
425  }
426  void register_module(
427  const std::string& name,
428  std::shared_ptr<Module> module) {
429  modules.insert(name, {name, std::move(module)});
430  }
431 
432  Method& create_method(
433  const std::string& name,
434  std::shared_ptr<Graph> graph,
435  std::vector<IValue*> member_inputs) {
436  AT_ASSERT(graph);
437  std::unique_ptr<Method> method(new Method(
438  this,
439  name,
440  optimize,
441  std::move(graph),
442  std::move(member_inputs),
443  nullptr));
444  return *methods.insert(name, std::move(method));
445  }
446 
447  Method& create_method(
448  const std::string& name,
449  std::function<void(Method&)> creator) {
450  std::unique_ptr<Method> method(new Method(
451  this,
452  name,
453  optimize,
454  std::make_shared<Graph>(),
455  {},
456  std::move(creator)));
457  return *methods.insert(name, std::move(method));
458  }
459 
460  IValue* parameter_slot(const std::string& name) const {
461  return parameters[name].slot();
462  }
463 
464  void set_parameter(const std::string& name, at::Tensor v) {
465  *parameter_slot(name) = std::move(v);
466  }
467 
468  autograd::Variable get_parameter(const std::string& name) const {
469  return autograd::as_variable_ref(parameter_slot(name)->toTensor());
470  }
471  autograd::Variable get_buffer(const std::string& name) const {
472  return autograd::as_variable_ref(attributes.find(name)->slot()->toTensor());
473  }
474 
475  // each module owns its method. The reference returned here
476  // is guarenteed to stay valid until this module has been destroyed
477  Method& get_method(const std::string& name) const {
478  return *methods[name];
479  }
480 
481  std::shared_ptr<Module> get_module(const std::string& name) const {
482  return modules[name].module;
483  }
484 
485  const torch::OrderedDict<std::string, NamedModule>& get_modules() const {
486  return modules;
487  }
488  const torch::OrderedDict<std::string, NamedIValue>& get_parameters()
489  const {
490  return parameters;
491  }
492  const torch::OrderedDict<std::string, NamedIValue>& get_attributes()
493  const {
494  return attributes;
495  }
497  const {
498  return methods;
499  }
500 
501  NamedIValue* find_parameter(const std::string& name) {
502  return parameters.find(name);
503  }
504  NamedIValue* find_attribute(const std::string& name) {
505  return attributes.find(name);
506  }
507  NamedIValue* find_buffer(const std::string& name) {
508  auto b = attributes.find(name);
509  if (b && b->type->isSubtypeOf(TensorType::get())) {
510  return b;
511  }
512  return nullptr;
513  }
514  NamedModule* find_module(const std::string& name) {
515  return modules.find(name);
516  }
517  Method* find_method(const std::string& name) {
518  if (auto* pm = methods.find(name)) {
519  return pm->get();
520  }
521  return nullptr;
522  }
523  void apply(std::function<void(Module&)> fn) {
524  for (auto& submod : get_modules()) {
525  submod.value().module->apply(fn);
526  }
527  fn(*this);
528  }
530  void train(bool on = true) {
531  for (auto& submod : get_modules()) {
532  submod->module->train(on);
533  }
534  register_buffer("training", torch::tensor(on ? 1 : 0, at::kLong));
535  }
538  void eval() {
539  train(/*on=*/false);
540  }
542  bool is_training() {
543  if (auto p = find_buffer("training")) {
544  return p->slot()->toTensor().item<int64_t>() == 1;
545  }
546  // We are in training mode by default
547  return true;
548  }
549 
556  TORCH_API void to(
557  at::Device device,
558  at::ScalarType dtype,
559  bool non_blocking = false);
560 
567  TORCH_API void to(at::ScalarType dtype, bool non_blocking = false);
568 
575  TORCH_API void to(at::Device device, bool non_blocking = false);
576 
590  template <typename... Types>
591  IValue run_method(const std::string& method_name, Types&&... args) {
592  return get_method(method_name)({IValue(std::forward<Types>(args))...});
593  }
594 
595  void save(
596  std::ostream& out,
597  const ExtraFilesMap& extra_files = ExtraFilesMap());
598 
599  void save(
600  const std::string& filename,
601  const ExtraFilesMap& extra_files = ExtraFilesMap());
602 
603  void copy_into(
604  ModuleLookup module_lookup,
605  // parameter_remap is needed when a parent module uses a parameter of a
606  // submodule
607  std::unordered_map<IValue*, IValue*>& parameter_remap,
608  std::vector<std::string> names = {}) const {
609  auto curr = module_lookup(names);
610  for (auto& kv : parameters) {
611  curr->register_parameter(
612  kv.key(),
613  kv.value().slot()->toTensor(),
614  /*is_buffer=*/false);
615  parameter_remap[kv.value().slot()] = curr->parameter_slot(kv.key());
616  }
617  for (auto& kv : attributes) {
618  if (!kv.value().type->isSubtypeOf(TensorType::get())) {
619  continue;
620  }
621  curr->register_buffer(
622  kv.key(),
623  kv.value().slot()->toTensor());
624  parameter_remap[kv.value().slot()] = curr->find_buffer(kv.key())->slot();
625  }
626  for (auto& kv : modules) {
627  names.push_back(kv.key());
628  // Submodules must be translated first, otherwise parameter_remap entries
629  // will not be filled in for methods of this module.
630  kv.value().module->copy_into(module_lookup, parameter_remap, names);
631  names.pop_back();
632  }
633  for (auto& kv : methods) {
634  std::vector<IValue*> initial_ivalues;
635  for (auto& p : kv.value()->initial_ivalues()) {
636  initial_ivalues.push_back(parameter_remap.at(p));
637  }
638  curr->create_method(
639  kv.key(), kv.value()->graph()->copy(), initial_ivalues);
640  }
641  }
642 
643  private:
644  void to_impl(
645  const c10::optional<at::Device>& device,
646  const c10::optional<at::ScalarType>& dtype,
647  bool non_blocking);
648 
649  // invariant: to ensure initial_ivalues of Methods stay valid,
650  // it is only legal to _add_ new modules and parameters.
651  // removing them will allow initial_ivalues to point to invalid parameters
652  // no such restriction exists for methods
657  bool optimize;
658 };
659 
660 // returns nullptr and fills in failure_messages if the callee does not
661 // match the functions schema
662 Value* try_emit_call_to(
663  Graph& graph,
664  const SourceRange& loc,
665  Method& callee,
668  ArrayRef<NamedValue> kwargs,
669  std::stringstream& failure_messages,
670  // when callee uses no parameters (e.g. it is a function in a compilation
671  // unit, and not a method), then nullptr can be passed as caller.
672  Method* caller,
673  bool conv_tensors_to_nums);
674 } // namespace script
675 } // namespace jit
676 } // namespace torch
bool is_training()
True if the module is in training mode.
Definition: module.h:542
Represents a a compute device on which a tensor is located.
Definition: Device.h:30
constexpr size_t size() const
size - Get the array size.
Definition: ArrayRef.h:138
void train(bool on=true)
Enables "training" mode.
Definition: module.h:530
Variable A Variable augments a Tensor with the ability to interact in our autograd machinery...
Definition: variable.h:85
Definition: jit_type.h:17
ArrayRef - Represent a constant reference to an array (0 or more elements consecutively in memory)...
Definition: ArrayRef.h:41
AT_CPP14_CONSTEXPR const T & at(size_t Index) const
Vector compatibility.
Definition: ArrayRef.h:186
IValue run_method(const std::string &method_name, Types &&...args)
Run a method from this module.
Definition: module.h:591
void eval()
Calls train(false) to enable "eval" mode.
Definition: module.h:538
An ordered dictionary implementation, akin to Python&#39;s OrderedDict.
Definition: ordered_dict.h:16