microsoft / onnxjs

ONNX.js: run ONNX models using JavaScript

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Conv3d not supported

hgarrereyn opened this issue · comments

I'm trying to export a pytorch model that uses Conv3d layers for in-browser inference with onnx.js. Here is an example model:

class Test(torch.nn.Module):
    def __init__(self):
        super(Test, self).__init__()
        self.k = torch.nn.Conv3d(2, 3, (3,3,3))

    def forward(self, x):
        return self.k(x)

Exported via:

test = Test()
inp = torch.zeros(1,2,8,8,8)

torch.onnx.export(
    test,
    inp,
    './test.onnx',
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names = ['input'],
    output_names = ['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

In the browser:

const session = new onnx.InferenceSession();
await session.loadModel('/test.onnx');

var t = new Tensor(new Float32Array(2*8*8*8), "float32", [1,2,8,8,8]);
var res = await session.run([t]);

This fails with:

Uncaught (in promise) Error: invalid inputs detected; op: unnamed_Conv_0
    at t.<anonymous> (execution-plan.ts:101)
    at onnx.min.js:14
    at Object.next (onnx.min.js:14)
    at onnx.min.js:14
    at new Promise (<anonymous>)
    at r (onnx.min.js:14)
    at execution-plan.ts:98
    at t.event (instrument.ts:294)
    at execution-plan.ts:98
    at onnx.min.js:14

If I change Conv3d to Conv2d and change the tensor dimensions to [1,2,8,8] everything works fine. Additionally, the conv3d onnx model works fine in the python onnx runtime so I think this is specifically an onnx.js issue. This may also be related to #161 (unfortunately here there is no workaround by moving up one dimension).

I would also love to see 3D ops supported by ONNX.js

For anyone stumbling on this issue, the current workaround I'm using is converting from PyTorch to ONNX to Tensorflow Frozen model to TensorFlow web. Then I can use tensorflow.js to run inference in a browser. It's not elegant but it works.

Looking forward for 3D models support.