Custom TorchScript Operators Example
This repository contains examples for writing, compiling and using custom TorchScript operators. See here for the accompanying tutorial.
Contents
There a few monuments in this repository you can visit. They are described in context in the tutorial, which you are encouraged to read. These monuments are:
example_app/warp_perspective/op.cpp: The custom operator implementation,example_app/main.cpp: An example application that loads and executes a serialized TorchScript model, which uses the custom operator, in C++,script.py: Example of using the custom operator in a scripted model,trace.py: Example of using the custom operator in a traced model,eager.py: Example of using the custom operator in vanilla eager PyTorch,load.py: Example of usingtorch.utils.cpp_extension.loadto build the custom operator,load.py: Example of usingtorch.utils.cpp_extension.load_inlineto build the custom operator,setup.py: Example of using setuptools to build the custom operator,test_setup.py: Example of using the custom operator built usingsetup.py.
To execute the C++ application, first run script.py to serialize a TorchScript model to a file called example.pt, then pass that file to the example_app/build/example_app binary.
Setup
For the smoothest experience when trying out these examples, we recommend building a docker container from this repository's Dockerfile. This will give you a clean, isolated Ubuntu Linux environment in which we guarantee everything to work perfectly. These steps should get you started:
$ git clone https://github.com/pytorch/extension-script
$ cd extension-script
$ docker build -t extension-script .
$ docker run -v $PWD:/home -it extension-script
$ [email protected]:/home# source /activate # Activate the Conda environment
$ cd example_app && mkdir build && cd build
$ cmake -DCMAKE_PREFIX_PATH=/libtorch ..
-- 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
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE
-- Found torch: /libtorch/lib/libtorch.so
-- Configuring done
-- Generating done
-- Build files have been written to: /home/example_app/build
$ make -j
Scanning dependencies of target warp_perspective
[ 25%] Building CXX object warp_perspective/CMakeFiles/warp_perspective.dir/op.cpp.o
[ 50%] Linking CXX shared library libwarp_perspective.so
[ 50%] Built target warp_perspective
Scanning dependencies of target example_app
[ 75%] Building CXX object CMakeFiles/example_app.dir/main.cpp.o
[100%] Linking CXX executable example_app
[100%] Built target example_app
This will create a shared library under /home/example_app/build/warp_perspective/libwarp_perspective.so containing the custom operator defined in example_app/warp_perspective/op.cpp. Then, you can run the examples, e.g.:
(base) [email protected]:/home# python script.py
graph(%x.1 : Dynamic
%y : Dynamic) {
%20 : int = prim::Constant[value=1]()
%16 : int[] = prim::Constant[value=[0, -1]]()
%14 : int = prim::Constant[value=6]()
%2 : int = prim::Constant[value=0]()
%7 : int = prim::Constant[value=42]()
%z.1 : int = prim::Constant[value=5]()
%z.2 : int = prim::Constant[value=10]()
%13 : int = prim::Constant[value=3]()
%4 : Dynamic = aten::select(%x.1, %2, %2)
%6 : Dynamic = aten::select(%4, %2, %2)
%8 : Dynamic = aten::eq(%6, %7)
%9 : bool = prim::TensorToBool(%8)
%z : int = prim::If(%9)
block0() {
-> (%z.1)
}
block1() {
-> (%z.2)
}
%17 : Dynamic = aten::eye(%13, %14, %2, %16)
%x : Dynamic = my_ops::warp_perspective(%x.1, %17)
%19 : Dynamic = aten::matmul(%x, %y)
%21 : Dynamic = aten::add(%19, %z, %20)
return (%21);
}
tensor([[11.6196, 12.0056, 11.6122, 12.9298, 7.0649],
[ 8.5063, 9.0621, 9.9925, 6.3741, 8.9668],
[12.5898, 6.5872, 8.1511, 10.0806, 11.9829],
[ 4.9142, 11.6614, 15.7161, 17.0538, 11.7243],
[10.0000, 10.0000, 10.0000, 10.0000, 10.0000],
[10.0000, 10.0000, 10.0000, 10.0000, 10.0000],
[10.0000, 10.0000, 10.0000, 10.0000, 10.0000],
[10.0000, 10.0000, 10.0000, 10.0000, 10.0000]])
60 Dec 08, 2022
1 Jan 10, 2022
576 Dec 29, 2022
19 Oct 22, 2022
2.1k Dec 27, 2022
4 May 05, 2022
44 Nov 01, 2022
262 Dec 13, 2022
129 Dec 24, 2022
16 Nov 28, 2022
16 Dec 15, 2022
52 Jan 04, 2023
57 May 14, 2021
1k Jan 05, 2023
1.1k Dec 30, 2022
109 Dec 28, 2022
254 Jan 08, 2023
32 Dec 24, 2022
22 Jan 04, 2023
1.6k Dec 31, 2022