From 70fe2204b4250c238be4c32e65f8038a297059cf Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 21 Mar 2019 14:48:40 +0200
Subject: [PATCH] multi_thread dataloader

---
 models.py         |  3 ---
 train.py          | 19 ++++++++++++-------
 utils/datasets.py | 11 ++++++++---
 utils/utils.py    |  2 +-
 4 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/models.py b/models.py
index ef417603..67e8a83d 100755
--- a/models.py
+++ b/models.py
@@ -174,9 +174,6 @@ class Darknet(nn.Module):
         self.module_defs[0]['cfg'] = cfg_path
         self.module_defs[0]['height'] = img_size
         self.hyperparams, self.module_list = create_modules(self.module_defs)
-        self.img_size = img_size
-        self.loss_names = ['loss', 'xy', 'wh', 'conf', 'cls', 'nT']
-        self.losses = []
 
     def forward(self, x, var=None):
         img_size = x.shape[-1]
diff --git a/train.py b/train.py
index 9efbb7f4..52f22b3e 100644
--- a/train.py
+++ b/train.py
@@ -1,6 +1,8 @@
 import argparse
 import time
 
+from torch.utils.data import DataLoader
+
 import test  # Import test.py to get mAP after each epoch
 from models import *
 from utils.datasets import *
@@ -17,6 +19,7 @@ def train(
         accumulate=1,
         multi_scale=False,
         freeze_backbone=False,
+        num_workers=0
 ):
     weights = 'weights' + os.sep
     latest = weights + 'latest.pt'
@@ -38,10 +41,11 @@ def train(
     lr0 = 0.001  # initial learning rate
     optimizer = torch.optim.SGD(model.parameters(), lr=lr0, momentum=.9)
 
-    # Get dataloader
-    dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, augment=True)
-    # from torch.utils.data import DataLoader
-    # dataloader = DataLoader(dataloader, batch_size=batch_size, num_workers=1)
+    # Dataloader
+    if num_workers > 0:
+        cv2.setNumThreads(0)  # to prevent OpenCV from multithreading
+    dataset = LoadImagesAndLabels(train_path, img_size=img_size, augment=True)
+    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
 
     cutoff = -1  # backbone reaches to cutoff layer
     start_epoch = 0
@@ -102,7 +106,6 @@ def train(
         ui = -1
         rloss = defaultdict(float)
         for i, (imgs, targets, _, _) in enumerate(dataloader):
-
             if targets.shape[1] == 100:  # multithreaded 100-size block
                 targets = targets.view((-1, 6))
                 targets = targets[targets[:, 5].nonzero().squeeze()]
@@ -150,8 +153,8 @@ def train(
 
             # Multi-Scale training (320 - 608 pixels) every 10 batches
             if multi_scale and (i + 1) % 10 == 0:
-                dataloader.img_size = random.choice(range(10, 20)) * 32
-                print('multi_scale img_size = %g' % dataloader.img_size)
+                dataset.img_size = random.choice(range(10, 20)) * 32
+                print('multi_scale img_size = %g' % dataset.img_size)
 
         # Update best loss
         if rloss['total'] < best_loss:
@@ -194,6 +197,7 @@ if __name__ == '__main__':
     parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608')
     parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
     parser.add_argument('--resume', action='store_true', help='resume training flag')
+    parser.add_argument('--num_workers', type=int, default=0, help='number of Pytorch DataLoader workers')
     opt = parser.parse_args()
     print(opt, end='\n\n')
 
@@ -208,4 +212,5 @@ if __name__ == '__main__':
         batch_size=opt.batch_size,
         accumulate=opt.accumulate,
         multi_scale=opt.multi_scale,
+        num_workers=opt.num_workers
     )
diff --git a/utils/datasets.py b/utils/datasets.py
index 31fc166b..2280c6a1 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -7,7 +7,6 @@ import cv2
 import numpy as np
 import torch
 
-# from torch.utils.data import Dataset
 from utils.utils import xyxy2xywh
 
 
@@ -114,10 +113,11 @@ class LoadImagesAndLabels:  # for training
 
     def __getitem__(self, index):
         imgs, labels0, img_paths, img_shapes = self.load_images(index, index + 1)
-        labels0[:,0] = index % self.batch_size
 
+        labels0[:, 0] = index % self.batch_size
         labels = torch.zeros(100, 6)
         labels[:min(len(labels0), 100)] = labels0  # max 100 labels per image
+
         return imgs.squeeze(0), labels, img_paths, img_shapes
 
     def __next__(self):
@@ -225,7 +225,12 @@ class LoadImagesAndLabels:  # for training
         img_all = np.ascontiguousarray(img_all, dtype=np.float32)  # uint8 to float32
         img_all /= 255.0  # 0 - 255 to 0.0 - 1.0
 
-        labels_all = torch.from_numpy(np.concatenate(labels_all, 0))
+        if len(labels_all) > 0:
+            labels_all = np.concatenate(labels_all, 0)
+        else:
+            labels_all = np.zeros((1, 6), dtype='float32')
+
+        labels_all = torch.from_numpy(labels_all)
         return torch.from_numpy(img_all), labels_all, img_paths, img_shapes
 
     def __len__(self):
diff --git a/utils/utils.py b/utils/utils.py
index 1e5e691f..7c3d9929 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -40,7 +40,7 @@ def model_info(model):
     print('\n%5s %38s %9s %12s %20s %12s %12s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
     for i, (name, p) in enumerate(model.named_parameters()):
         name = name.replace('module_list.', '')
-        print('%5g %38s %9s %12g %20s %12.3g %12.3g' % (
+        print('%5g %40s %9s %12g %20s %10.3g %10.3g' % (
             i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
     print('Model Summary: %g layers, %g parameters, %g gradients' % (i + 1, n_p, n_g))