一、添加set_seeds函数
在utils/general.py
文件找到init_seeds
函数
def init_seeds(seed=0):
# Initialize random number generator (RNG) seeds
random.seed(seed)
np.random.seed(seed)
init_torch_seeds(seed)
在该函数上方添加新的函数set_seeds
,代码如下:
import pkg_resources as pkg
def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False, verbose=False):
# Check version vs. required version
current, minimum = (pkg.parse_version(x) for x in (current, minimum))
result = (current == minimum) if pinned else (current >= minimum) # bool
return result
def set_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)
二、开启随机种子设置
在train.py
中Ctrl+F
搜索init_seeds
函数并将其注释,导入set_seeds
函数:
# init_seeds(2 + rank)
set_seeds(2 + rank, deterministic=True)
这样随机种子就设置完成,确保每次运行时的随机数生成是可重复的。不过请注意,设置种子只能保证模型在相同的种子值下生成相同的随机数序列,但不能保证模型的训练结果完全相同。(少量数据集下训练结果是完全相同的)