Browse Source

调整训练脚本,更改权重文件保存路径

liyan 1 year ago
parent
commit
2589f8d9d3
1 changed files with 1 additions and 1 deletions
  1. 1 1
      block/train_get.py

+ 1 - 1
block/train_get.py

@@ -139,7 +139,7 @@ def train_get(args, data_dict, model_dict, loss):
             model_dict['train_loss'] = train_loss
             model_dict['val_loss'] = val_loss
             model_dict['val_m_ap'] = m_ap
-            torch.save(model_dict, 'last.pt' if not args.prune else 'prune_last.pt')  # 保存最后一次训练的模型
+            torch.save(model_dict, args.weight if not args.prune else args.prune_save)  # 保存最后一次训练的模型
             if m_ap > 0.1 and m_ap > model_dict['standard']:
                 model_dict['standard'] = m_ap
                 save_path = args.save_path if not args.prune else args.prune_save