瀏覽代碼

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
modifyDataloader
Glenn Jocher GitHub 2 年之前
父節點
當前提交
4de8b24881
沒有發現已知的金鑰在資料庫的簽署中 GPG 金鑰 ID: 4AEE18F83AFDEB23
共有 1 個檔案被更改,包括 7 行新增7 行删除
  1. +7
    -7
      utils/torch_utils.py

+ 7
- 7
utils/torch_utils.py 查看文件

@@ -9,6 +9,7 @@ import os
import platform
import subprocess
import time
import warnings
from contextlib import contextmanager
from copy import deepcopy
from pathlib import Path
@@ -25,6 +26,9 @@ try:
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):
@@ -293,13 +297,9 @@ class EarlyStopping:


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):

Loading…
取消
儲存