修改配置和路径。

train/train_embedding.py

@@ -48,8 +48,10 @@ TASKS = {
         embedding_dim=512,
         batch_size=16,
         lr=1e-3,
-        triplet_margin=0.3,
-        center_loss_weight=0.1,
+        # triplet_margin=0.3,
+        triplet_margin=0.5,
+        # center_loss_weight=0.1,
+        center_loss_weight=0.5,
         aug_strength="medium",
     ),
     "whole_ingredient": TaskConfig(
@@ -59,7 +61,8 @@ TASKS = {
         embedding_dim=512,
         batch_size=32,
         lr=8e-4,
-        triplet_margin=0.35,
+        # triplet_margin=0.35,
+        triplet_margin=0.5,
         center_loss_weight=0.1,
         aug_strength="medium",
     ),
@@ -385,7 +388,7 @@ class EarlyStopping:
 
 
 def train_epoch(model, train_loader, triplet_criterion, center_criterion, 
-                optimizer, center_optimizer, device, epoch):
+                optimizer, center_optimizer, device, epoch,CENTER_LOSS_WEIGHT):
     """
     训练一个epoch
     
@@ -430,7 +433,11 @@ def train_epoch(model, train_loader, triplet_criterion, center_criterion,
         center_loss = center_criterion(anchor_emb, labels)
         
         # 总损失
-        loss = triplet_loss + settings.CENTER_LOSS_WEIGHT * center_loss  # 可配置的中心损失权重
+        # loss = triplet_loss + settings.CENTER_LOSS_WEIGHT * center_loss  # 可配置的中心损失权重
+        loss = triplet_loss + CENTER_LOSS_WEIGHT * center_loss  # 可配置的中心损失权重
+        # print('triplet_loss',triplet_loss)
+        # print('center_loss',CENTER_LOSS_WEIGHT * center_loss)
+        # print('CENTER_LOSS_WEIGHT',CENTER_LOSS_WEIGHT)
         
         # 反向传播
         optimizer.zero_grad()
@@ -709,7 +716,7 @@ def main(task_key: str = "dish"):
         # 训练
         train_loss, train_triplet_loss, train_center_loss = train_epoch(
             model, train_loader, triplet_criterion, center_criterion,
-            optimizer, center_optimizer, device, epoch
+            optimizer, center_optimizer, device, epoch,CENTER_LOSS_WEIGHT
         )
         
         # 验证
@@ -811,4 +818,5 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("task", choices=list(TASKS.keys()), nargs="?", default="dish")
     args = parser.parse_args()
-    main("whole_ingredient")
+    # main("whole_ingredient")
+    main("dish")