Make torch_script_custom_classes tutorial runnable

I also fixed some warnings in the tutorial, and fixed some minor bitrot
(e.g., torch::script::Module to torch::jit::Module)

I also added some missing quotes around some bash expansions.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

Author: ezyang

.gitignore

@@ -4,6 +4,7 @@ intermediate
 advanced
 pytorch_basics
 recipes
+prototype
 
 #data things
 _data/
@@ -117,3 +118,6 @@ ENV/
 .DS_Store
 cleanup.sh
 *.swp
+
+# PyTorch things
+*.pt

advanced_source/torch_script_custom_classes.rst

@@ -0,0 +1,15 @@
+cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
+project(infer)
+
+find_package(Torch REQUIRED)
+
+add_subdirectory(custom_class_project)
+
+# Define our library target
+add_executable(infer infer.cpp)
+set(CMAKE_CXX_STANDARD 14)
+# Link against LibTorch
+target_link_libraries(infer "${TORCH_LIBRARIES}")
+# This is where we link in our libcustom_class code, making our
+# custom class available in our binary.
+target_link_libraries(infer -Wl,--no-as-needed custom_class)

advanced_source/torch_script_custom_classes/CMakeLists.txt

@@ -0,0 +1,10 @@
+cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
+project(custom_class)
+
+find_package(Torch REQUIRED)
+
+# Define our library target
+add_library(custom_class SHARED class.cpp)
+set(CMAKE_CXX_STANDARD 14)
+# Link against LibTorch
+target_link_libraries(custom_class "${TORCH_LIBRARIES}")

advanced_source/torch_script_custom_classes/custom_class_project/CMakeLists.txt

