<p dir="auto">I’m trying to convert my model to ONNX format for further deployment in TensorRT. Here is a sample code to illustrate my problem in layer_norm here.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch
from torch import nn

class ExportModel(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        # n, c, h, w = x.shape
        # y = nn.functional.layer_norm(x, [c, h, w])       # not working
        # y = nn.functional.layer_norm(x, x.size()[1:])     # not working
        y = nn.functional.layer_norm(x, [16, 32, 128])

        return y

def main():
    model = ExportModel()

    dummy_input = torch.randn(64, 16, 32, 128)
    input_names = [ &quot;input&quot; ]
    output_names = [ &quot;output&quot; ]

    with torch.no_grad():
        torch.onnx.export(
            model, dummy_input, &quot;sample.onnx&quot;, verbose=True,
            input_names=input_names, output_names=output_names
        )
    return

if __name__ == '__main__':
    main()"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span> <span class="pl-k">import</span> <span class="pl-s1">nn</span>

<span class="pl-k">class</span> <span class="pl-v">ExportModel</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>):
        <span class="pl-en">super</span>().<span class="pl-en">__init__</span>()

    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-c"># n, c, h, w = x.shape</span>
        <span class="pl-c"># y = nn.functional.layer_norm(x, [c, h, w])       # not working</span>
        <span class="pl-c"># y = nn.functional.layer_norm(x, x.size()[1:])     # not working</span>
        <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-s1">functional</span>.<span class="pl-en">layer_norm</span>(<span class="pl-s1">x</span>, [<span class="pl-c1">16</span>, <span class="pl-c1">32</span>, <span class="pl-c1">128</span>])

        <span class="pl-k">return</span> <span class="pl-s1">y</span>

<span class="pl-k">def</span> <span class="pl-en">main</span>():
    <span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-v">ExportModel</span>()

    <span class="pl-s1">dummy_input</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randn</span>(<span class="pl-c1">64</span>, <span class="pl-c1">16</span>, <span class="pl-c1">32</span>, <span class="pl-c1">128</span>)
    <span class="pl-s1">input_names</span> <span class="pl-c1">=</span> [ <span class="pl-s">"input"</span> ]
    <span class="pl-s1">output_names</span> <span class="pl-c1">=</span> [ <span class="pl-s">"output"</span> ]

    <span class="pl-k">with</span> <span class="pl-s1">torch</span>.<span class="pl-en">no_grad</span>():
        <span class="pl-s1">torch</span>.<span class="pl-s1">onnx</span>.<span class="pl-en">export</span>(
            <span class="pl-s1">model</span>, <span class="pl-s1">dummy_input</span>, <span class="pl-s">"sample.onnx"</span>, <span class="pl-s1">verbose</span><span class="pl-c1">=</span><span class="pl-c1">True</span>,
            <span class="pl-s1">input_names</span><span class="pl-c1">=</span><span class="pl-s1">input_names</span>, <span class="pl-s1">output_names</span><span class="pl-c1">=</span><span class="pl-s1">output_names</span>
        )
    <span class="pl-k">return</span>

<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">'__main__'</span>:
    <span class="pl-en">main</span>()</pre></div>
<p dir="auto">It could only work when the parameter of layer_norm is constant number. If not, the following error will occur.</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;sample.py&quot;, line 31, in &lt;module&gt;
    main()
  File &quot;sample.py&quot;, line 26, in main
    verbose=True, input_names=input_names, output_names=output_names
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/__init__.py&quot;, line 148, in export
    strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py&quot;, line 66, in export
    dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py&quot;, line 409, in _export
    fixed_batch_size=fixed_batch_size)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py&quot;, line 289, in _model_to_graph
    fixed_batch_size=fixed_batch_size)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py&quot;, line 132, in _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/__init__.py&quot;, line 179, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py&quot;, line 647, in _run_symbolic_function
    return op_fn(g, *inputs, **attrs)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py&quot;, line 128, in wrapper
    args = [_parse_arg(arg, arg_desc) for arg, arg_desc in zip(args, arg_descriptors)]
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py&quot;, line 128, in &lt;listcomp&gt;
    args = [_parse_arg(arg, arg_desc) for arg, arg_desc in zip(args, arg_descriptors)]
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py&quot;, line 81, in _parse_arg
    &quot;', since it's not constant, please try to make &quot;
