import torch.nn as nn
# 步骤1: 加载ONNX模型
onnx_path = r"E:\2022Work\RedCellCount\红细胞网络\除100\tt.pth"
model = torch.load(onnx_path)
print(model)
del model.model_resnet.fc
model.model_resnet.fc = nn.Identity()
print(model)

model.eval()
dummy_input = torch.randn(1, 3, 512, 512).to('cuda')
input_names = ['input']
output_names = ['fetures']

yy = model(dummy_input)

onnx_path = r"E:\2022Work\RedCellCount\红细胞网络\除100\tt1.onnx"
torch.onnx.export(model, dummy_input, onnx_path,export_params=True, verbose=False, input_names=input_names, output_names=output_names, opset_version=9,
                    dynamic_axes={
                      "input": {0: "batch_size"},
                      "fetures": {0: "batch_size"}
                    }
                    )