Custom TorchScript Operators Example

This repository contains examples for writing, compiling and using custom TorchScript operators. See here for the accompanying tutorial.

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 using torch.utils.cpp_extension.load to build the custom operator,
load.py: Example of using torch.utils.cpp_extension.load_inline to 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 using setup.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]])

Example repository for custom C++/CUDA operators for TorchScript

Related tags

Overview

Custom TorchScript Operators Example

Contents

Setup

Owner

Toolchain to build Yoshi's Island from source code

Lightweight plotting to the terminal. 4x resolution via Unicode.

Official implementation of the paper 'High-Resolution Photorealistic Image Translation in Real-Time: A Laplacian Pyramid Translation Network' in CVPR 2021

The VeriNet toolkit for verification of neural networks

Learning Spatio-Temporal Transformer for Visual Tracking

QuickAI is a Python library that makes it extremely easy to experiment with state-of-the-art Machine Learning models.

[CIKM 2019] Code and dataset for "Fi-GNN: Modeling Feature Interactions via Graph Neural Networks for CTR Prediction"

Open-source implementation of Google Vizier for hyper parameters tuning

Code for the Population-Based Bandits Algorithm, presented at NeurIPS 2020.

Pytorch implementation of our method for high-resolution (e.g. 2048x1024) photorealistic video-to-video translation.

Monocular 3D pose estimation. OpenVINO. CPU inference or iGPU (OpenCL) inference.

[ICLR 2021] Heteroskedastic and Imbalanced Deep Learning with Adaptive Regularization

Training deep models using anime, illustration images.

Neural Logic Inductive Learning

Personalized Transfer of User Preferences for Cross-domain Recommendation (PTUPCDR)

Generalized hybrid model for mode-locked laser diodes with an extended passive cavity

Code to reproduce the experiments in the paper "Transformer Based Multi-Source Domain Adaptation" (EMNLP 2020)

Finetuner allows one to tune the weights of any deep neural network for better embeddings on search tasks

Softlearning is a reinforcement learning framework for training maximum entropy policies in continuous domains. Includes the official implementation of the Soft Actor-Critic algorithm.

RL and distillation in CARLA using a factorized world model