Update train.py

Author: bubbliiiing

train.py

@@ -177,6 +177,12 @@
         lr          = Freeze_lr
         start_epoch = Init_Epoch
         end_epoch   = Freeze_Epoch
+                        
+        epoch_step      = num_train // batch_size
+        epoch_step_val  = num_val // batch_size
+        
+        if epoch_step == 0 or epoch_step_val == 0:
+            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
 
         optimizer       = optim.Adam(model_train.parameters(), lr, weight_decay = 5e-4)
         if Cosine_lr:
@@ -190,12 +196,6 @@
                                     drop_last=True, collate_fn=yolo_dataset_collate)
         gen_val         = DataLoader(val_dataset  , shuffle = True, batch_size = batch_size, num_workers = num_workers, pin_memory=True, 
                                     drop_last=True, collate_fn=yolo_dataset_collate)
-                        
-        epoch_step      = num_train // batch_size
-        epoch_step_val  = num_val // batch_size
-        
-        if epoch_step == 0 or epoch_step_val == 0:
-            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
 
         #------------------------------------#
         #   冻结一定部分训练
@@ -214,6 +214,12 @@
         lr          = Unfreeze_lr
         start_epoch = Freeze_Epoch
         end_epoch   = UnFreeze_Epoch
+                        
+        epoch_step      = num_train // batch_size
+        epoch_step_val  = num_val // batch_size
+        
+        if epoch_step == 0 or epoch_step_val == 0:
+            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
 
         optimizer       = optim.Adam(model_train.parameters(), lr, weight_decay = 5e-4)
         if Cosine_lr:
@@ -227,12 +233,6 @@
                                     drop_last=True, collate_fn=yolo_dataset_collate)
         gen_val         = DataLoader(val_dataset  , shuffle = True, batch_size = batch_size, num_workers = num_workers, pin_memory=True, 
                                     drop_last=True, collate_fn=yolo_dataset_collate)
-                        
-        epoch_step      = num_train // batch_size
-        epoch_step_val  = num_val // batch_size
-        
-        if epoch_step == 0 or epoch_step_val == 0:
-            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
 
         #------------------------------------#
         #   冻结一定部分训练