pt文件转onnx确定不同的输入宽高如何转换？

现实中训练好的模型部署时经常会有这样的需求，就是模型的输入需要改变，不使用训练时的输入大小，如yolo系列的模型训练时一般都是输入的图片是640x640，但是部署时我希望输入到模型的分辨率是1920x1080（调整为1920x1088，必须为32的整数倍），那么如何通过pt文件导出对应的onnx模型文件呢？静态batch如何导出呢？动态batch如何导出呢？

先说说本人的环境：

Ultralytics YOLOv8.0.53 Python-3.9.13 torch-1.12.1+cu116 CPU

ONNX: starting export with onnx 1.13.0

在yolov8中，官方作者已经写好了转换代码，具体如下：

output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
dynamic = self.args.dynamic
if dynamic:
    dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
    if isinstance(self.model, SegmentationModel):
        dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
        dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
    elif isinstance(self.model, DetectionModel):
        dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)

torch.onnx.export(
    self.model.cpu() if dynamic else self.model,  # --dynamic only compatible with cpu
    self.im.cpu() if dynamic else self.im,
    f,
    verbose=False,
    opset_version=self.args.opset or get_latest_opset(),
    do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
    input_names=['images'],
    output_names=output_names,
    dynamic_axes=dynamic or None)

# Checks
model_onnx = onnx.load(f)  # load onnx model
# onnx.checker.check_model(model_onnx)  # check onnx model

# Simplify
if self.args.simplify:
    try:
        import onnxsim

        LOGGER.info(f'{prefix} simplifying with onnxsim {onnxsim.__version__}...')
        # subprocess.run(f'onnxsim {f} {f}', shell=True)
        model_onnx, check = onnxsim.simplify(model_onnx)
        assert check, 'Simplified ONNX model could not be validated'
    except Exception as e:
        LOGGER.info(f'{prefix} simplifier failure: {e}')

# Metadata
for k, v in self.metadata.items():
    meta = model_onnx.metadata_props.add()
    meta.key, meta.value = k, str(v)

onnx.save(model_onnx, f)

如果通过动态batch进行导出，我们希望导出输入是batch x 3 x 1920 x 1088 ，输出是batch x 3 x number1 x number2 ,其中 number1和2应该是自动确定的，但是当使用上的代码导出时会发现模型的输入虽然可以，但是输出不对：

想要达到要求，就不能使用动态batch导出，只能使用静态batch导出，具体如下：

import onnxsim
# batch 根据需要进行设置，一般都是1， 这里为了演示使用了10
x = torch.randn([10, 3, 1920, 1088])
print(x.shape)
torch.onnx.export(self.model, x, "Net.onnx", input_names=['images'],output_names=['output0', 'output1'],verbose=False)
model_onnx1 = onnx.load("Net.onnx")  # load onnx model
# 需要使用onnxsim.simplify进行简化才能得到想要的结果
model_onnx1, check = onnxsim.simplify(model_onnx1)
onnx.save(model_onnx1, "Net.onnx")

结果：

打印过程：

ONNX: starting export with onnx 1.13.0...
torch.Size([10, 3, 1920, 1088])
11
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
WARNING: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.
ONNX: simplifying with onnxsim 0.4.17...
ONNX: export success  27.7s, saved as yolov8n-seg-test.onnx (13.1 MB)

Export complete (29.1s)
Results saved to D:\workspace\yolov8\ultralytics
Predict:         yolo predict task=segment model=yolov8n-seg-test.onnx imgsz=640
Validate:        yolo val task=segment model=yolov8n-seg-test.onnx imgsz=640 data=coco.yaml
Visualize:       https://netron.app