@@ -0,0 +1,145 @@
+// BEGIN class
+// This header is all you need to do the C++ portions of this
+// tutorial
+#include <torch/script.h>
+// This header is what defines the custom class registration
+// behavior specifically. script.h already includes this, but
+// we include it here so you know it exists in case you want
+// to look at the API or implementation.
+#include <torch/custom_class.h>
+
+#include <string>
+#include <vector>
+
+template <class T>
+struct MyStackClass : torch::CustomClassHolder {
+  std::vector<T> stack_;
+  MyStackClass(std::vector<T> init) : stack_(init.begin(), init.end()) {}
+
+  void push(T x) {
+    stack_.push_back(x);
+  }
+  T pop() {
+    auto val = stack_.back();
+    stack_.pop_back();
+    return val;
+  }
+
+  c10::intrusive_ptr<MyStackClass> clone() const {
+    return c10::make_intrusive<MyStackClass>(stack_);
+  }
+
+  void merge(const c10::intrusive_ptr<MyStackClass>& c) {
+    for (auto& elem : c->stack_) {
+      push(elem);
+    }
+  }
+};
+// END class
+
+#ifdef NO_PICKLE
+
+// BEGIN binding
+// Notice a few things:
+// - We pass the class to be registered as a template parameter to
+//   `torch::class_`. In this instance, we've passed the
+//   specialization of the MyStackClass class ``MyStackClass<std::string>``.
+//   In general, you cannot register a non-specialized template
+//   class. For non-templated classes, you can just pass the
+//   class name directly as the template parameter.
+// - The arguments passed to the constructor make up the "qualified name"
+//   of the class. In this case, the registered class will appear in
+//   Python and C++ as `torch.classes.my_classes.MyStackClass`. We call
+//   the first argument the "namespace" and the second argument the
+//   actual class name.
+static auto testStack =
+  torch::class_<MyStackClass<std::string>>("my_classes", "MyStackClass")
+      // The following line registers the contructor of our MyStackClass
+      // class that takes a single `std::vector<std::string>` argument,
+      // i.e. it exposes the C++ method `MyStackClass(std::vector<T> init)`.
+      // Currently, we do not support registering overloaded
+      // constructors, so for now you can only `def()` one instance of
+      // `torch::init`.
+      .def(torch::init<std::vector<std::string>>())
+      // The next line registers a stateless (i.e. no captures) C++ lambda
+      // function as a method. Note that a lambda function must take a
+      // `c10::intrusive_ptr<YourClass>` (or some const/ref version of that)
+      // as the first argument. Other arguments can be whatever you want.
+      .def("top", [](const c10::intrusive_ptr<MyStackClass<std::string>>& self) {
+        return self->stack_.back();
+      })
+      // The following four lines expose methods of the MyStackClass<std::string>
+      // class as-is. `torch::class_` will automatically examine the
+      // argument and return types of the passed-in method pointers and
+      // expose these to Python and TorchScript accordingly. Finally, notice
+      // that we must take the *address* of the fully-qualified method name,
+      // i.e. use the unary `&` operator, due to C++ typing rules.
+      .def("push", &MyStackClass<std::string>::push)
+      .def("pop", &MyStackClass<std::string>::pop)
+      .def("clone", &MyStackClass<std::string>::clone)
+      .def("merge", &MyStackClass<std::string>::merge);
+// END binding
+
+#else
+
+// BEGIN pickle_binding
+static auto testStack =
+  torch::class_<MyStackClass<std::string>>("my_classes", "MyStackClass")
+      .def(torch::init<std::vector<std::string>>())
+      .def("top", [](const c10::intrusive_ptr<MyStackClass<std::string>>& self) {
+        return self->stack_.back();
+      })
+      .def("push", &MyStackClass<std::string>::push)
+      .def("pop", &MyStackClass<std::string>::pop)
+      .def("clone", &MyStackClass<std::string>::clone)
+      .def("merge", &MyStackClass<std::string>::merge)
+      // class_<>::def_pickle allows you to define the serialization
+      // and deserialization methods for your C++ class.
+      // Currently, we only support passing stateless lambda functions
+      // as arguments to def_pickle
+      .def_pickle(
+            // __getstate__
+            // This function defines what data structure should be produced
+            // when we serialize an instance of this class. The function
+            // must take a single `self` argument, which is an intrusive_ptr
+            // to the instance of the object. The function can return
+            // any type that is supported as a return value of the TorchScript
+            // custom operator API. In this instance, we've chosen to return
+            // a std::vector<std::string> as the salient data to preserve
+            // from the class.
+            [](const c10::intrusive_ptr<MyStackClass<std::string>>& self)
+                -> std::vector<std::string> {
+              return self->stack_;
+            },
+            // __setstate__
+            // This function defines how to create a new instance of the C++
+            // class when we are deserializing. The function must take a
+            // single argument of the same type as the return value of
+            // `__getstate__`. The function must return an intrusive_ptr
+            // to a new instance of the C++ class, initialized however
+            // you would like given the serialized state.
+            [](std::vector<std::string> state)
+                -> c10::intrusive_ptr<MyStackClass<std::string>> {
+              // A convenient way to instantiate an object and get an
+              // intrusive_ptr to it is via `make_intrusive`. We use
+              // that here to allocate an instance of MyStackClass<std::string>
+              // and call the single-argument std::vector<std::string>
+              // constructor with the serialized state.
+              return c10::make_intrusive<MyStackClass<std::string>>(std::move(state));
+            });
+// END pickle_binding
+
+// BEGIN free_function
+c10::intrusive_ptr<MyStackClass<std::string>> manipulate_instance(const c10::intrusive_ptr<MyStackClass<std::string>>& instance) {
+  instance->pop();
+  return instance;
+}
+
+static auto instance_registry = torch::RegisterOperators().op(
+torch::RegisterOperators::options()
+    .schema(
+        "foo::manipulate_instance(__torch__.torch.classes.my_classes.MyStackClass x) -> __torch__.torch.classes.my_classes.MyStackClass Y")
+    .catchAllKernel<decltype(manipulate_instance), &manipulate_instance>());
+// END free_function
+
+#endif

advanced_source/torch_script_custom_classes/custom_class_project/class.cpp

@@ -0,0 +1,48 @@
+import torch
+
+# `torch.classes.load_library()` allows you to pass the path to your .so file
+# to load it in and make the custom C++ classes available to both Python and
+# TorchScript
+torch.classes.load_library("build/libcustom_class.so")
+# You can query the loaded libraries like this:
+print(torch.classes.loaded_libraries)
+# prints {'/custom_class_project/build/libcustom_class.so'}
+
+# We can find and instantiate our custom C++ class in python by using the
+# `torch.classes` namespace:
+#
+# This instantiation will invoke the MyStackClass(std::vector<T> init)
+# constructor we registered earlier
+s = torch.classes.my_classes.MyStackClass(["foo", "bar"])
+
+# We can call methods in Python
+s.push("pushed")
+assert s.pop() == "pushed"
+
+# Returning and passing instances of custom classes works as you'd expect
+s2 = s.clone()
+s.merge(s2)
+for expected in ["bar", "foo", "bar", "foo"]:
+    assert s.pop() == expected
+
+# We can also use the class in TorchScript
+# For now, we need to assign the class's type to a local in order to
+# annotate the type on the TorchScript function. This may change
+# in the future.
+MyStackClass = torch.classes.my_classes.MyStackClass
+
+
+@torch.jit.script
+def do_stacks(s: MyStackClass):  # We can pass a custom class instance
+    # We can instantiate the class
+    s2 = torch.classes.my_classes.MyStackClass(["hi", "mom"])
+    s2.merge(s)  # We can call a method on the class
+    # We can also return instances of the class
+    # from TorchScript function/methods
+    return s2.clone(), s2.top()
+
+
+stack, top = do_stacks(torch.classes.my_classes.MyStackClass(["wow"]))
+assert top == "wow"
+for expected in ["wow", "mom", "hi"]:
+    assert stack.pop() == expected

