pyTorch模型轉換為ONNX時遇到問題

Question

我想將pyTorch模型轉換為ONNX。 但是，我得到一個錯誤，說

RuntimeError：提供的輸入名稱的數量（9）超過了輸入的數量（7）但是，如果我從模型中刪除兩個Dropout層，則我的代碼可以完美運行。

為什么是這樣？

這是我的代碼：

# Define the model
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Dropout(0.2),  # problem with dropout layer
    torch.nn.Linear(H, H),
    torch.nn.LeakyReLU(),
    torch.nn.Dropout(0.2),  # problem with dropout layer
    torch.nn.Linear(H, D_out),
    torch.nn.Sigmoid()
)

checkpoint = torch.load("./saved_pytorch_model.pth")  # load pyTorch model
model.load_state_dict(checkpoint['state_dict'])
features = torch.Tensor(df_X.values[0])

# Convert pyTorch model to ONNX
input_names = ['input_1']
output_names = ['output_1']
for key, module in model._modules.items():
    input_names.append("l_{}_".format(key) + module._get_name())
torch_out = torch.onnx.export(model, 
                         features, 
                         "onnx_model.onnx", 
                         export_params = True, 
                         verbose = True, 
                         input_names = input_names, 
                         output_names = output_names,
                        )

我如何做才能使其包含Dropout導出到ONNX？

Answer 1

在我看來，這似乎是個錯誤。 您可以注釋掉輸入名稱參數。

# Convert pyTorch model to ONNX
input_names = ['input_1']
output_names = ['output_1']
for key, module in model._modules.items():
    input_names.append("l_{}_".format(key) + module._get_name())
torch_out = torch.onnx.export(model, 
                         features, 
                         "onnx_model.onnx", 
                         export_params = True, 
                         verbose = True, 
                         #input_names = input_names, 
                         output_names = output_names,
                        )

您將獲得帶有%輸入名稱，如下所示：

graph(%input.1 : Float(10, 3, 224, 10)
      %1 : Float(10, 10)
      %2 : Float(10)
      %3 : Float(10, 10)
      %4 : Float(10)
      %5 : Float(10, 10)
      %6 : Float(10)) {
  %7 : Float(10!, 10!) = onnx::Transpose[perm=[1, 0]](%1), scope: Sequential/Linear[0]
  %8 : Float(10, 3, 224, 10) = onnx::MatMul(%input.1, %7), scope: Sequential/Linear[0]
  %9 : Float(10, 3, 224, 10) = onnx::Add(%8, %2), scope: Sequential/Linear[0]
  %10 : Float(10, 3, 224, 10) = onnx::Relu(%9), scope: Sequential/ReLU[1]
  %11 : Float(10, 3, 224, 10), %12 : Tensor = onnx::Dropout[ratio=0.2](%10), scope: Sequential/Dropout[2]
  %13 : Float(10!, 10!) = onnx::Transpose[perm=[1, 0]](%3), scope: Sequential/Linear[3]
  %14 : Float(10, 3, 224, 10) = onnx::MatMul(%11, %13), scope: Sequential/Linear[3]
  %15 : Float(10, 3, 224, 10) = onnx::Add(%14, %4), scope: Sequential/Linear[3]
  %16 : Float(10, 3, 224, 10) = onnx::LeakyRelu[alpha=0.01](%15), scope: Sequential/LeakyReLU[4]
  %17 : Float(10, 3, 224, 10), %18 : Tensor = onnx::Dropout[ratio=0.2](%16), scope: Sequential/Dropout[5]
  %19 : Float(10!, 10!) = onnx::Transpose[perm=[1, 0]](%5), scope: Sequential/Linear[6]
  %20 : Float(10, 3, 224, 10) = onnx::MatMul(%17, %19), scope: Sequential/Linear[6]
  %21 : Float(10, 3, 224, 10) = onnx::Add(%20, %6), scope: Sequential/Linear[6]
  %output_1 : Float(10, 3, 224, 10) = onnx::Sigmoid(%21), scope: Sequential/Sigmoid[7]
  return (%output_1);
}