Update train.py

train.py

@@ -14,7 +14,7 @@ from torch.utils.data import DataLoader
 from utils.dataloader import yolo_dataset_collate, YoloDataset
 from nets.yolo_training import YOLOLoss,Generator
 from nets.yolo4 import YoloBody
-
+from tqdm import tqdm
 
 #---------------------------------------------------#
 #   获得类和先验框
@@ -33,67 +33,78 @@ def get_anchors(anchors_path):
     anchors = [float(x) for x in anchors.split(',')]
     return np.array(anchors).reshape([-1,3,2])[::-1,:,:]
 
-def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda):
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+def fit_one_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda):
     total_loss = 0
     val_loss = 0
     start_time = time.time()
-    for iteration, batch in enumerate(gen):
-        if iteration >= epoch_size:
-            break
-        images, targets = batch[0], batch[1]
-        with torch.no_grad():
-            if cuda:
-                images = Variable(torch.from_numpy(images).type(torch.FloatTensor)).cuda()
-                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
-            else:
-                images = Variable(torch.from_numpy(images).type(torch.FloatTensor))
-                targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
-        optimizer.zero_grad()
-        outputs = net(images)
-        losses = []
-        for i in range(3):
-            loss_item = yolo_losses[i](outputs[i], targets)
-            losses.append(loss_item[0])
-        loss = sum(losses)
-        loss.backward()
-        optimizer.step()
-
-        total_loss += loss
-        waste_time = time.time() - start_time
-        print('\nEpoch:'+ str(epoch+1) + '/' + str(Epoch))
-        print('iter:' + str(iteration) + '/' + str(epoch_size) + ' || Total Loss: %.4f || %.4fs/step' % (total_loss/(iteration+1),waste_time))
-        start_time = time.time()
-
-    print('Start Validation')
-    for iteration, batch in enumerate(genval):
-        if iteration >= epoch_size_val:
-            break
-        images_val, targets_val = batch[0], batch[1]
-
-        with torch.no_grad():
-            if cuda:
-                images_val = Variable(torch.from_numpy(images_val).type(torch.FloatTensor)).cuda()
-                targets_val = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets_val]
-            else:
-                images_val = Variable(torch.from_numpy(images_val).type(torch.FloatTensor))
-                targets_val = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets_val]
+    with tqdm(total=epoch_size,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
+        for iteration, batch in enumerate(gen):
+            if iteration >= epoch_size:
+                break
+            images, targets = batch[0], batch[1]
+            with torch.no_grad():
+                if cuda:
+                    images = Variable(torch.from_numpy(images).type(torch.FloatTensor)).cuda()
+                    targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
+                else:
+                    images = Variable(torch.from_numpy(images).type(torch.FloatTensor))
+                    targets = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets]
             optimizer.zero_grad()
-            outputs = net(images_val)
+            outputs = net(images)
             losses = []
             for i in range(3):
-                loss_item = yolo_losses[i](outputs[i], targets_val)
+                loss_item = yolo_losses[i](outputs[i], targets)
                 losses.append(loss_item[0])
             loss = sum(losses)
-            val_loss += loss
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss
+            waste_time = time.time() - start_time
+            
+            pbar.set_postfix(**{'total_loss': total_loss.item() / (iteration + 1), 
+                                'lr'        : get_lr(optimizer),
+                                'step/s'    : waste_time})
+            pbar.update(1)
+
+            start_time = time.time()
+
+    print('Start Validation')
+    with tqdm(total=epoch_size_val, desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
+        for iteration, batch in enumerate(genval):
+            if iteration >= epoch_size_val:
+                break
+            images_val, targets_val = batch[0], batch[1]
+
+            with torch.no_grad():
+                if cuda:
+                    images_val = Variable(torch.from_numpy(images_val).type(torch.FloatTensor)).cuda()
+                    targets_val = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets_val]
+                else:
+                    images_val = Variable(torch.from_numpy(images_val).type(torch.FloatTensor))
+                    targets_val = [Variable(torch.from_numpy(ann).type(torch.FloatTensor)) for ann in targets_val]
+                optimizer.zero_grad()
+                outputs = net(images_val)
+                losses = []
+                for i in range(3):
+                    loss_item = yolo_losses[i](outputs[i], targets_val)
+                    losses.append(loss_item[0])
+                loss = sum(losses)
+                val_loss += loss
+            pbar.set_postfix(**{'total_loss': val_loss.item() / (iteration + 1)})
+            pbar.update(1)
+
     print('Finish Validation')
-    print('\nEpoch:'+ str(epoch+1) + '/' + str(Epoch))
+    print('Epoch:'+ str(epoch+1) + '/' + str(Epoch))
     print('Total Loss: %.4f || Val Loss: %.4f ' % (total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
 
     print('Saving state, iter:', str(epoch+1))
     torch.save(model.state_dict(), 'logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.pth'%((epoch+1),total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
 
-
-
 #----------------------------------------------------#
 #   检测精度mAP和pr曲线计算参考视频
 #   https://www.bilibili.com/video/BV1zE411u7Vw
@@ -209,7 +220,7 @@ if __name__ == "__main__":
             param.requires_grad = False
 
         for epoch in range(Init_Epoch,Freeze_Epoch):
-            fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch,Cuda)
+            fit_one_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch,Cuda)
             lr_scheduler.step()
 
     if True:
@@ -246,5 +257,5 @@ if __name__ == "__main__":
             param.requires_grad = True
 
         for epoch in range(Freeze_Epoch,Unfreeze_Epoch):
-            fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Unfreeze_Epoch,Cuda)
+            fit_one_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,gen_val,Unfreeze_Epoch,Cuda)
             lr_scheduler.step()