更新预测功能

train.py

@@ -181,8 +181,8 @@ if __name__ == "__main__":
     #                           Adam可以使用相对较小的UnFreeze_Epoch
     #   Unfreeze_batch_size     模型在解冻后的batch_size
     #------------------------------------------------------------------#
-    UnFreeze_Epoch      = 500
-    Unfreeze_batch_size = 6
+    UnFreeze_Epoch      = 100
+    Unfreeze_batch_size = 4
     #------------------------------------------------------------------#
     #   Freeze_Train    是否进行冻结训练
     #                   默认先冻结主干训练后解冻训练。
@@ -237,14 +237,14 @@ if __name__ == "__main__":
     #                   开启后会加快数据读取速度，但是会占用更多内存
     #                   内存较小的电脑可以设置为2或者0  
     #------------------------------------------------------------------#
-    num_workers         = 20
+    num_workers         = 6
 
     #------------------------------------------------------#
     #   train_annotation_path   训练图片路径和标签
     #   val_annotation_path     验证图片路径和标签
     #------------------------------------------------------#
-    train_annotation_path   = '2007_train.txt'
-    val_annotation_path     = '2007_val.txt'
+    train_annotation_path   = 'model_data/2007_train.txt'
+    val_annotation_path     = 'model_data/2007_val.txt'
 
     seed_everything(seed)
     #------------------------------------------------------#
@@ -261,6 +261,7 @@ if __name__ == "__main__":
             print("Gpu Device Count : ", ngpus_per_node)
     else:
         device          = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        print("\033[1;33;44mRuning on {}\033[0m".format(device))
         local_rank      = 0
         rank            = 0