opengait中common代码+注释

前言

在学习opengait时，发现代码很多都看不懂，因此写了注释，以便日后查看

代码+注释

import copy
import os
import inspect
import logging
import torch
import numpy as np
import torch.nn as nn
import torch.autograd as autograd
import yaml
import random
from torch.nn.parallel import DistributedDataParallel as DDP
from collections import OrderedDict, namedtuple


# 功能：在调用实例的任何属性或方法时不执行任何操作，而是简单地忽略调用
class NoOp:
    # __getattr__(self, *args) 方法是一个特殊方法，当尝试获取实例中不存在的属性时会调用它
    # 在这个类中，__getattr__被重写以返回一个匿名函数 no_op，这个函数在被调用时不做任何事情
    def __getattr__(self, *args):
        def no_op(*args, **kwargs): pass  # 接受任意数量的位置参数和关键字参数，但在函数体中没有任何操作，通过 pass 语句表示空操作

        return no_op


# 该类继承自 OrderedDict 类，它添加了一个 append 方法，该方法用于将另一个字典对象的内容附加到当前字典中
class Odict(OrderedDict):
    def append(self, odict):
        # 获取当前字典的键列表
        dst_keys = self.keys()
        # 遍历要附加的字典对象的键值对
        for k, v in odict.items():
            # 如果值不是列表，则转换为单元素列表
            if not is_list(v):
                v = [v]
            # 如果当前键已经存在于当前字典中
            if k in dst_keys:
                # 如果当前键对应的值是列表，则将要附加的值合并到列表中
                if is_list(self[k]):
                    self[k] += v
                # 如果当前键对应的值不是列表，则将其转换为列表，并将要附加的值添加到列表中
                else:
                    self[k] = [self[k]] + v
            # 如果当前键不存在于当前字典中，则直接将键值对添加到当前字典中
            else:
                self[k] = v


def Ntuple(description, keys, values):
    # 如果 keys 不是列表或元组，则将其转换为单元素列表，并相应调整 values
    if not is_list_or_tuple(keys):
        keys = [keys]
        values = [values]
    # 使用 namedtuple 创建一个新的命名元组类型
    Tuple = namedtuple(description, keys)
    # 使用 values 来填充新创建的命名元组，并返回结果
    return Tuple._make(values)


def get_valid_args(obj, input_args, free_keys=[]):
    """
    从输入参数中提取与函数或类构造函数期望的参数匹配的参数。
    Args:
        obj: 要检查参数的函数或类对象。
        input_args (dict): 输入参数字典。
        free_keys (list, optional): 允许自由传递的参数键列表，默认为空列表。
    Returns:
        dict: 与函数或类构造函数期望的参数匹配的参数字典。
    Raises:
        ValueError: 如果传递的对象不是函数或类对象，则引发 ValueError。
    """
    if inspect.isfunction(obj):
        # 如果传入的对象是函数，则获取函数的参数列表
        expected_keys = inspect.getfullargspec(obj)[0]
    elif inspect.isclass(obj):
        # 如果传入的对象是类对象，则获取类构造函数的参数列表
        expected_keys = inspect.getfullargspec(obj.__init__)[0]
    else:
        # 如果传入的既不是函数也不是类对象，则引发 ValueError
        raise ValueError('Just support function and class object!')

    unexpect_keys = list()
    expected_args = {}

    for k, v in input_args.items():
        if k in expected_keys:
            expected_args[k] = v
        elif k in free_keys:
            pass
        else:
            unexpect_keys.append(k)

    if unexpect_keys != []:
        logging.info("Find Unexpected Args(%s) in the Configuration of - %s -" %
                     (', '.join(unexpect_keys), obj.__name__))

    return expected_args


def get_attr_from(sources, name):
    try:
        return getattr(sources[0], name)
    except:
        return get_attr_from(sources[1:], name) if len(sources) > 1 else getattr(sources[0], name)


# 检查一个对象是否为列表或元组类型
def is_list_or_tuple(x):
    return isinstance(x, (list, tuple))


# 检查一个对象是否为布尔类型
def is_bool(x):
    return isinstance(x, bool)


# 检查一个对象是否为字符串类型
def is_str(x):
    return isinstance(x, str)


# 检查一个对象是否为列表类型或模块列表类型
def is_list(x):
    return isinstance(x, list) or isinstance(x, nn.ModuleList)


