@@ -29,7 +29,7 @@ state in a member variable.
29
29
#include <vector>
30
30
31
31
template <class T>
32
- struct MyStackClass : torch::jit:: CustomClassHolder {
32
+ struct MyStackClass : torch::CustomClassHolder {
33
33
std::vector<T> stack_;
34
34
MyStackClass(std::vector<T> init) : stack_(init.begin(), init.end()) {}
35
35
@@ -63,7 +63,7 @@ There are several things to note:
63
63
is to ensure consistent lifetime management of the object instances between languages
64
64
(C++, Python and TorchScript).
65
65
- The second thing to notice is that the user-defined class must inherit from
66
- ``torch::jit:: CustomClassHolder ``. This ensures that everything is set up to handle
66
+ ``torch::CustomClassHolder ``. This ensures that everything is set up to handle
67
67
the lifetime management system previously mentioned.
68
68
69
69
Now let's take a look at how we will make this class visible to TorchScript, a process called
@@ -73,24 +73,25 @@ Now let's take a look at how we will make this class visible to TorchScript, a p
73
73
74
74
// Notice a few things:
75
75
// - We pass the class to be registered as a template parameter to
76
- // `torch::jit:: class_`. In this instance, we've passed the
76
+ // `torch::class_`. In this instance, we've passed the
77
77
// specialization of the MyStackClass class ``MyStackClass<std::string>``.
78
78
// In general, you cannot register a non-specialized template
79
79
// class. For non-templated classes, you can just pass the
80
80
// class name directly as the template parameter.
81
- // - The single parameter to ``torch::jit::class_()`` is a
82
- // string indicating the name of the class. This is the name
83
- // the class will appear as in both Python and TorchScript.
84
- // For example, our MyStackClass class would appear as ``torch.classes.MyStackClass``.
81
+ // - The arguments passed to the constructor make up the "qualified name"
82
+ // of the class. In this case, the registered class will appear in
83
+ // Python and C++ as `torch.classes.my_classes.MyStackClass`. We call
84
+ // the first argument the "namespace" and the second argument the
85
+ // actual class name.
85
86
static auto testStack =
86
- torch::jit:: class_<MyStackClass<std::string>>("MyStackClass")
87
+ torch::class_<MyStackClass<std::string>>("my_classes", "MyStackClass")
87
88
// The following line registers the contructor of our MyStackClass
88
89
// class that takes a single `std::vector<std::string>` argument,
89
90
// i.e. it exposes the C++ method `MyStackClass(std::vector<T> init)`.
90
91
// Currently, we do not support registering overloaded
91
92
// constructors, so for now you can only `def()` one instance of
92
- // `torch::jit:: init`.
93
- .def(torch::jit:: init<std::vector<std::string>>())
93
+ // `torch::init`.
94
+ .def(torch::init<std::vector<std::string>>())
94
95
// The next line registers a stateless (i.e. no captures) C++ lambda
95
96
// function as a method. Note that a lambda function must take a
96
97
// `c10::intrusive_ptr<YourClass>` (or some const/ref version of that)
@@ -99,7 +100,7 @@ Now let's take a look at how we will make this class visible to TorchScript, a p
99
100
return self->stack_.back();
100
101
})
101
102
// The following four lines expose methods of the MyStackClass<std::string>
102
- // class as-is. `torch::jit:: class_` will automatically examine the
103
+ // class as-is. `torch::class_` will automatically examine the
103
104
// argument and return types of the passed-in method pointers and
104
105
// expose these to Python and TorchScript accordingly. Finally, notice
105
106
// that we must take the *address* of the fully-qualified method name,
@@ -217,7 +218,7 @@ demonstrates that:
217
218
#
218
219
# This instantiation will invoke the MyStackClass(std::vector<T> init) constructor
219
220
# we registered earlier
220
- s = torch.classes.MyStackClass([" foo" , " bar" ])
221
+ s = torch.classes.my_classes. MyStackClass([" foo" , " bar" ])
221
222
222
223
# We can call methods in Python
223
224
s.push(" pushed" )
@@ -233,16 +234,16 @@ demonstrates that:
233
234
# For now, we need to assign the class's type to a local in order to
234
235
# annotate the type on the TorchScript function. This may change
235
236
# in the future.
236
- MyStackClass = torch.classes.MyStackClass
237
+ MyStackClass = torch.classes.my_classes. MyStackClass
237
238
238
239
@torch.jit.script
239
240
def do_stacks(s : MyStackClass): # We can pass a custom class instance to TorchScript
240
- s2 = torch.classes.MyStackClass([" hi" , " mom" ]) # We can instantiate the class
241
+ s2 = torch.classes.my_classes. MyStackClass([" hi" , " mom" ]) # We can instantiate the class
241
242
s2.merge(s) # We can call a method on the class
242
243
return s2.clone (), s2.top () # We can also return instances of the class
243
244
# from TorchScript function/methods
244
245
245
- stack, top = do_stacks(torch.classes.MyStackClass([" wow" ]))
246
+ stack, top = do_stacks(torch.classes.my_classes. MyStackClass([" wow" ]))
246
247
assert top == " wow"
247
248
for expected in [" wow" , " mom" , " hi" ]:
248
249
assert stack.pop () == expected
@@ -265,7 +266,7 @@ instantiates and calls a method on our MyStackClass class:
265
266
super().__init__()
266
267
267
268
def forward(self, s : str) -> str:
268
- stack = torch.classes.MyStackClass(["hi", "mom"])
269
+ stack = torch.classes.my_classes. MyStackClass(["hi", "mom"])
269
270
return stack.pop() + s
270
271
271
272
scripted_foo = torch.jit.script(Foo())
@@ -307,7 +308,7 @@ Let's populate ``infer.cpp`` with the following:
307
308
# include <memory>
308
309
309
310
int main(int argc, const char* argv[]) {
310
- torch::jit:: script::Module module;
311
+ torch::script::Module module;
311
312
try {
312
313
// Deserialize the ScriptModule from a file using torch::jit::load ().
313
314
module = torch::jit::load(" foo.pt" );
@@ -394,6 +395,31 @@ And now we can run our exciting C++ binary:
394
395
395
396
Incredible!
396
397
398
+ Moving Custom Classes To/From IValues
399
+ -------------------------------------
400
+
401
+ It's also possible that you may need to move custom classes into or out of
402
+ ` ` IValue` ` s, such as when you take or return ` ` IValue` ` s from TorchScript methods
403
+ or you want to instantiate a custom class attribute in C++. For creating an
404
+ ` ` IValue` ` from a custom C++ class instance:
405
+
406
+ - ` ` torch::make_custom_class<T>()` ` provides an API similar to c10::intrusive_ptr<T>
407
+ in that it will take whatever set of arguments you provide to it, call the constructor
408
+ of T that matches that set of arguments, and wrap that instance up and return it.
409
+ However, instead of returning just a pointer to a custom class object, it returns
410
+ an ` ` IValue` ` wrapping the object. You can then pass this ` ` IValue` ` directly to
411
+ TorchScript.
412
+ - In the event that you already have an ` ` intrusive_ptr` ` pointing to your class, you
413
+ can directly construct an IValue from it using the constructor ` ` IValue(intrusive_ptr<T>)` ` .
414
+
415
+ For converting ` ` IValue` ` s back to custom classes:
416
+
417
+ - ` ` IValue::toCustomClass<T>()` ` will return an ` ` intrusive_ptr<T>` ` pointing to the
418
+ custom class that the ` ` IValue` ` contains. Internally, this function is checking
419
+ that ` ` T` ` is registered as a custom class and that the ` ` IValue` ` does in fact contain
420
+ a custom class. You can check whether the ` ` IValue` ` contains a custom class manually by
421
+ calling ` ` isCustomClass()` ` .
422
+
397
423
Defining Serialization/Deserialization Methods for Custom C++ Classes
398
424
---------------------------------------------------------------------
399
425
@@ -410,7 +436,7 @@ an attribute, you'll get the following error:
410
436
class Foo(torch.nn.Module):
411
437
def __init__(self):
412
438
super().__init__()
413
- self.stack = torch.classes.MyStackClass(["just", "testing"])
439
+ self.stack = torch.classes.my_classes. MyStackClass(["just", "testing"])
414
440
415
441
def forward(self, s : str) -> str:
416
442
return self.stack.pop() + s
@@ -422,7 +448,7 @@ an attribute, you'll get the following error:
422
448
.. code-block:: shell
423
449
424
450
$ python export_attr.py
425
- RuntimeError: Cannot serialize custom bound C++ class __torch__.torch.classes.MyStackClass. Please define serialization methods via torch::jit::pickle_ for this class. (pushIValueImpl at ../torch/csrc/jit/pickler.cpp:128)
451
+ RuntimeError: Cannot serialize custom bound C++ class __torch__.torch.classes.my_classes. MyStackClass. Please define serialization methods via def_pickle for this class. (pushIValueImpl at ../torch/csrc/jit/pickler.cpp:128)
426
452
427
453
This is because TorchScript cannot automatically figure out what information
428
454
save from your C++ class. You must specify that manually. The way to do that
@@ -441,8 +467,8 @@ Here is an example of how we can update the registration code for our
441
467
.. code-block:: cpp
442
468
443
469
static auto testStack =
444
- torch::jit:: class_<MyStackClass<std::string>>("MyStackClass")
445
- .def(torch::jit:: init<std::vector<std::string>>())
470
+ torch::class_<MyStackClass<std::string>>("my_classes", "MyStackClass")
471
+ .def(torch::init<std::vector<std::string>>())
446
472
.def("top", [](const c10::intrusive_ptr<MyStackClass<std::string>>& self) {
447
473
return self->stack_.back();
448
474
})
@@ -503,7 +529,7 @@ now run successfully:
503
529
class Foo(torch.nn.Module):
504
530
def __init__(self):
505
531
super().__init__()
506
- self.stack = torch.classes.MyStackClass(["just", "testing"])
532
+ self.stack = torch.classes.my_classes. MyStackClass(["just", "testing"])
507
533
508
534
def forward(self, s : str) -> str:
509
535
return self.stack.pop() + s
@@ -537,7 +563,7 @@ example of how to do that:
537
563
static auto instance_registry = torch::RegisterOperators().op(
538
564
torch::RegisterOperators::options()
539
565
.schema(
540
- "foo::manipulate_instance(__torch__.torch.classes.MyStackClass x) -> __torch__.torch.classes.MyStackClass Y")
566
+ "foo::manipulate_instance(__torch__.torch.classes.my_classes. MyStackClass x) -> __torch__.torch.classes.my_classes .MyStackClass Y")
541
567
.catchAllKernel<decltype(manipulate_instance), &manipulate_instance>());
542
568
543
569
Refer to the ` custom op tutorial < https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html> ` _
@@ -550,7 +576,7 @@ Once this is done, you can use the op like the following example:
550
576
class TryCustomOp(torch.nn.Module):
551
577
def __init__(self):
552
578
super(TryCustomOp, self).__init__()
553
- self.f = torch.classes.MyStackClass(["foo", "bar"])
579
+ self.f = torch.classes.my_classes. MyStackClass(["foo", "bar"])
554
580
555
581
def forward(self):
556
582
return torch.ops.foo.manipulate_instance(self.f)
0 commit comments