new train

2026-07-28 18:06:10 +08:00
parent 44b0d6b756
commit a55d449028
13 changed files with 224 additions and 27 deletions
@@ -1,24 +1,29 @@
 import os
 import torch
 import sys
-from core.nets.resnet import resnet
+# from core.nets.resnet import resnet
+from core.nets.resnet18 import resnet18

 # 加载 pth
-net = resnet()  # 实例化你的模型
+# net = resnet()  # 实例化你的模型
+net = resnet18()
+

 script_dir = os.path.dirname(os.path.abspath(__file__))
 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, "resnet_epoch_100.pth"), map_location="cpu"))
+
+net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50_bak2.pth"), map_location="cpu"))
 net.eval()

+
 # 导出 ONNX
 dummy_input = torch.randn(1, 3, 224, 224)
 torch.onnx.export(
    net,
    dummy_input,
-    os.path.join(model_dir, "resnet_epoch_100.onnx"),
+    os.path.join(model_dir, "resnet18_epoch_50_bak2.onnx"),
    export_params=True,
    opset_version=11,
    input_names=["input"],