The output tensor's shape of a scripted model is fixed by example's shape?

jit
1

In a pytorch network model, if I use nn.ConvTranspose2d and do traced_script_module = torch.jit.trace(model, example) after constructing the model. Whatever the shape size of example tensor is, I can get the output tensor, of which the shape is same as the input image, from traced_script_module.
But if I use nn.Upsample instead of nn.ConvTranspose2d in the model. Strange things happen. I only can get the output tensor which the shape of is the same as example tensor that I used in torch.jit.trace(model, example).
Can anyone explain this strange thing. Is that related to weak_script_method decorator or something?

2

Hi @cfanyyx, thanks for raising the issue. The problem on this is because nn.Upsample is a module that contains data dependent control flow (which means the traced model totally depends on the input you provide during tracing and will not generalize to the new input). To make nn.Upsample working in JIT, you will need to use TorchScript scripting instead of tracing. You can refer more to the doc here

Let me know if that’s clear or not.

3

Thanks for your reply. It helps me a lot. I will have a try. Actually, I tried to use TorchScript scripting before, but I didn’t make it so I changed to use tracing instead. I will try TorchScript scripting again. And I will come back if I meet any new question.

4

What’s the solution? It seems to not be just putting @torch.jit.script as a decorator on the function that calls upsample.