[onnx] crash when exporting a model with Sequence module (encodeBlock: Assertion failled )

[gkossakowski]

🐛 Bug

The tracer in onnx exporter is failing with an assertion error in the exporter.

To Reproduce

Steps to reproduce the behavior:

1. define a model:

class SomeModel(nn.Module):

    def __init__(self):
        super(SomeModel, self).__init__()
        dim = 5
        self.emb = nn.Embedding(10, dim)
        self.lin1 = nn.Linear(dim, 1)
        self.seq = nn.Sequential(
            self.emb,
            self.lin1,
        )

    def forward(self, input):
        return self.seq(input)

model = SomeModel()

2. run export:

dummy_input = torch.tensor([2], dtype=torch.long)
dummy_input

torch.onnx.export(model, dummy_input, "model.onnx", verbose=True)

fails with:

[redacted]/vendor/local/lib/python2.7/site-packages/torch/onnx/utils.pyc in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, propagate)
    230     defer_weight_export = export_type is not ExportTypes.PROTOBUF_FILE
    231     if export_params:
--> 232         proto, export_map = graph.export(params, _onnx_opset_version, defer_weight_export, operator_export_type)
    233     else:
    234         proto, export_map = graph.export([], _onnx_opset_version, False, operator_export_type)

RuntimeError: torch/csrc/jit/export.cpp:296: encodeBlock: Assertion `b->inputs().size() >= num_initializers` failed.

Expected behavior

Model is exported to ONNX

Environment

Collecting environment information...
PyTorch version: 0.4.1
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.11) 5.4.0 20160609
CMake version: Could not collect

Python version: 2.7
Is CUDA available: No
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA

Versions of relevant libraries:
[pip] msgpack-numpy==0.4.3.1
[pip] numpy==1.14.1
[pip] torch==0.4.1
[conda] Could not collect

Additional context

As a workaround one can inline layer definitions into the sequence:

class SomeModel(nn.Module):

    def __init__(self):
        super(SomeModel, self).__init__()
        dim = 5
        self.seq = nn.Sequential(
            nn.Embedding(10, dim),
            nn.Linear(dim, 1),
        )

    def forward(self, input):
        return self.seq(input)

model = SomeModel()

[gkossakowski]

I just checked that pytorch 1.0 doesn't fix this issue. I get the same crash.

The workaround I posted where I inline all layer definitions into a Sequential is workable but awkward to use with more complex architectures. Can you think of other workaround?

[rgommers]

Here's a script to run both this test case and the test case in gh-20329 (same issue) in one go and clean up the generated files again:

import os

import torch
from torch import nn


class SomeModel(nn.Module):

    def __init__(self):
        super(SomeModel, self).__init__()
        dim = 5
        self.emb = nn.Embedding(10, dim)
        self.lin1 = nn.Linear(dim, 1)
        self.seq = nn.Sequential(
            self.emb,
            self.lin1,
        )

    def forward(self, input):
        return self.seq(input)


model = SomeModel()


dummy_input = torch.tensor([2], dtype=torch.long)
dummy_input

torch.onnx.export(model, dummy_input, "model.onnx", verbose=True)
os.remove('model.onnx')


print('\n\nTEST CASE 2\n\n')


n = 8
dim = 10

class SomeModel(nn.Module):

    def __init__(self):
        super(SomeModel, self).__init__()
        self.embedding = nn.Embedding(n, dim)
        self.seq = nn.Sequential(
            self.embedding,
            nn.Linear(dim, 1),
            nn.Sigmoid()
        )

    def forward(self, indices):
        return self.seq(indices)

model = SomeModel()

dummy_input = torch.LongTensor([2])
torch.onnx.export(model, dummy_input, "foo.onnx", verbose=True)

os.remove('foo.onnx')

This indeed gives a RuntimeError for PyTorch 1.0, however it works in 1.1, 1.3.1 and current master; gives output:

graph(%input.1 : Long(1),
      %emb.weight : Float(10, 5),
      %lin1.weight : Float(1, 5),
      %lin1.bias : Float(1)):
  %4 : Float(1, 5) = onnx::Gather(%emb.weight, %input.1), scope: SomeModel/Sequential[seq]/Embedding[0]
  %5 : Float(1, 1) = onnx::Gemm[alpha=1, beta=1, transB=1](%4, %lin1.weight, %lin1.bias), scope: SomeModel/Sequential[seq]/Embedding[0]
  return (%5)



TEST CASE 2


graph(%input.1 : Long(1),
      %embedding.weight : Float(8, 10),
      %seq.1.weight : Float(1, 10),
      %seq.1.bias : Float(1)):
  %4 : Float(1, 10) = onnx::Gather(%embedding.weight, %input.1), scope: SomeModel/Sequential[seq]/Embedding[0]
  %5 : Float(1, 1) = onnx::Gemm[alpha=1, beta=1, transB=1](%4, %seq.1.weight, %seq.1.bias), scope: SomeModel/Sequential[seq]/Embedding[0]
  %6 : Float(1, 1) = onnx::Sigmoid(%5), scope: SomeModel/Sequential[seq]/Sigmoid[2]
  return (%6)

This seems fixed, output looks right to me. Can you confirm that this works as expected now @gkossakowski?

There's several unit tests for both Sequential and Embedding under test/onnx/, but this exact case isn't tested yet, so adding a regression test before closing this may still be a good idea.