我制作了一个简单的PyTorch MLP(GAN生成器),并按照教程(https://www.youtube.com/watch?v=Vs730jsRgO8)将其转换为ONNX,我的代码略有不同,但我无法捕捉到错误。
class Generator(nn.Module):def __init__(self, g_input_dim, g_output_dim): super(Generator, self).__init__() # g_input = 100 self.net = nn.Sequential( nn.Linear(g_input_dim, 256), nn.LeakyReLU(.2), nn.Linear(256, 512), nn.LeakyReLU(.2), nn.Linear(512, 1024), nn.LeakyReLU(.2), nn.Linear(1024, 784), nn.Tanh() )# forward methoddef forward(self, x): return self.net(x)
训练后,我将模型导出为ONNX格式。
torch.save(G.state_dict(), "pytorch_model.pth")import torch.onnxmodel = Generator(z_dim,mnist_dim)state_dict = torch.load("pytorch_model.pth")model.load_state_dict(state_dict)model.eval()dummy_input = torch.zeros(100)torch.onnx.export(model, dummy_input, "onnx_model.onnx", verbose=True)
这会生成以下ONNX图,看起来是准确的。
graph(%input.1 : Float(100), %net.0.bias : Float(256), %net.2.bias : Float(512), %net.4.bias : Float(1024), %net.6.bias : Float(784), %25 : Float(100, 256), %26 : Float(256, 512), %27 : Float(512, 1024), %28 : Float(1024, 784)): %10 : Float(256) = onnx::MatMul(%input.1, %25) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1612:0 %11 : Float(256) = onnx::Add(%10, %net.0.bias) %12 : Float(256) = onnx::LeakyRelu[alpha=0.20000000000000001](%11) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1239:0 %14 : Float(512) = onnx::MatMul(%12, %26) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1612:0 %15 : Float(512) = onnx::Add(%14, %net.2.bias) %16 : Float(512) = onnx::LeakyRelu[alpha=0.20000000000000001](%15) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1239:0 %18 : Float(1024) = onnx::MatMul(%16, %27) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1612:0 %19 : Float(1024) = onnx::Add(%18, %net.4.bias) %20 : Float(1024) = onnx::LeakyRelu[alpha=0.20000000000000001](%19) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1239:0 %22 : Float(784) = onnx::MatMul(%20, %28) # /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1612:0 %23 : Float(784) = onnx::Add(%22, %net.6.bias) %24 : Float(784)
然后我将代码导入到JavaScript中。
<html> <body> <script src="./onnx.min.js"></script> <script> async function test() { const sess = new onnx.InferenceSession() await sess.loadModel('./onnx_model.onnx') const input = new onnx.Tensor(new Float32Array(100), 'float32', [100]) const outputMap = await sess.run([input]) const outputTensor = outputMap.values().next().value console.log(`Output tensor: ${outputTensor.data}`) } test() </script> </body></html>
我知道输入维度是正确的,但ONNX给我以下错误。
onnx.min.js:8 Uncaught (in promise) Error: Can't use matmul on the given tensors at e.createProgramInfo (onnx.min.js:8) at t.run (onnx.min.js:8) at e.run (onnx.min.js:8) at t.<anonymous> (onnx.min.js:14) 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 onnx.min.js:14
我也知道matmul是ONNX支持的操作符,但我无法弄清楚我的输入张量是否正确,或者如何正确设置它。
回答:
我认为matmul操作符期望输入是二维的。当我在输入中添加一个批次大小维度(批次大小为1)时,它似乎可以工作:
之前:dummy_input = torch.zeros(100)
之后:dummy_input = torch.zeros(1, 100)
之前:const input = new onnx.Tensor(new Float32Array(100), 'float32', [100])
之后:const input = new onnx.Tensor(new Float32Array(100), 'float32', [1, 100])