microsoft / onnxjs

ONNX.js: run ONNX models using JavaScript
Other
1.75k stars 130 forks source link

Conv3d not supported #229

Open hgarrereyn opened 4 years ago

hgarrereyn commented 4 years ago

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).

dsuess commented 3 years ago

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

hgarrereyn commented 3 years ago

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.

neso613 commented 3 years ago

Looking forward for 3D models support.