ultralytics 8.0.104 bug fixes and thop dependency removal (#2665)

Co-authored-by: Kevin Abraham <5976139+abraha2d@users.noreply.github.com>
Co-authored-by: Kevin Abraham <abraha2d@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>

ultralytics/yolo/data/utils.py

@@ -324,6 +324,7 @@ class HUBDatasetStats():
 
     def __init__(self, path='coco128.yaml', task='detect', autodownload=False):
         """Initialize class."""
+        LOGGER.info(f'Starting HUB dataset checks for {path}....')
         zipped, data_dir, yaml_path = self._unzip(Path(path))
         try:
             # data = yaml_load(check_yaml(yaml_path))  # data dict

ultralytics/yolo/engine/trainer.py

@@ -9,7 +9,7 @@ import os
 import subprocess
 import time
 from copy import deepcopy
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
 
 import numpy as np
@@ -181,8 +181,6 @@ class BaseTrainer:
             # Command
             cmd, file = generate_ddp_command(world_size, self)
             try:
-                LOGGER.info('Pre-caching dataset to avoid NCCL timeout before running DDP command')
-                deepcopy(self)._setup_train(world_size=0)
                 LOGGER.info(f'Running DDP command {cmd}')
                 subprocess.run(cmd, check=True)
             except Exception as e:
@@ -197,7 +195,11 @@ class BaseTrainer:
         torch.cuda.set_device(RANK)
         self.device = torch.device('cuda', RANK)
         LOGGER.info(f'DDP settings: RANK {RANK}, WORLD_SIZE {world_size}, DEVICE {self.device}')
-        dist.init_process_group('nccl' if dist.is_nccl_available() else 'gloo', rank=RANK, world_size=world_size)
+        os.environ['NCCL_BLOCKING_WAIT'] = '1'  # set to enforce timeout
+        dist.init_process_group('nccl' if dist.is_nccl_available() else 'gloo',
+                                timeout=timedelta(seconds=3600),
+                                rank=RANK,
+                                world_size=world_size)
 
     def _setup_train(self, world_size):
         """

ultralytics/yolo/utils/torch_utils.py

@@ -11,7 +11,6 @@ from pathlib import Path
 from typing import Union
 
 import numpy as np
-import thop
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -21,6 +20,11 @@ import torchvision
 from ultralytics.yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, __version__
 from ultralytics.yolo.utils.checks import check_version
 
+try:
+    import thop
+except ImportError:
+    thop = None
+
 TORCHVISION_0_10 = check_version(torchvision.__version__, '0.10.0')
 TORCH_1_9 = check_version(torch.__version__, '1.9.0')
 TORCH_1_11 = check_version(torch.__version__, '1.11.0')
@@ -193,7 +197,7 @@ def get_flops(model, imgsz=640):
         p = next(model.parameters())
         stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32  # max stride
         im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
-        flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1E9 * 2  # stride GFLOPs
+        flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1E9 * 2 if thop else 0  # stride GFLOPs
         imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz]  # expand if int/float
         flops = flops * imgsz[0] / stride * imgsz[1] / stride  # 640x640 GFLOPs
         return flops
@@ -378,7 +382,7 @@ def profile(input, ops, n=10, device=None):
             m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
             tf, tb, t = 0, 0, [0, 0, 0]  # dt forward, backward
             try:
-                flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1E9 * 2  # GFLOPs
+                flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1E9 * 2 if thop else 0  # GFLOPs
             except Exception:
                 flops = 0
 

ultralytics/yolo/v8/classify/train.py

@@ -19,6 +19,8 @@ class ClassificationTrainer(BaseTrainer):
         if overrides is None:
             overrides = {}
         overrides['task'] = 'classify'
+        if overrides.get('imgsz') is None:
+            overrides['imgsz'] = 224
         super().__init__(cfg, overrides, _callbacks)
 
     def set_model_attributes(self):
@@ -40,10 +42,6 @@ class ClassificationTrainer(BaseTrainer):
         for p in model.parameters():
             p.requires_grad = True  # for training
 
-        # Update defaults
-        if self.args.imgsz == 640:
-            self.args.imgsz = 224
-
         return model
 
     def setup_model(self):