From c5f5b80c04585a9df6dcde88c072dbff5ed1cf90 Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Fri, 25 Nov 2022 02:06:14 +0530 Subject: [PATCH] Deterministic training (#53) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- ultralytics/yolo/engine/trainer.py | 7 +++++-- ultralytics/yolo/utils/configs/default.yaml | 1 + ultralytics/yolo/utils/torch_utils.py | 17 +++++++++++++++++ 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 4b0ac60..8f1f7b8 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -28,16 +28,19 @@ from ultralytics.yolo.utils import LOGGER, ROOT, TQDM_BAR_FORMAT from ultralytics.yolo.utils.checks import print_args from ultralytics.yolo.utils.files import increment_path, save_yaml from ultralytics.yolo.utils.modeling import get_model -from ultralytics.yolo.utils.torch_utils import ModelEMA, de_parallel, one_cycle +from ultralytics.yolo.utils.torch_utils import ModelEMA, de_parallel, init_seeds, one_cycle DEFAULT_CONFIG = ROOT / "yolo/utils/configs/default.yaml" +RANK = int(os.getenv('RANK', -1)) class BaseTrainer: def __init__(self, config=DEFAULT_CONFIG, overrides={}): - self.console = LOGGER self.args = self._get_config(config, overrides) + init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic) + + self.console = LOGGER self.validator = None self.model = None self.callbacks = defaultdict(list) diff --git a/ultralytics/yolo/utils/configs/default.yaml b/ultralytics/yolo/utils/configs/default.yaml index a57a6c2..fdb5d39 100644 --- a/ultralytics/yolo/utils/configs/default.yaml +++ b/ultralytics/yolo/utils/configs/default.yaml @@ -22,6 +22,7 @@ pretrained: False optimizer: 'SGD' # choices=['SGD', 'Adam', 'AdamW', 'RMSProp'] verbose: False seed: 0 +deterministic: True local_rank: -1 single_cls: False # train multi-class data as single-class image_weights: False # use weighted image selection for training diff --git a/ultralytics/yolo/utils/torch_utils.py b/ultralytics/yolo/utils/torch_utils.py index b48160d..645b860 100644 --- a/ultralytics/yolo/utils/torch_utils.py +++ b/ultralytics/yolo/utils/torch_utils.py @@ -1,11 +1,13 @@ import math import os import platform +import random import time from contextlib import contextmanager from copy import deepcopy from pathlib import Path +import numpy as np import thop import torch import torch.distributed as dist @@ -199,6 +201,21 @@ def one_cycle(y1=0.0, y2=1.0, steps=100): return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1 +def init_seeds(seed=0, deterministic=False): + # Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) # for Multi-GPU, exception safe + # torch.backends.cudnn.benchmark = True # AutoBatch problem https://github.com/ultralytics/yolov5/issues/9287 + if deterministic and check_version(torch.__version__, '1.12.0'): # https://github.com/ultralytics/yolov5/pull/8213 + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.deterministic = True + os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' + os.environ['PYTHONHASHSEED'] = str(seed) + + class ModelEMA: """ 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)