修改训练代码，兼容三个模型的训练。

2025-10-16 15:26:58 +08:00
parent e74c42b5d4
commit 25f8e0a267
2 changed files with 123 additions and 30 deletions
@@ -52,4 +52,7 @@ NUM_CLASSES = 4

 # 其他配置
 NUM_WORKERS = 0  # Windows下建议设为0
-DEVICE = 'cuda'  # 'cuda' 或 'cpu'，程序会自动检测可用性
+DEVICE = 'cuda'  # 'cuda' 或 'cpu'，程序会自动检测可用性
+
+# 中心损失权重
+CENTER_LOSS_WEIGHT = 0.1