train

backend/app.go

@@ -44,7 +44,6 @@ func (a *App) Startup(ctx context.Context) {
 	a.ctx = ctx
 
 	// 设置 ONNX Runtime DLL 路径
-	// ort.SetSharedLibraryPath("onnxruntime.dll")
 	exeDir, _ := os.Executable()
 	println("exeDir", exeDir)
 	ort.SetSharedLibraryPath(filepath.Join(filepath.Dir(exeDir), "onnxruntime.dll"))
@@ -160,7 +159,6 @@ func (a *App) GormDB() (*gorm.DB, error) {
 }
 
 func (a *App) UploadImage(data []byte, filename string) Response {
-	println("UploadImage")
 	uploadsDir := publicImagePath
 	err := os.MkdirAll(uploadsDir, 0755)
 	if err != nil {

core/const/const.py

@@ -26,22 +26,25 @@ if mode == "toy":
     ]
     num_classes = len(label_name)
 elif mode == "benchmark":
-    epoch = 100  # convnext->100, resnet18->50
-    lr = 1e-4
+    epoch = 50  # resnet=80, resnet18=50, convnext=100
+    lr = 1e-4  # resnet, resnet18=1e-4
+    weight_decay = 1e-4  # resnet=1e-3, resnet18=1e-4
     batch_size = 8
     input_size = 224
 
     # 分类
     label_name = [
-        "american_shorthair",
-        "bengal",
-        "british_shorthair",
-        "exotic_shorthair",
-        "maine_coon",
-        "ragdoll",
-        "scottish_fold",
-        "siamese",
-        "sphynx",
-        "turkish_van",
+        "american_shorthair",  # 美国短毛猫
+        "british_shorthair",  # 英国短毛猫
+        "ragdoll",  # 布偶猫
+        "exotic_shorthair",  # 异国短毛猫
+        "maine_coon",  # 缅因猫
+        "siamese",  # 暹罗猫
+        "sphynx",  # 斯芬克斯猫
+        "turkish_van",  # 土耳其梵猫
+        "bengal",  # 孟加拉豹猫
+        "scottish_fold",  # 苏格兰折耳猫
+        "none",  # 风景人物
+        "other",  # 其他动物
     ]
     num_classes = len(label_name)

core/dataloader/dataloader.py

@@ -76,4 +76,5 @@ print("test_dataset", len(test_dataset))
 
 
 train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
-test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
+test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
+

core/static/dataset_stats.csv

@@ -1,67 +0,0 @@
-class,image_count,avg_width,avg_height,min_width,min_height,max_width,max_height,formats,corrupt_files
-abyssinian,200,148,117,88,46,162,140,"jpeg, png",0
-cyprus,200,159,100,93,40,162,140,"jpeg, png",0
-lykoi,200,142,122,68,55,162,140,"jpeg, png",0
-donskoy,200,141,123,67,54,162,140,jpeg,0
-chausie,200,143,119,78,46,162,140,"jpeg, png",0
-european_shorthair,200,149,118,70,49,162,140,"jpeg, png",0
-turkish_van,200,149,120,87,56,162,140,"jpeg, png",0
-pixie_bob,200,147,119,78,49,162,140,"jpeg, png",0
-ragdoll,199,149,116,78,50,162,140,jpeg,0
-german_rex,199,138,125,72,50,162,140,"jpeg, png",0
-american_shorthair,199,147,114,65,50,162,140,"jpeg, png",0
-sokoke,199,145,122,70,49,162,140,"jpeg, png",0
-khao_manee,198,144,120,78,49,162,140,"jpeg, png",0
-thai,198,143,118,68,60,162,140,"jpeg, png",0
-cymric,198,148,120,63,43,162,140,"jpeg, png",0
-oriental_shorthair,197,141,123,78,54,162,140,jpeg,0
-cornish_rex,197,147,117,75,56,162,140,"jpeg, png",0
-burmese,197,147,117,45,58,162,140,"jpeg, png",0
-savannah,197,153,106,78,40,162,140,"jpeg, png",0
-american_wirehair,196,149,117,76,50,162,140,"jpeg, png",0
-peterbald,196,144,123,77,50,162,140,"jpeg, png",0
-karelian_bobtail,196,145,125,67,65,162,140,"jpeg, png",0
-tonkinese,195,148,116,70,53,162,140,"jpeg, png",0
-balinese,195,154,112,79,54,162,140,jpeg,0
-japanese_bobtail,194,149,118,87,49,162,140,"jpeg, png",0
-nebelung,194,143,119,64,46,162,140,jpeg,0
-selkirk_rex,192,144,122,78,76,162,140,jpeg,0
-persian,192,150,114,78,46,162,140,jpeg,0
-manx,192,155,113,86,50,162,140,"jpeg, png",0
-himalayan,192,158,109,93,40,162,140,jpeg,0
-munchkin,191,147,117,61,48,162,140,jpeg,0
-bengal,189,151,115,78,71,162,140,jpeg,0
-turkish_angora,188,145,119,66,52,162,140,jpeg,0
-vankedisi,187,147,116,63,47,162,140,"jpeg, png",0
-scottish_fold,184,143,120,78,50,162,140,jpeg,0
-egyptian_mau,184,144,121,72,74,162,140,"jpeg, png",0
-ocicat,182,150,115,62,44,162,140,"jpeg, png",0
-ragamuffin,182,149,116,78,44,162,140,"jpeg, png",0
-serengeti,175,159,100,78,53,300,140,"jpeg, png",0
-british_shorthair,174,148,117,78,50,162,140,jpeg,0
-toyger,160,145,120,78,50,162,140,"jpeg, png",0
-siberian,159,153,116,78,55,162,140,jpeg,0
-havana_brown,159,133,127,75,34,300,140,"jpeg, png",0
-exotic_shorthair,157,148,118,93,55,300,140,"jpeg, png",0
-bombay,154,139,123,46,50,162,140,"jpeg, png",0
-korat,152,147,117,93,50,162,140,"jpeg, png",0
-safari,150,158,105,79,38,300,140,"jpeg, png",0
-american_bobtail,140,155,114,64,48,162,140,jpeg,0
-mekong_bobtail,140,149,118,44,50,162,140,"jpeg, png",0
-korean_bobtail,139,140,129,63,81,162,140,jpeg,0
-siamese,139,152,115,78,53,300,140,"jpeg, png",0
-somali,139,150,112,61,53,162,140,jpeg,0
-devon_rex,138,150,116,78,55,162,140,jpeg,0
-american_curl,138,155,112,91,51,300,140,"jpeg, png",0
-ural_rex,137,144,121,48,65,162,140,"jpeg, png",0
-singapura,136,158,109,93,54,300,140,"jpeg, png",0
-ukrainian_levkoy,134,140,123,78,59,162,140,jpeg,0
-maine_coon,133,141,122,66,79,162,140,jpeg,0
-birman,131,155,113,93,56,162,140,jpeg,0
-oregon_rex,121,147,123,85,66,300,140,"jpeg, png",0
-kurilian_bobtail,120,149,121,78,81,162,140,jpeg,0
-laperm,120,149,116,78,55,162,140,"jpeg, png",0
-sphynx,120,149,117,47,63,162,140,jpeg,0
-chartreux,114,146,117,75,61,162,140,jpeg,0
-russian_blue,109,152,111,88,36,162,140,jpeg,0
-norwegian_forest_cat,97,150,114,78,56,162,140,jpeg,0

core/test/test_onnx.py

@@ -4,13 +4,12 @@ import glob
 import torch
 import numpy as np
 from PIL import Image
-import onnxruntime as ort  # 【修改1】替换torch导入
-from torchvision import transforms  # 保留transforms，仍用于预处理
+import onnxruntime as ort
+from torchvision import transforms
 from core.const import mode, label_name, input_size
 
 
 def test():
-    # 【修改2】选择ONNX Runtime的执行提供程序，自动选择CPU或CUDA
     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
     print(f"使用设备: {providers[0]}")
 
@@ -19,7 +18,7 @@ def test():
     model_dir = os.path.join(core_dir, "models")
     dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
 
-    # 【修改3】加载ONNX模型，替代原来的PyTorch模型加载
+    # 加载ONNX模型
     session = ort.InferenceSession(os.path.join(model_dir, "resnet_epoch_100.onnx"), providers=providers)
 
     # 获取输入名称（用于后续推理时指定输入）
@@ -28,7 +27,7 @@ def test():
     im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
     np.random.shuffle(im_list)
 
-    # 预处理完全不变
+    # 预处理
     test_transform = transforms.Compose([
         transforms.Resize(input_size),
         transforms.CenterCrop(input_size),
@@ -42,15 +41,15 @@ def test():
         inputs = test_transform(im_data)
         inputs = torch.unsqueeze(inputs, dim=0)  # 这里还在用torch，下面会改
 
-        # 【修改4】将输入转为numpy，ONNX Runtime需要numpy输入
+        # 将输入转为numpy，ONNX Runtime需要numpy输入
         inputs = inputs.numpy()
         if providers[0] == 'CPUExecutionProvider':
             inputs = inputs.astype(np.float32)
 
-        # 【修改5】ONNX Runtime推理，输出直接是numpy数组
+        # ONNX Runtime推理，输出直接是numpy数组
         outputs = session.run(None, {input_name: inputs})[0]
 
-        # 【修改6】解析结果，直接用numpy操作
+        # 解析结果，直接用numpy操作
         pred = np.argmax(outputs, axis=1)
         print(label_name[pred[0]], " ", im_path)
 

core/test/test_resnet.py

@@ -21,7 +21,7 @@ def test():
     print("model_dir", model_dir)
 
     net = resnet()
-    net.load_state_dict(torch.load(os.path.join(model_dir, "resnet_epoch_99.pth"), weights_only=True))
+    net.load_state_dict(torch.load(os.path.join(model_dir, "resnet_epoch_80.pth"), weights_only=True))
 
     print("111")
 

core/test/test_resnet18.py

@@ -18,10 +18,8 @@ def test():
     model_dir = os.path.join(core_dir, "models")
     dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
 
-    print("model_dir", model_dir)
-
     net = resnet18()
-    net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50_bak2.pth"), weights_only=True))
+    net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50.pth"), weights_only=True))
 
     im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
     np.random.shuffle(im_list)
@@ -35,6 +33,8 @@ def test():
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
 
+    ary = []
+    errors = []
     for im_path in im_list:
         net.eval()
         im_data = Image.open(im_path)
@@ -46,13 +46,25 @@ def test():
         outputs = net.forward(inputs)
 
         _, pred = torch.max(outputs.data, dim=1)
-        print(label_name[pred.cpu().numpy()[0]], " ", im_path)
+        # print(label_name[pred.cpu().numpy()[0]], " ", im_path)
 
-        img = np.asarray(im_data)
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        img = cv2.resize(img, (200, 200))
-        cv2.imshow("img", img)
-        cv2.waitKey(0)
+        # img = np.asarray(im_data)
+        # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        # img = cv2.resize(img, (200, 200))
+        # cv2.imshow("img", img)
+        # cv2.waitKey(0)
+        result = label_name[pred.cpu().numpy()[0]]
+
+        if result in im_path:
+            ary.append(True)
+        else:
+            ary.append(False)
+            print(f"{label_name[pred.cpu().numpy()[0]]} {im_path}\n")
+            errors.append(f"{label_name[pred.cpu().numpy()[0]]} {im_path}")
+
+    print("ary", ary)
+    # print("errors", errors)
+    print(sum(ary) / len(ary))
 
 
 if __name__ == "__main__":

core/train/train_resnet.py

@@ -2,7 +2,7 @@ import os
 import torch
 from core.nets.resnet import resnet
 from core.dataloader.dataloader import train_dataloader
-from core.const import epoch, lr, batch_size
+from core.const import epoch, lr, weight_decay, batch_size
 
 
 def train():
@@ -13,9 +13,9 @@ def train():
 
     loss_func = torch.nn.CrossEntropyLoss()
 
-    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
+    optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
 
-    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
 
     for e in range(epoch):
         print("epoch: ", e)

core/train/train_resnet18.py

@@ -2,7 +2,7 @@ import os
 import torch
 from core.nets.resnet18 import resnet18
 from core.dataloader.dataloader import train_dataloader
-from core.const import epoch, lr, batch_size
+from core.const import epoch, lr, weight_decay, batch_size
 
 
 def train():
@@ -13,7 +13,7 @@ def train():
 
     loss_func = torch.nn.CrossEntropyLoss()
 
-    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
+    optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
 
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6) # 余弦退火