pytorch 使用動態shape輸出onnx - jenhaoyang/ml_blog GitHub Wiki

torch.onnx.export(SumModule(), (torch.ones(2, 2),), "onnx.pb",
                  input_names=["x"], output_names=["sum"],
                  dynamic_axes={
                      # dict value: manually named axes
                      "x": {0: "my_custom_axis_name"},
                      # list value: automatic names
                      "sum": [0],
                  })

參考:
https://pytorch.org/docs/stable/onnx.html#functions