# 检查一个对象是否为字典类型、有序字典类型
def is_dict(x):
    return isinstance(x, dict) or isinstance(x, OrderedDict) or isinstance(x, Odict)


# 检查一个对象是否为张量类型
def is_tensor(x):
    return isinstance(x, torch.Tensor)


# 检查一个对象是否为 NumPy 数组类型
def is_array(x):
    return isinstance(x, np.ndarray)


# 将张量转换为 NumPy 数组
def ts2np(x):
    return x.cpu().data.numpy()


# 将张量转换为 PyTorch 变量，并移到 GPU 上
def ts2var(x, **kwargs):
    return autograd.Variable(x, **kwargs).cuda()


# 将 NumPy 数组转换为 PyTorch 变量，并移到 GPU 上。
def np2var(x, **kwargs):
    return ts2var(torch.from_numpy(x), **kwargs)


# 将列表转换为 PyTorch 变量，并移到 GPU 上。
def list2var(x, **kwargs):
    return np2var(np.array(x), **kwargs)


#  如果路径不存在，则创建目录。
def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)


def MergeCfgsDict(src, dst):
    """
    合并两个配置字典。
    Args:
        src (dict): 源配置字典。
        dst (dict): 目标配置字典，合并结果将存储在其中。
    Returns:
        None
    """
    for k, v in src.items():
        # 如果源字典中的键在目标字典中不存在，或者对应的值不是字典类型，则直接将源字典中的键值对添加到目标字典中
        if (k not in dst.keys()) or (type(v) != type(dict())):
            dst[k] = v
        # 如果源字典和目标字典中对应的键都是字典类型，则递归地将这两个字典合并
        else:
            if is_dict(src[k]) and is_dict(dst[k]):
                MergeCfgsDict(src[k], dst[k])
            else:
                dst[k] = v


def clones(module, N):
    """
    复制模块 N 次并返回一个模块列表。
    参数:
        module: 要复制的模块对象。
        N: 要复制的次数。
    返回:
        一个包含 N 个相同模块的模块列表。
    """
    # 使用列表推导式复制模块 N 次，并将结果放入 ModuleList 中返回
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


def config_loader(path):
    """
    加载配置文件并将其与默认配置文件合并。
    参数:
        path: 配置文件的路径。
    返回:
        合并后的配置字典。
    """
    # 打开指定路径的配置文件，并加载其中的配置
    with open(path, 'r') as stream:
        src_cfgs = yaml.safe_load(stream)
    # 打开默认配置文件，并加载其中的配置
    with open("./configs/default.yaml", 'r') as stream:
        dst_cfgs = yaml.safe_load(stream)
    # 将加载的配置与默认配置合并
    MergeCfgsDict(src_cfgs, dst_cfgs)
    # 返回合并后的配置
    return dst_cfgs


def init_seeds(seed=0, cuda_deterministic=True):
    """
    初始化随机种子以实现实验的可重复性，并设置 PyTorch 的 CUDA 相关参数以控制 CUDA 的随机性。
    参数:
        seed: 随机种子值，用于初始化各种随机数生成器。
        cuda_deterministic: 是否启用 CUDA 的确定性计算，影响 CUDA 的随机性和性能。
    注意:
        当 cuda_deterministic 设置为 True 时，CUDA 将会以确定性方式运行，这意味着相同的输入将产生相同的输出，但可能会导致性能下降。
        当 cuda_deterministic 设置为 False 时，CUDA 将以更快的方式运行，但输出可能会有微小的变化。
    """
    # 随机种子值是一个初始的整数值，它作为随机数生成器的起始点。
    # 在计算机程序中，随机数生成器通常不是真正的随机，而是通过一定的算法生成的伪随机数。
    # 这个算法接受一个种子作为输入，并基于种子计算出一系列看似随机的数值。
    # 如果给定相同的种子值，随机数生成器将会生成相同的序列，这样就可以实现可重复的随机性。
    # 设置 Python 标准库中的随机数生成器的种子
    random.seed(seed)
    # 设置 NumPy 库中的随机数生成器的种子
    np.random.seed(seed)
    # 设置 PyTorch 库中的随机数生成器的种子
    torch.manual_seed(seed)
    # 设置所有可用的 CUDA 设备上的随机数生成器的种子
    torch.cuda.manual_seed_all(seed)

    # 根据 cuda_deterministic 参数设置 CUDA 的随机性和性能参数
    if cuda_deterministic:
        # 当启用 CUDA 的确定性计算时，关闭 CUDA 的随机性，以实现更高的重现性
        torch.backends.cudnn.deterministic = True
        # 关闭 CUDA 的性能优化，以实现更高的重现性，但可能会导致性能下降
        torch.backends.cudnn.benchmark = False
    else:
        # 当禁用 CUDA 的确定性计算时，启用 CUDA 的随机性，以获得更快的运行速度
        torch.backends.cudnn.deterministic = False
        # 启用 CUDA 的性能优化，以获得更快的运行速度，但输出可能会有微小的变化
        torch.backends.cudnn.benchmark = True