advanced_source/torch_script_custom_classes/custom_class_project/custom_test.py

@@ -0,0 +1,21 @@
+# export_attr.py
+import torch
+
+torch.classes.load_library('build/libcustom_class.so')
+
+
+class Foo(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.stack = torch.classes.my_classes.MyStackClass(["just", "testing"])
+
+    def forward(self, s: str) -> str:
+        return self.stack.pop() + s
+
+
+scripted_foo = torch.jit.script(Foo())
+
+scripted_foo.save('foo.pt')
+loaded = torch.jit.load('foo.pt')
+
+print(loaded.stack.pop())

advanced_source/torch_script_custom_classes/custom_class_project/export_attr.py

@@ -0,0 +1,18 @@
+import torch
+
+torch.classes.load_library('build/libcustom_class.so')
+
+
+class Foo(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, s: str) -> str:
+        stack = torch.classes.my_classes.MyStackClass(["hi", "mom"])
+        return stack.pop() + s
+
+
+scripted_foo = torch.jit.script(Foo())
+print(scripted_foo.graph)
+
+scripted_foo.save('foo.pt')

advanced_source/torch_script_custom_classes/custom_class_project/save.py

@@ -0,0 +1,20 @@
+#include <torch/script.h>
+
+#include <iostream>
+#include <memory>
+
+int main(int argc, const char* argv[]) {
+  torch::jit::Module module;
+  try {
+    // Deserialize the ScriptModule from a file using torch::jit::load().
+    module = torch::jit::load("foo.pt");
+  }
+  catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n";
+    return -1;
+  }
+
+  std::vector<c10::IValue> inputs = {"foobarbaz"};
+  auto output = module.forward(inputs).toString();
+  std::cout << output->string() << std::endl;
+}

advanced_source/torch_script_custom_classes/infer.cpp

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -ex
+
+rm -rf build
+rm -rf custom_class_project/build
+
+pushd custom_class_project
+  mkdir build
+  (cd build && cmake CXXFLAGS="-DNO_PICKLE" -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..)
+  (cd build && make)
+  python custom_test.py
+  python save.py
+  ! python export_attr.py
+popd
+
+mkdir build
+(cd build && cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..)
+(cd build && make)
+mv custom_class_project/foo.pt build/foo.pt
+(cd build && ./infer)

advanced_source/torch_script_custom_classes/run.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+set -ex
+
+rm -rf build
+rm -rf custom_class_project/build
+
+pushd custom_class_project
+  mkdir build
+  (cd build && cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..)
+  (cd build && make)
+  python export_attr.py
+popd

advanced_source/torch_script_custom_classes/run2.sh

@@ -212,7 +212,7 @@ To now build our operator, we can run the following commands from our
 
   $ mkdir build
   $ cd build
-  $ cmake -DCMAKE_PREFIX_PATH=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)') ..
+  $ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
   -- The C compiler identification is GNU 5.4.0
   -- The CXX compiler identification is GNU 5.4.0
   -- Check for working C compiler: /usr/bin/cc
@@ -609,7 +609,7 @@ At this point, we should be able to build the application:
 
   $ mkdir build
   $ cd build
-  $ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
+  $ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
   -- The C compiler identification is GNU 5.4.0
   -- The CXX compiler identification is GNU 5.4.0
   -- Check for working C compiler: /usr/bin/cc
@@ -752,7 +752,7 @@ library. In the top level ``example_app`` directory:
 
   $ mkdir build
   $ cd build
-  $ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
+  $ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
   -- The C compiler identification is GNU 5.4.0
   -- The CXX compiler identification is GNU 5.4.0
   -- Check for working C compiler: /usr/bin/cc