pytorch 模型转ONNX模型问题请教

[hejing-maker]

下面是我的转换代码，只是训练是单GPU训练。
"""
This code is used to convert the pytorch models into an onnx format models.
"""
import torch.onnx
from collections import OrderedDict
from model import carn_m
import os
import json
import time
import importlib
import argparse
import numpy as np
from collections import OrderedDict
import torch
import torch.nn as nn
from torch.autograd import Variable
from dataset import TestDataset
from PIL import Image

def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str)
parser.add_argument("--ckpt_path", type=str)
parser.add_argument("--group", type=int, default=1)
parser.add_argument("--scale", type=int, default=4)

return parser.parse_args()

def main(cfg):
module = importlib.import_module("model.{}".format(cfg.model))
net = module.Net(multi_scale=True,
group=cfg.group).cuda()
print(json.dumps(vars(cfg), indent=4, sort_keys=True))
net.load_state_dict(torch.load(cfg.ckpt_path))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = net.to(device)

#转换后模型保存路径
model_path = f"../checkpoint/onnx/carn.onnx"

dummy_input = torch.randn(1, 3, 256, 256).to("cuda")
torch.onnx.export(net, dummy_input, model_path, export_params=True, 
                      keep_initializers_as_inputs=True, # store the trained parameter weights inside the model file
                      opset_version=10,  # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose = True,
                      input_names=['input'],
                      `output_names=['enhance_image'])`

if name == "main":
cfg = parse_args()
main(cfg)

出现下面报错：
(base) hejing@tonly-Super-Server:~/TL_project/CARN-pytorch$ python carn/pytorch_to_onnx.py --model carn_m --scale 2 --ckpt_path ./checkpoint/carn_m_2000.pth --group 4
{
"ckpt_path": "./checkpoint/carn_m_2000.pth",
"group": 4,
"model": "carn_m",
"scale": 2
}
Traceback (most recent call last):
File "carn/pytorch_to_onnx.py", line 53, in
main(cfg)
File "carn/pytorch_to_onnx.py", line 48, in main
output_names=['enhance_image'])
File "/home/hejing/.local/lib/python3.6/site-packages/torch/onnx/init.py", line 148, in export
strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 66, in export
dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 416, in _export
fixed_batch_size=fixed_batch_size)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 279, in _model_to_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 236, in _trace_and_get_graph_from_model
trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(model, args, _force_outplace=True, _return_inputs_states=True)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/jit/init.py", line 277, in _get_trace_graph
outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in call
result = self.forward(*input, **kwargs)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/jit/init.py", line 360, in forward
self._force_outplace,
File "/home/hejing/.local/lib/python3.6/site-packages/torch/jit/init.py", line 347, in wrapper
outs.append(self.inner(*trace_inputs))
File "/home/hejing/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 530, in call
result = self._slow_forward(*input, **kwargs)
File "/home/hejing/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward
result = self.forward(*input, **kwargs)
**TypeError: forward() missing 1 required positional argument: 'scale'**
这是什么原因？