Simple Conv2d Function cannot be scripted and reports Runtime Error.

username: Yang_Hu
cooked: <h2><a name="simple-conv2d-function-cannot-be-scripted-and-reports-runtime-error-1" class="anchor" href="#simple-conv2d-function-cannot-be-scripted-and-reports-runtime-error-1"></a>Simple Conv2d Function cannot be scripted and reports Runtime Error.</h2>
<p>Here is my simple Conv2d module, I want to script it using <code>torch.jit.script</code>.</p>
<pre><code class="lang-auto">import torch

class Conv2dCell(torch.nn.Module):
  def __init__(self):
    super(Conv2dCell, self).__init__()

  def forward(self, x):
    conv = torch.nn.Conv2d(1, 3, 3, stride=1)
    output = conv(x)
    return output

m = Conv2dCell()
scripted_m = torch.jit.script(m)
</code></pre>
<p>Running this piece of code will give the following error message:</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “conv2d.py”, line 13, in <br>
scripted_m = torch.jit.script(m)<br>
File “/mnt/ssd/maxhy/py36/lib/python3.6/site-packages/torch/jit/<strong>init</strong>.py”, line 1261, in script<br>
return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)<br>
File “/mnt/ssd/maxhy/py36/lib/python3.6/site-packages/torch/jit/_recursive.py”, line 305, in create_script_module<br>
return create_script_module_impl(nn_module, concrete_type, stubs_fn)<br>
File “/mnt/ssd/maxhy/py36/lib/python3.6/site-packages/torch/jit/_recursive.py”, line 361, in create_script_module_impl<br>
create_methods_from_stubs(concrete_type, stubs)<br>
File “/mnt/ssd/maxhy/py36/lib/python3.6/site-packages/torch/jit/_recursive.py”, line 279, in create_methods_from_stubs<br>
concrete_type._create_methods(defs, rcbs, defaults)<br>
File “/mnt/ssd/maxhy/py36/lib/python3.6/site-packages/torch/jit/<strong>init</strong>.py”, line 1108, in _compile_and_register_class<br>
_jit_script_class_compile(qualified_name, ast, rcb)<br>
RuntimeError:<br>
Arguments for call are not valid.<br>
The following variants are available:</p>
<p>_pair(float[2] x) → (float<span class="chcklst-box fa fa-square-o fa-fw"></span>):<br>
Expected a value of type ‘List[float]’ for argument ‘x’ but instead found type ‘Tensor’.</p>
<p>_pair(int[2] x) → (int<span class="chcklst-box fa fa-square-o fa-fw"></span>):<br>
Expected a value of type ‘List[int]’ for argument ‘x’ but instead found type ‘Tensor’.</p>
<p>The original call is:<br>
File “/mnt/ssd/maxhy/py36/lib/python3.6/site-packages/torch/nn/modules/conv.py”, line 336<br>
padding=0, dilation=1, groups=1,<br>
bias=True, padding_mode=‘zeros’):<br>
kernel_size = _pair(kernel_size)<br>
~~~~~ &lt;— HERE<br>
stride = _pair(stride)<br>
padding = _pair(padding)<br>
‘Conv2d.<strong>init</strong>’ is being compiled since it was called from ‘Conv2d’<br>
File “conv2d.py”, line 8<br>
def forward(self, x):<br>
conv = torch.nn.Conv2d(1, 3, 3, stride=1)<br>
~~~~~~~~~~~~~~~ &lt;— HERE<br>
output = conv(x)<br>
return output<br>
‘Conv2d’ is being compiled since it was called from ‘Conv2dCell.forward’<br>
File “conv2d.py”, line 8<br>
def forward(self, x):<br>
conv = torch.nn.Conv2d(1, 3, 3, stride=1)<br>
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;— HERE<br>
output = conv(x)<br>
return output</p>
</blockquote>
<p>I am using <strong>PyTorch 1.5.1</strong> and <strong>python 3.6.13</strong>, could someone help me identify the problem?</p>

username: ptrblck
cooked: <p>You are recreating a conv layer in the forward, which is most likely wrong.<br>
(The usual way would be to initialize it in the <code>__init__</code> method and use it in the <code>forward</code>)<br>
Note that this would initialize a new layer in each forward pass and thus it won’t be trained.<br>
If that’s intended, could you explain your use case a bit more?</p>

username: Yang_Hu
cooked: <p>Hi ptrblck,<br>
No special use case here. I just want to know if this (cannot define Conv2d layer in <code>forward</code> function) is a restriction in Pytorch? Is there formal documentation for this? And also I notice that we can actually define some simple ops in <code>forward</code> function like <code>add</code> or <code>relu</code>, what is the difference between simple ops and more complicated ops like <code>Conv2d</code>?</p>

username: ptrblck
cooked: <p>I wouldn’t call it a restriction, but maybe an unwanted usage.<br>
Modules containing parameters are initialized before they can be used.<br>
E.g. conv and linear layers contain trainable parameters and you would create objects first:</p>
<pre><code class="lang-python">conv = nn.Conv2d(3, 6, 3, 1, 1)
lin = nn.Linear(10, 10)
</code></pre>
<p>Later in your code you would then use these modules and feed an activation to them:</p>
<pre><code class="lang-python">out_conv = conv(in_conv)
out_lin = lin(in_lin)
</code></pre>
<p>The common use case it to create these layers in the <code>__init__</code> method of your custom module and use them in the <code>forward</code>.<br>
However, you could of course create them outside of the <code>__init__</code> and pass them to it or even pass them to the <code>forward</code> method.<br>
The issue in your code is created, since you are recreating the modules in each forward pass, which will thus reset the parameters every time the forward pass is executed.<br>
The <a href="https://pytorch.org/tutorials/beginner/nn_tutorial.html"><code>nn</code> tutorial</a> might explain it in more detail.</p>

username: Yang_Hu
cooked: <p>Modules are computation blocks with state, can we put it this way? But simple ops like <code>add</code> do not hold states. And what you explained makes sense because during training we need to save states for every iteration.</p>
<p>But Torchscript is often used for inference. I don’t see why we cannot script or trace a forward function with a submodule defined in it. From my understanding, scripting or tracing a Function(or a Module) does not save states like weight or bias right? It will just record a pure function path.</p>

username: ptrblck
cooked: <p>Yes, the first explanation of “stateful” modules makes sense.</p>
<p>I’m not sure how TorchScript is related to this. Note that you surely can re-initialize modules in the <code>forward</code> pass, if you explicitly don’t want to train these layers and want to create new random parameters.<br>
A scripted model should respect this workflow (even if it’s wrong from the point of view of training the model).</p>