RuntimeError: Failed to export an ONNX attribute 'onnx::Gather', since it's not constant, please try to make things (e.g., kernel size) static if possible"><pre class="notranslate">Traceback (most recent call last):
  File <span class="pl-s"><span class="pl-pds">"</span>sample.py<span class="pl-pds">"</span></span>, line 31, <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
    <span class="pl-en">main</span>()
  File <span class="pl-s"><span class="pl-pds">"</span>sample.py<span class="pl-pds">"</span></span>, line 26, <span class="pl-k">in</span> main
    verbose=True, input_names=input_names, output_names=output_names
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/__init__.py<span class="pl-pds">"</span></span>, line 148, <span class="pl-k">in</span> <span class="pl-k">export</span>
    strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py<span class="pl-pds">"</span></span>, line 66, <span class="pl-k">in</span> <span class="pl-k">export</span>
    dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py<span class="pl-pds">"</span></span>, line 409, <span class="pl-k">in</span> _export
    fixed_batch_size=fixed_batch_size)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py<span class="pl-pds">"</span></span>, line 289, <span class="pl-k">in</span> _model_to_graph
    fixed_batch_size=fixed_batch_size)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py<span class="pl-pds">"</span></span>, line 132, <span class="pl-k">in</span> _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/__init__.py<span class="pl-pds">"</span></span>, line 179, <span class="pl-k">in</span> _run_symbolic_function
    <span class="pl-k">return</span> utils._run_symbolic_function(<span class="pl-k">*</span>args, <span class="pl-k">**</span>kwargs)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py<span class="pl-pds">"</span></span>, line 647, <span class="pl-k">in</span> _run_symbolic_function
    <span class="pl-k">return</span> op_fn(g, <span class="pl-k">*</span>inputs, <span class="pl-k">**</span>attrs)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py<span class="pl-pds">"</span></span>, line 128, <span class="pl-k">in</span> wrapper
    args = [_parse_arg(arg, arg_desc) <span class="pl-k">for</span> <span class="pl-smi">arg, arg_desc</span> <span class="pl-k">in</span> zip(args, arg_descriptors)]
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py<span class="pl-pds">"</span></span>, line 128, <span class="pl-k">in</span> <span class="pl-k">&lt;</span>listcomp<span class="pl-k">&gt;</span>
    args = [_parse_arg(arg, arg_desc) <span class="pl-k">for</span> <span class="pl-smi">arg, arg_desc</span> <span class="pl-k">in</span> zip(args, arg_descriptors)]
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py<span class="pl-pds">"</span></span>, line 81, <span class="pl-k">in</span> _parse_arg
    <span class="pl-s"><span class="pl-pds">"</span>', since it's not constant, please try to make <span class="pl-pds">"</span></span>
RuntimeError: Failed to <span class="pl-k">export</span> an ONNX attribute <span class="pl-s"><span class="pl-pds">'</span>onnx::Gather<span class="pl-pds">'</span></span>, since it<span class="pl-s"><span class="pl-pds">'</span>s not constant, please try to make things (e.g., kernel size) static if possible</span></pre></div>
<p dir="auto">I have few code blocks in my model have layer_norm op. It would turn into some ugly code if I explicitly mark all parameters constant number. Is there any “best practice” of how to use dynamic shape for this kind of use case?</p>
<p dir="auto">Also, I have posted the same issue on <a href="https://discuss.pytorch.org/t/how-to-convert-layer-norm-layer-to-onnx/66841" rel="nofollow">forum</a>. I'm not sure where is the better place for this kind of quesion, so I duplicate the issue here.</p>
<p dir="auto">Thanks in advance.</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/houseroad/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/houseroad">@houseroad</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/spandantiwari/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/spandantiwari">@spandantiwari</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/lara-hdr/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/lara-hdr">@lara-hdr</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/BowenBao/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/BowenBao">@BowenBao</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/neginraoof/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/neginraoof">@neginraoof</a></p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">I have encountered the following error in the backward pass of a simple 2 convolution layer model using nn.Conv3d.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="torch.Size([1, 1, 256, 256, 256])
torch.Size([1, 1, 256, 256, 256])
Forward pass successful
Loss computed successfully
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-1-04985c43e191&gt; in &lt;module&gt;
     28 loss = criterion(output, y)
     29 optimizer.zero_grad()
