Unverified 提交 4de8b248 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

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 https://github.com/ultralytics/yolov5/issues/6692
上级 de9c25b3
...@@ -9,6 +9,7 @@ import os ...@@ -9,6 +9,7 @@ import os
import platform import platform
import subprocess import subprocess
import time import time
import warnings
from contextlib import contextmanager from contextlib import contextmanager
from copy import deepcopy from copy import deepcopy
from pathlib import Path from pathlib import Path
...@@ -25,6 +26,9 @@ try: ...@@ -25,6 +26,9 @@ try:
except ImportError: except ImportError:
thop = None thop = None
# Suppress PyTorch warnings
warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')
@contextmanager @contextmanager
def torch_distributed_zero_first(local_rank: int): def torch_distributed_zero_first(local_rank: int):
...@@ -293,13 +297,9 @@ class EarlyStopping: ...@@ -293,13 +297,9 @@ class EarlyStopping:
class ModelEMA: class ModelEMA:
""" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models """ Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
Keep a moving average of everything in the model state_dict (parameters and buffers). Keeps a moving average of everything in the model state_dict (parameters and buffers)
This is intended to allow functionality like For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
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.
""" """
def __init__(self, model, decay=0.9999, updates=0): def __init__(self, model, decay=0.9999, updates=0):
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论