Loading UNET model into C++ using the new Torch Script feature of Pytorch 1.0


I am trying the new feature from Pytorch (dev version) which enables loading a Pytorch model in C++ without any Python dependencies. I am using the tracing method:

import torch
import torchvision

# An instance of your model.
model = A UNET MODEL FROM FASTAI which has hooks as required by UNET

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

However, I got a value error:

ValueError: Modules that have hooks assigned can't be compiled

A brief look at the error trace:

   1118         if orig._backward_hooks or orig._forward_hooks or orig._forward_pre_hooks:
-> 1119             raise ValueError("Modules that have hooks assigned can't be compiled")
   1121         for name, submodule in orig._modules.items():

I am using it for UNET MODEL from FASTAI. I was wondering if anyone else has tried it and would like to share their experiences if it worked. Also, does anyone if this new feature comes with a limitation of only working for very simple models?


(Jeremy Howard (Admin)) #2

I’m not sure about this - I’ve asked on twitter; hopefully a Pytorch guru will be able to help here. The issue is that Unet relies on a forward hook for the cross connections, and hooks aren’t currently supported by that new pytorch feature.


(Even Oldridge) #3

Will that be an issue for the library? Isn’t a lot of the functionality of v1 built around hooks?



I was wondering if we could try this feature with other fastai models as well. Its a really cool feature so would be nice to find a workaround for it.

1 Like

(Katharina) #5

[quote=“shbkan, post:1, topic:23810”]
ValueError: Modules that have hooks assigned can't be compiled
has anyone allready tried to use torch.jit.trace() or something similar to a fastai model?
I am trying to wrap my ResNet of fastai. I am using the new fast.ai version 1 but I am getting some errors. So if somone was allready successful, I would like to know what I have to change to make it work.


(Katharina) #6

I just used learn.model.cpu() to convert it and at least I did not get a error… I will try to make it work in c++ now

1 Like


Any update on this issue if someone else has tried it?


Converting UNET to Torchscript format for AWS Lambda
(Bob) #8

Here is a snippet of text I found in the following book that seems to provide the answer:

PyTorch Deep Learning Hands-On
By Sherin Thomas, Sudhanshu Passi
April 2019

PyTorch allows you to make a TorchScript IR through two methods. The easiest is by tracing, just like ONNX. You can pass the model (even a function) to torch.jit.trace with a dummy input. PyTorch runs the dummy input through the model/function and traces the operations while it runs the input.

The traced functions (PyTorch operations) then can be converted to the optimized IR, which is also called a static single assignment IR. Like an ONNX graph, instructions in this graph also have primitive operators that A TENsor library (ATen, the backend of PyTorch) would understand.

This is really easy but comes with a cost. Tracing-based inference has the basic problem ONNX had: it can’t handle the model structure changes that are dependent on the data, that is, an if / else condition check or a loop (sequence data). For handling such cases, PyTorch introduced scripting mode.

Scripting mode can be enabled by using the torch.jit.script decorator for normal functions and torch.jit.script_method for methods on the PyTorch model. By this decorator, the content inside a function/method will be directly converted to TorchScript. Another important thing to remember while using torch.jit.script_method for model classes is about the parent class. Normally, we inherit from torch.nn.Module , but for making TorchScript, we inherit from torch.jit.ScriptModule . This helps PyTorch to avoid using pure Python methods, which can’t be converted to TorchScript. Right now, TorchScript doesn’t support all Python features, but it has all the necessary features to support data-dependent tensor operations.


(Jeremy COCHOY) #9

Just to let you know, I managed to compiler and run a fastai DynamicUnet using torch.jit.
The model use hooks, but it seams the behavior is deterministic and induced by the shape of the dummy input. So I just commented pytorch’s code to force the algorithm to go on, and obtained a nice functional traced module.

It’s the line raise ValueError("Modules that have hooks assigned can't be compiled") that should be removed.

( The link to my SO question just in case https://stackoverflow.com/questions/56242857/how-can-i-force-torch-jit-trace-to-compule-my-module-by-ignoring-hooks )

Well, maybe it’s dirty, but at least you can compile it and run it in production :slight_smile:


(Alexander Spirin) #10

I’ve been following an awesome DeOldify notebook and got the same error -
ValueError("Modules that have hooks assigned can't be compiled")
The reason for this was the tensorboard callback, which I didn’t need at the time.

So, commenting the
learn_gen.callback_fns.append(partial(ImageGenTensorboardWriter, base_dir=TENSORBOARD_PATH, name='GenPre')) line did the job for me.

Hope it can help someone who’s also trying this notebook’s approach.
Source: https://github.com/jantic/DeOldify


(Eugene Ware) #11

FYI, this has recently been fixed in pytorch to allow tracing through forward passes. Tracing through backwards passes will still throw the error though, which is only an issue for training, not inference, which I guess is what most people are trying to do with tracing with UNETs.

This PR solves it. Not sure when it will land in an official release or nightly, but it’s now in master.



Hi Jeremycochoy

I followed your approach to compile unet, and got an error. here is my code

learn = unet_learner(data, models.resnet34, metrics=dice, wd=wd).to_fp16()

# after training
learn.to_fp32()    # put back to full floating point

trace_input = torch.ones(1,3,599,599).cuda()
jit_model = torch.jit.trace(learn.model, trace_input)
torch.jit.save(jit_model, f'models/{model_file}')

I got following warning

/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/fastai/vision/models/unet.py:32: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can’t record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if ssh != up_out.shape[-2:]:

and this error

~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/fastai/layers.py in forward(self, x)
    170         self.dense=dense
--> 172     def forward(self, x): return torch.cat([x,x.orig], dim=1) if self.dense else (x+x.orig)
    174 def res_block(nf, dense:bool=False, norm_type:Optional[NormType]=NormType.Batch, bottle:bool=False, **conv_kwargs):

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 1. Got 599 and 600 in dimension 2 at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/generic/THCTensorMath.cu:71

the shape is 599x600 instead of 599x599.

did you run into any issues when you compile your model?