Update train.py

bubbliiiing · web-flow · commit 9b9ed8a40caa · 2020-06-04T20:35:58.000+08:00
diff --git a/train.py b/train.py
@@ -34,17 +34,18 @@ def get_anchors(anchors_path):
 def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda):
     total_loss = 0
     val_loss = 0
-    for iteration in range(epoch_size):
+    for iteration, batch in enumerate(gen):
+        if iteration == epoch_size:
+            break
         start_time = time.time()
-        images, targets = next(gen)
+        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]
-        # print(images)
         optimizer.zero_grad()
         outputs = net(images)
         losses = []
@@ -61,8 +62,10 @@ def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epo
         print('iter:' + str(iteration) + '/' + str(epoch_size) + ' || Total Loss: %.4f || %.4fs/step' % (total_loss/(iteration+1),waste_time))
 
     print('Start Validation')
-    for iteration in range(epoch_size_val):
-        images_val, targets_val = next(genval)
+    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:
