fix sampler bugs and update dataloader

bubbliiiing · Apr 15, 2022 · 71c94df · 71c94df
1 parent 5524b5e
commit 71c94df
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 12 deletions.
diff --git a/train.py b/train.py
@@ -458,6 +458,12 @@
             gen.dataset.epoch_now       = epoch
             gen_val.dataset.epoch_now   = epoch
 
+            if distributed:
+                train_sampler.set_epoch(epoch)
+
             set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
 
             fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, UnFreeze_Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rank)
+
+        if local_rank == 0:
+            loss_history.writer.close()
diff --git a/utils/dataloader.py b/utils/dataloader.py
@@ -2,6 +2,7 @@
 
 import cv2
 import numpy as np
+import torch
 from PIL import Image
 from torch.utils.data.dataset import Dataset
 
@@ -354,5 +355,6 @@ def yolo_dataset_collate(batch):
     for img, box in batch:
         images.append(img)
         bboxes.append(box)
-    images = np.array(images)
-    return images, bboxes
+    images = torch.from_numpy(np.array(images)).type(torch.FloatTensor)
+    bboxes = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in bboxes]
+    return images, bboxes
diff --git a/utils/utils_fit.py b/utils/utils_fit.py
@@ -21,11 +21,8 @@ def fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch,
         images, targets = batch[0], batch[1]
         with torch.no_grad():
             if cuda:
-                images  = torch.from_numpy(images).type(torch.FloatTensor).cuda()
-                targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets]
-            else:
-                images  = torch.from_numpy(images).type(torch.FloatTensor)
-                targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
+                images  = images.cuda()
+                targets = [ann.cuda() for ann in targets]
         #----------------------#
         #   清零梯度
         #----------------------#
@@ -94,11 +91,8 @@ def fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch,
         images, targets = batch[0], batch[1]
         with torch.no_grad():
             if cuda:
-                images  = torch.from_numpy(images).type(torch.FloatTensor).cuda()
-                targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets]
-            else:
-                images  = torch.from_numpy(images).type(torch.FloatTensor)
-                targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
+                images  = images.cuda()
+                targets = [ann.cuda() for ann in targets]
             #----------------------#
             #   清零梯度
             #----------------------#