YOLOv6-4.0部分代码阅读笔记-envs.py
envs.py
yolov6\utils\envs.py
目录
envs.py
1.所需的库和模块
2.def get_envs():
3.def select_device(device):
4.def set_random_seed(seed, deterministic=False):
1.所需的库和模块
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import random
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from yolov6.utils.events import LOGGER
2.def get_envs():
def get_envs():
# 从系统环境中获取 PyTorch 所需的环境。
"""Get PyTorch needed environments from system envirionments."""
# 获取名为 LOCAL_RANK 的环境变量,并将其值转换为整数。如果该环境变量不存在,则默认值为 -1 。
local_rank = int(os.getenv('LOCAL_RANK', -1))
# 获取名为 RANK 的环境变量,并将其值转换为整数。如果该环境变量不存在,则默认值为 -1 。
rank = int(os.getenv('RANK', -1))
# 获取名为 WORLD_SIZE 的环境变量,并将其值转换为整数。如果该环境变量不存在,则默认值为 1 。
world_size = int(os.getenv('WORLD_SIZE', 1))
# 函数返回一个包含三个整数值的元组,分别代表 local_rank 、 rank 和 world_size 。
return local_rank, rank, world_size
3.def select_device(device):
# 用于设置程序使用的设备信息,可以是 CPU 或者特定的 GPU 设备。
# 1.device :该参数是一个字符串,用于指定设备类型。
def select_device(device):
# 将设备的信息设置到程序中。
"""Set devices' information to the program.
Args:
device: a string, like 'cpu' or '1,2,3,4'
Returns:
torch.device
"""
# 如果 device 参数为 'cpu' ,则设置环境变量 CUDA_VISIBLE_DEVICES 为 '-1' ,这会导致 CUDA 不可见,即不使用任何 GPU。
if device == 'cpu':
# os.environ
# os.environ 是 Python 标准库 os 模块中的一个特殊对象,它表示当前运行环境中的环境变量。 environ 是一个全局变量,它是一个代表用户环境的字典对象,其中包含了所有环境变量的键值对。
# 以下是 os.environ 的一些常用操作 :
# 1. 读取环境变量 :你可以通过 os.environ['VAR_NAME'] 来获取名为 VAR_NAME 的环境变量的值。如果环境变量不存在,会引发 KeyError 。
# 2. 设置环境变量 :可以通过赋值操作来设置或修改环境变量,例如 os.environ['VAR_NAME'] = 'value' 。
# 3. 删除环境变量 :使用 del os.environ['VAR_NAME'] 可以删除名为 VAR_NAME 的环境变量。
# 4. 遍历环境变量 :可以通过 for key in os.environ 来遍历所有的环境变量。
# 5. 检查环境变量是否存在 :使用 if 'VAR_NAME' in os.environ 来检查某个环境变量是否存在。
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
# 使用日志记录器输出信息,表明程序将使用 CPU 进行训练。
LOGGER.info('Using CPU for training... ')
# 如果 device 参数非空,说明用户指定了 GPU 设备。
elif device:
# 设置环境变量 CUDA_VISIBLE_DEVICES 为用户指定的设备字符串,例如 '1,2,3,4' 。
os.environ['CUDA_VISIBLE_DEVICES'] = device
# 确保 CUDA 可用,即系统支持 CUDA。
assert torch.cuda.is_available()
# 计算用户指定的 GPU 数量。
nd = len(device.strip().split(','))
# 使用日志记录器输出信息,表明程序将使用指定数量的 GPU 进行训练。
LOGGER.info(f'Using {nd} GPU for training... ')
# 判断是否使用 CUDA 设备。
cuda = device != 'cpu' and torch.cuda.is_available()
# 根据是否使用 CUDA 设备,创建一个 torch.device 对象,指定使用 'cuda:0' 或 'cpu' 。
device = torch.device('cuda:0' if cuda else 'cpu')
# 函数返回一个 torch.device 对象,表示程序将使用的设备。
return device
4.def set_random_seed(seed, deterministic=False):
# 是在不同的库中设置随机种子,以确保实验的可重复性。
# 1.seed :一个整数值,用作随机种子。
# 2.deterministic :一个布尔值,指示是否设置确定性行为,默认为 False 。
def set_random_seed(seed, deterministic=False):
# 将随机状态设置为随机库、numpy、pytorch 和 cudnn。
""" Set random state to random libray, numpy, torch and cudnn.
Args:
seed: int value.
deterministic: bool value.
"""
# 为 Python 标准库中的 random 模块设置随机种子。
random.seed(seed)
# 为 NumPy 库设置随机种子。
np.random.seed(seed)
# 为 PyTorch 库设置随机种子。
torch.manual_seed(seed)
# 如果 deterministic 参数为 True 。
if deterministic:
# 设置 PyTorch 的 CUDA 神经网络库(cuDNN)为确定性模式。这会限制某些 CUDA 操作的使用,以确保结果的可重复性,但可能会影响性能。
cudnn.deterministic = True
# 关闭 cuDNN 的基准测试功能,因为在确定性模式下不需要。
cudnn.benchmark = False
# 如果 deterministic 参数为 False 。
else:
# 关闭确定性模式,允许 cuDNN 使用所有 CUDA 操作。
cudnn.deterministic = False
# 开启 cuDNN 的基准测试功能,这可以提高性能,因为它允许 cuDNN 自动选择最优的算法。
cudnn.benchmark = True