---&gt; 30 loss.backward()

/home/software/miniconda3/envs/torch_gpu/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    164                 products. Defaults to ``False``.
    165         &quot;&quot;&quot;
--&gt; 166         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    167
    168     def register_hook(self, hook):

/home/software/miniconda3/envs/torch_gpu/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
---&gt; 99         allow_unreachable=True)  # allow_unreachable flag
    100
    101

RuntimeError: cuDNN error: CUDNN_STATUS_NOT_SUPPORTED. This error may appear if you passed in a non-contiguous input."><pre class="notranslate"><code class="notranslate">torch.Size([1, 1, 256, 256, 256])
torch.Size([1, 1, 256, 256, 256])
Forward pass successful
Loss computed successfully
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-1-04985c43e191&gt; in &lt;module&gt;
     28 loss = criterion(output, y)
     29 optimizer.zero_grad()
---&gt; 30 loss.backward()

/home/software/miniconda3/envs/torch_gpu/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    164                 products. Defaults to ``False``.
    165         """
--&gt; 166         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    167
    168     def register_hook(self, hook):

/home/software/miniconda3/envs/torch_gpu/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
---&gt; 99         allow_unreachable=True)  # allow_unreachable flag
    100
    101

RuntimeError: cuDNN error: CUDNN_STATUS_NOT_SUPPORTED. This error may appear if you passed in a non-contiguous input.
</code></pre></div>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">This error is encountered when executing the following code for a simple 2 layer 3D convolutional neural network. An input tensor of size [1, 1, 256, 256, 256] is passed into the network. The first layer performs a 3D convolution with a 3x3x3 filter and padding=1, and the second layer performs a dilated 3D convolution with a 3x3x3 filter, dilation=2 and padding=2. The output tensors of the two layers are also of size [1, 1, 256, 256, 256].</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
from torch import nn

class Simple2Conv(nn.Module):
    def __init__(self):
        super(Simple2Conv, self).__init__()
        self.layer_1 = nn.Conv3d(1, 1, 3, padding=1)
        self.layer_2 = nn.Conv3d(1, 1, 3, padding=2, dilation=2)

    def forward(self, x):
        layer_1 = self.layer_1(x)        
        layer_2 = self.layer_2(layer_1)  
        
        print(layer_1.shape)
        print(layer_2.shape)
        return layer_2

model = Simple2Conv().cuda()

isize = 256
osize = isize
x = torch.FloatTensor(1,1,isize,isize,isize).cuda()
y = torch.LongTensor(1,osize ,osize ,osize).cuda()
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

output = model(x)
print('Forward pass successful')

loss = criterion(output, y)
print('Loss computed successfully')
optimizer.zero_grad()
loss.backward()
print('Backward pass successful')"><pre class="notranslate"><code class="notranslate">import torch
from torch import nn

class Simple2Conv(nn.Module):
    def __init__(self):
        super(Simple2Conv, self).__init__()
        self.layer_1 = nn.Conv3d(1, 1, 3, padding=1)
        self.layer_2 = nn.Conv3d(1, 1, 3, padding=2, dilation=2)

    def forward(self, x):
        layer_1 = self.layer_1(x)        
        layer_2 = self.layer_2(layer_1)  
        
        print(layer_1.shape)
        print(layer_2.shape)
        return layer_2

model = Simple2Conv().cuda()

isize = 256
osize = isize
x = torch.FloatTensor(1,1,isize,isize,isize).cuda()
y = torch.LongTensor(1,osize ,osize ,osize).cuda()
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

output = model(x)
print('Forward pass successful')

loss = criterion(output, y)
print('Loss computed successfully')
optimizer.zero_grad()
loss.backward()
print('Backward pass successful')
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">The expected output of the code snippet without the error is:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="torch.Size([1, 1, 256, 256, 256])
torch.Size([1, 1, 256, 256, 256])
Forward pass successful
Loss computed successfully
Backward pass successful"><pre class="notranslate"><code class="notranslate">torch.Size([1, 1, 256, 256, 256])
torch.Size([1, 1, 256, 256, 256])
Forward pass successful
Loss computed successfully
Backward pass successful
</code></pre></div>
<p dir="auto">Any of the following changes will resolve the error:</p>
<ol dir="auto">
<li>replace the dilated convolution layer with a regular convolution layer (i.e. duplicate the first layer)</li>
<li>swap the order of the two convolution layers (i.e. place the dilated convolution layer first)</li>
<li>Having only a single (dilated) convolution layer (i.e. remove either convolution layer)</li>
<li>for the dilated convolution layer, change padding=2 to padding=1 (note: this requires modification of <code class="notranslate">osize = isize</code> to <code class="notranslate">osize=isize-2</code>)</li>
<li>use CPU instead of GPU (by removing .cuda() throughout)</li>
<li>reduce isize to any number &lt;=255</li>
</ol>
<h2 dir="auto">Environment</h2>
<p dir="auto">(from running python collect_env.py)</p>
<p dir="auto">PyTorch version: 1.3.1<br>
Is debug build: No<br>
CUDA used to build PyTorch: 10.1.243</p>
<p dir="auto">OS: CentOS Linux 7 (Core)<br>
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-39)<br>
CMake version: Could not collect (manually added from <code class="notranslate">cmake --version</code>: 3.14.6)</p>
<p dir="auto">Python version: 3.7<br>
Is CUDA available: Yes<br>
CUDA runtime version: Could not collect (manually added from <code class="notranslate">nvidia-smi</code>: 10.1)<br>
GPU models and configuration:<br>
GPU 0: Tesla V100-SXM2-32GB</p>
<p dir="auto">Nvidia driver version: 418.87.01<br>
cuDNN version: Could not collect (manually added from <code class="notranslate">numba -s | grep cudnn</code>: 7.6.3)</p>
<p dir="auto">Versions of relevant libraries:<br>
[pip3] numpy==1.16.5<br>
[conda] _pytorch_select           0.2                       gpu_0<br>
[conda] blas                      1.0                         mkl<br>
[conda] mkl                       2019.4                      243<br>
[conda] mkl-service               2.3.0            py37he904b0f_0<br>
[conda] mkl_fft                   1.0.15           py37ha843d7b_0<br>
[conda] mkl_random                1.1.0            py37hd6b4f25_0<br>
[conda] pytorch                   1.3.1           py3.7_cuda10.1.243_cudnn7.6.3_0    pytorch<br>
[conda] pytorch-gpu               1.3.1                         0<br>
[conda] torchsummary              1.5.1                    pypi_0    pypi<br>
[conda] torchvision               0.4.2                py37_cu101    pytorch</p>
<h2 dir="auto">Additional context</h2>
<p dir="auto">Note in relation to issue <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="463904135" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/22496" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/22496/hovercard" href="https://github.com/pytorch/pytorch/issues/22496">#22496</a>, tensor sizes in this example are only 256^3 = 2^24, i.e. &lt; 2^31, thus do not violate the cudnn limit of 2^31. This error has been replicated on other systems with different versions of CUDA and PyTorch.</p>