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

This code provides a PyTorch implementation for OTTER (Optimal Transport distillation for Efficient zero-shot Recognition), as described in the paper.

QueryInst: Parallelly Supervised Mask Query for Instance Segmentation

This repository contains various models targetting multimodal representation learning, multimodal fusion for downstream tasks such as multimodal sentiment analysis.

A tutorial on DataFrames.jl prepared for JuliaCon2021

Unofficial Pytorch Lightning implementation of Contrastive Syn-to-Real Generalization (ICLR, 2021)

Official repository for: Continuous Control With Ensemble DeepDeterministic Policy Gradients

This repo contains the code for the paper "Efficient hierarchical Bayesian inference for spatio-temporal regression models in neuroimaging" that has been accepted to NeurIPS 2021.

An Efficient Training Approach for Very Large Scale Face Recognition or F²C for simplicity.

Automatically align face images 🙃→🙂. Can also do windowing and warping.

Losslandscapetaxonomy - Taxonomizing local versus global structure in neural network loss landscapes

[BMVC 2021] Official PyTorch Implementation of Self-supervised learning of Image Scale and Orientation Estimation

This repository contains the official implementation code of the paper Transformer-based Feature Reconstruction Network for Robust Multimodal Sentiment Analysis

TipToiDog - Tip Toi Dog With Python

Optimizing Deeper Transformers on Small Datasets

Checking fibonacci - Generating the Fibonacci sequence is a classic recursive problem

Python package for visualizing the loss landscape of parameterized quantum algorithms.

Implementation of QuickDraw - an online game developed by Google, combined with AirGesture - a simple gesture recognition application

[CVPRW 21] "BNN - BN = ? Training Binary Neural Networks without Batch Normalization", Tianlong Chen, Zhenyu Zhang, Xu Ouyang, Zechun Liu, Zhiqiang Shen, Zhangyang Wang

This is the workbook I created while I was studying for the Qiskit Associate Developer exam. I hope this becomes useful to others as it was for me :)

Jittor Medical Segmentation Lib -- The assignment of Pattern Recognition course (2021 Spring) in Tsinghua University