pytorch 转onnx模型

# import torch
# import torchvision.models as models
# import torch.onnx

# model = models.resnet18(pretrained=True)
# model.eval()

# # 步骤 2: 提供示例输入数据并运行模型
# dummy_input = torch.randn(1, 3, 640, 640)  # 输入数据的形状，这里使用ResNet的标准输入形状
# output = model(dummy_input)
# # 步骤 3: 将运行时图转换为ONNX图
# onnx_path = "yolox.onnx"
# torch.onnx.export(model, dummy_input, onnx_path, verbose=True, input_names=['input'], output_names=['output'])
# # # 步骤 4: 保存ONNX图到文件
# print(f"ONNX模型已保存到 {onnx_path}")

# import torch 

# batch_size = 1  
# input = torch.rand(batch_size, 3, 256, 256)  # 初始化输入
# input = input.cuda()
# model = model.cuda()  # 加载模型
# model.eval()  # 把模型设置成推理模式
# output = model(input)  # 推理
# print("Beginning to export onnx ...")
# torch.onnx.export(model,  # 要转换的模型
#                   input,  # 输入
#                   "./yolox.onnx", # 到处的模型，以onnx格式保存