torch.onnx.export

 torch.onnx.export

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)

x = torch.onnx.export(model,  # 网络模型
                torch.randn(1, 3, 224, 224), # 用于确定输入大小和类型，其中的值可以是随机的。
                export_onnx_file,  # 输出onnx的名称
                verbose=False,      # 是否以字符串的形式显示计算图
                input_names=["input"],  # 输入节点的名称，可以是一个list
                output_names=["output"], # 输出节点的名称
                opset_version=10,   # onnx 支持采用的operator set
                do_constant_folding=True, # 是否压缩常量
                #设置动态维度，此处指明input节点的第0维度可变，命名为batch_size
                dynamic_axes={"input":{0: "batch_size", 2: "h", 3: "w"}, "output":{0: "batch_size"}} 
                )

model (torch.nn.Module) – 要导出的模型.
args (tuple of arguments) – 模型的输入, 任何非Tensor参数都将硬编码到导出的模型中；任何Tensor参数都将成为导出的模型的输入，并按照他们在args中出现的顺序输入。因为export运行模型，所以我们需要提供一个输入张量x。只要是正确的类型和大小，其中的值就可以是随机的。请注意，除非指定为动态轴，否则输入尺寸将在导出的ONNX图形中固定为所有输入尺寸。在此示例中，我们使用输入batch_size 1导出模型，但随后dynamic_axes 在torch.onnx.export()。因此，导出的模型将接受大小为[batch_size，3、100、100]的输入，其中batch_size可以是可变的。
export_params (bool, default True) – 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False.
verbose (bool, default False) - 如果指定，我们将打印出一个导出轨迹的调试描述。
training (bool, default False) - 在训练模式下导出模型。目前，ONNX导出的模型只是为了做推断，所以你通常不需要将其设置为True。
input_names (list of strings, default empty list) – 按顺序分配名称到图中的输入节点
output_names (list of strings, default empty list) –按顺序分配名称到图中的输出节点
dynamic_axes – {‘input’ : {0 : ‘batch_size’}, ‘output’ : {0 : ‘batch_size’}}
 

 onnx INT64 weights error

TensorRT解析PyTorch导出的onnx模型时:

Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.

pip install onnx-simplifier
python -m onnxsim model_old.onnx model_sim_new.onnx