Skip to content

Commit

Permalink
Suppress torch AMP-CPU warnings (ultralytics#6706)
Browse files Browse the repository at this point in the history
This is a torch bug, but they seem unable or unwilling to fix it so I'm creating a suppression in YOLOv5. 

Resolves ultralytics#6692
  • Loading branch information
glenn-jocher committed Feb 19, 2022
1 parent 4c1cfe8 commit a054a88
Showing 1 changed file with 7 additions and 7 deletions.
14 changes: 7 additions & 7 deletions utils/torch_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import platform
import subprocess
import time
import warnings
from contextlib import contextmanager
from copy import deepcopy
from pathlib import Path
Expand All @@ -25,6 +26,9 @@
except ImportError:
thop = None

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


@contextmanager
def torch_distributed_zero_first(local_rank: int):
Expand Down Expand Up @@ -293,13 +297,9 @@ def __call__(self, epoch, fitness):


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

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

0 comments on commit a054a88

Please sign in to comment.