Skip to content

Commit 4de8b24

Browse files
authored
Suppress torch AMP-CPU warnings (#6706)
This is a torch bug, but they seem unable or unwilling to fix it so I'm creating a suppression in YOLOv5. Resolves #6692
1 parent de9c25b commit 4de8b24

File tree

1 file changed

+7
-7
lines changed

1 file changed

+7
-7
lines changed

utils/torch_utils.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import platform
1010
import subprocess
1111
import time
12+
import warnings
1213
from contextlib import contextmanager
1314
from copy import deepcopy
1415
from pathlib import Path
@@ -25,6 +26,9 @@
2526
except ImportError:
2627
thop = None
2728

29+
# Suppress PyTorch warnings
30+
warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')
31+
2832

2933
@contextmanager
3034
def torch_distributed_zero_first(local_rank: int):
@@ -293,13 +297,9 @@ def __call__(self, epoch, fitness):
293297

294298

295299
class ModelEMA:
296-
""" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
297-
Keep a moving average of everything in the model state_dict (parameters and buffers).
298-
This is intended to allow functionality like
299-
https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
300-
A smoothed version of the weights is necessary for some training schemes to perform well.
301-
This class is sensitive where it is initialized in the sequence of model init,
302-
GPU assignment and distributed training wrappers.
300+
""" Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
301+
Keeps a moving average of everything in the model state_dict (parameters and buffers)
302+
For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
303303
"""
304304

305305
def __init__(self, model, decay=0.9999, updates=0):

0 commit comments

Comments
 (0)