/torchscript_atlas_sr

Primary LanguageC++

This repo demonstrates how to use torchscript to wrap a pytorch model in a such that it can be embedded in a C++ program.

Python

trace_model.py loads a serialized pytorch model @ model_fp and performs source analysis on it in order to embed. It has a command-line interface, i.e., it can be run by executing

$ python trace_model.py --model=dbpn --upscale=2

CPP

Compiling

The CPP dependencies are OpenCV 4.1.2 and LibTorch built against CUDA 10.1:

image

Adjust CMakeLists.txt appropriately according to your OpenCV and LibTorch install paths. Then

$ cmake -DCMAKE_BUILD_TYPE=Debug
$ make

Run

main.cpp is a CLI app that ingests a traced pytorch model and upscales cat.jpg. After compiling it can be run

$ ./main \
  -w /home/maksim/dev_projects/torchscript/traced_dbpn_2x.pt \
  -i /home/maksim/dev_projects/torchscript/cat.jpg \
  -o /home/maksim/dev_projects/torchscript/sr_cat.jpg

where -w is the file path of the traced model as produced by trace_model.py.