# 这个好像是Linux中用的，signum 是一个参数，用于指示发生了哪种操作系统信号
def handler(signum, frame):
    logging.info('Ctrl+c/z pressed')
    os.system(
        "kill $(ps aux | grep main.py | grep -v grep | awk '{print $2}') ")
    logging.info('process group flush!')


def ddp_all_gather(features, dim=0, requires_grad=True):
    """
    使用 PyTorch 分布式工具包进行全局聚合
    Args:
        features (Tensor): 当前进程中的特征数据，形状为 [n, ...]
        dim (int, optional): 聚合的维度，默认为0
        requires_grad (bool, optional): 是否需要梯度信息，默认为True
    Returns:
        Tensor: 全局聚合后的特征数据，形状与输入features相同
    """
    # 获取当前分布式环境中的进程总数
    world_size = torch.distributed.get_world_size()
    # 获取当前进程在分布式环境中的唯一标识符（进程的 rank）
    rank = torch.distributed.get_rank()
    # 创建一个列表，用于存储从所有进程收集到的特征数据
    feature_list = [torch.ones_like(features) for _ in range(world_size)]
    # 使用 torch.distributed.all_gather 函数将所有进程中的特征数据收集到 feature_list 列表中
    torch.distributed.all_gather(feature_list, features.contiguous())  # .contiguous() 方法用于返回一个张量的连续版本
    # 如果 requires_grad 为 True，则将当前进程的特征数据替换为输入的 features，以确保只有当前进程的数据具有梯度信息
    if requires_grad:
        feature_list[rank] = features
    # 将收集到的所有特征数据按照指定的维度 dim 进行拼接，得到全局聚合后的特征数据
    feature = torch.cat(feature_list, dim=dim)
    # 返回全局聚合后的特征数据
    return feature


# https://github.com/pytorch/pytorch/issues/16885
class DDPPassthrough(DDP):
    def __getattr__(self, name):
        """
        在属性查找过程中被调用的特殊方法。
        Args:
            name (str): 要查找的属性或方法的名称。
        Returns:
            object: 返回对象的属性或方法。
        """
        try:
            # 尝试从父类中获取属性或方法
            return super().__getattr__(name)
        except AttributeError:
            # 如果父类中不存在对应的属性或方法，则从 self.module 中获取
            return getattr(self.module, name)


# 将一个 PyTorch 模型转换为支持分布式数据并行的模型，以便在多个 GPU 上进行训练
def get_ddp_module(module, find_unused_parameters=False, **kwargs):
    """
    获取用于分布式训练的 DDPPassthrough 模块。
    Args:
        module (nn.Module): 要进行分布式训练的模型。
        find_unused_parameters (bool): 是否查找未使用的参数。默认为 False。
        **kwargs: 其他传递给 DDPPassthrough 构造函数的参数。
    Returns:
        DDPPassthrough: 分布式训练的 DDPPassthrough 模块。
    """
    if len(list(module.parameters())) == 0:
        # 对于没有参数的损失模块的情况，直接返回原始模块
        return module
    device = torch.cuda.current_device()
    # 创建 DDPPassthrough 模块，将当前设备作为唯一的设备，并传递其他参数
    module = DDPPassthrough(module, device_ids=[device], output_device=device,
                            find_unused_parameters=find_unused_parameters, **kwargs)
    return module


# 获取模型参数个数
def params_count(net):
    """
    计算模型的参数数量。
    Args:
        net (nn.Module): 要计算参数数量的模型。
    Returns:
        str: 包含参数数量的字符串，格式为 "Parameters Count: {参数数量}M"。
    """
    # 统计模型参数总数
    n_parameters = sum(p.numel() for p in net.parameters())
    # 将参数数量转换为以百万为单位的格式，并返回包含参数数量的字符串
    return 'Parameters Count: {:.5f}M'.format(n_parameters / 1e6)