lots change

core/train/to_onnx.py

@@ -1,11 +1,8 @@
 import os
 import torch
-import sys
-# from core.nets.resnet import resnet
 from core.nets.resnet18 import resnet18
 
 # 加载 pth
-# net = resnet()  # 实例化你的模型
 net = resnet18()
 
 
@@ -14,7 +11,7 @@ core_dir = os.path.dirname(script_dir)
 model_dir = os.path.join(core_dir, "models")
 
 
-net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50_bak2.pth"), map_location="cpu"))
+net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50.pth"), map_location="cpu"))
 net.eval()
 
 
@@ -23,7 +20,7 @@ dummy_input = torch.randn(1, 3, 224, 224)
 torch.onnx.export(
     net,
     dummy_input,
-    os.path.join(model_dir, "resnet18_epoch_50_bak2.onnx"),
+    os.path.join(model_dir, "resnet18_epoch_50.onnx"),
     export_params=True,
     opset_version=11,
     input_names=["input"],