Update train_with_tensorboard.py

Author: bubbliiiing

train_with_tensorboard.py

@@ -15,6 +15,7 @@
 from nets.yolo_training import YOLOLoss,Generator
 from nets.yolo4 import YoloBody
 from tensorboardX import SummaryWriter
+from tqdm import tqdm
 
 #---------------------------------------------------#
 #   获得类和先验框
@@ -33,76 +34,91 @@ def get_anchors(anchors_path):
     anchors = [float(x) for x in anchors.split(',')]
     return np.array(anchors).reshape([-1,3,2])[::-1,:,:]
 
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
 def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda,writer):
     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()
-        # 将loss写入tensorboard，每一步都写
-        writer.add_scalar('Train_loss', loss, (epoch*epoch_size + iteration))
-
-        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()
-    # 将loss写入tensorboard，下面注释的是每个世代保存一次
-    # writer.add_scalar('Train_loss', total_loss/(iteration+1), epoch)
-
-    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写入tensorboard, 下面注释的是每一步都写
-            # writer.add_scalar('Val_loss',val_loss/(epoch_size_val+1), (epoch*epoch_size_val + iteration))
+            loss.backward()
+            optimizer.step()
+            # 将loss写入tensorboard，每一步都写
+            writer.add_scalar('Train_loss', loss, (epoch*epoch_size + iteration))
+
+            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()
+        
+    # 将loss写入tensorboard，下面注释的是每个世代保存一次
+    # writer.add_scalar('Train_loss', total_loss/(iteration+1), epoch)
+
+    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
+                # 将loss写入tensorboard, 下面注释的是每一步都写
+                # writer.add_scalar('Val_loss',val_loss/(epoch_size_val+1), (epoch*epoch_size_val + iteration))
+
+            pbar.set_postfix(**{'total_loss': val_loss.item() / (iteration + 1)})
+            pbar.update(1)
+            
     # 将loss写入tensorboard，每个世代保存一次
     writer.add_scalar('Val_loss',val_loss/(epoch_size_val+1), epoch)
     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)))
 
 
-
 if __name__ == "__main__":
     #-------------------------------#
     #   输入的shape大小
@@ -248,4 +264,4 @@ def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epo
 
         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,writer)
-            lr_scheduler.step()
+            lr_scheduler.step()