Add PyTorch AMP check (#7917)

* Add PyTorch AMP check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup

* Cleanup

* Cleanup

* Robust for DDP

* Fixes

* Add amp enabled boolean to check_train_batch_size

* Simplify

* space to prefix

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

utils/autobatch.py

@@ -7,15 +7,14 @@ from copy import deepcopy
 
 import numpy as np
 import torch
-from torch.cuda import amp
 
 from utils.general import LOGGER, colorstr
 from utils.torch_utils import profile
 
 
-def check_train_batch_size(model, imgsz=640):
+def check_train_batch_size(model, imgsz=640, amp=True):
     # Check YOLOv5 training batch size
-    with amp.autocast():
+    with torch.cuda.amp.autocast(amp):
         return autobatch(deepcopy(model).train(), imgsz)  # compute optimal batch size
 
 

utils/general.py

@@ -36,9 +36,11 @@ import yaml
 from utils.downloads import gsutil_getsize
 from utils.metrics import box_iou, fitness
 
-# Settings
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[1]  # YOLOv5 root directory
+RANK = int(os.getenv('RANK', -1))
+
+# Settings
 DATASETS_DIR = ROOT.parent / 'datasets'  # YOLOv5 datasets directory
 NUM_THREADS = min(8, max(1, os.cpu_count() - 1))  # number of YOLOv5 multiprocessing threads
 AUTOINSTALL = str(os.getenv('YOLOv5_AUTOINSTALL', True)).lower() == 'true'  # global auto-install mode
@@ -505,6 +507,27 @@ def check_dataset(data, autodownload=True):
     return data  # dictionary
 
 
+def check_amp(model):
+    # Check PyTorch Automatic Mixed Precision (AMP) functionality. Return True on correct operation
+    from models.common import AutoShape
+
+    if next(model.parameters()).device.type == 'cpu':  # get model device
+        return False
+    prefix = colorstr('AMP: ')
+    im = cv2.imread(ROOT / 'data' / 'images' / 'bus.jpg')[..., ::-1]  # OpenCV image (BGR to RGB)
+    m = AutoShape(model, verbose=False)  # model
+    a = m(im).xyxy[0]  # FP32 inference
+    m.amp = True
+    b = m(im).xyxy[0]  # AMP inference
+    if (a.shape == b.shape) and torch.allclose(a, b, atol=1.0):  # close to 1.0 pixel bounding box
+        LOGGER.info(emojis(f'{prefix}checks passed ✅'))
+        return True
+    else:
+        help_url = 'https://github.com/ultralytics/yolov5/issues/7908'
+        LOGGER.warning(emojis(f'{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}'))
+        return False
+
+
 def url2file(url):
     # Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt
     url = str(Path(url)).replace(':/', '://')  # Pathlib turns :// -> :/