torch.nn.utils.convert_parameters 的源代码

import torch
from typing import Iterable, Optional


def parameters_to_vector(parameters: Iterable[torch.Tensor]) -> torch.Tensor:
    r"""将可迭代的参数展平为一个单一的向量。

    参数:
        parameters (Iterable[Tensor]): 一个包含模型参数的Tensor的可迭代对象。

    返回:
        由单一向量表示的参数
    """
    # 参数所在设备的标志
    param_device = None

    vec = []
    for param in parameters:
        # 确保参数位于同一设备上
        param_device = _check_param_device(param, param_device)

        vec.append(param.view(-1))
    return torch.cat(vec)


def vector_to_parameters(vec: torch.Tensor, parameters: Iterable[torch.Tensor]) -> None:
    r"""将向量的切片复制到可迭代的参数中。

    参数:
        vec (Tensor): 表示模型参数的单一向量。
        parameters (Iterable[Tensor]): 一个包含模型参数的Tensor的可迭代对象。
    """
    # 确保vec是Tensor类型
    if not isinstance(vec, torch.Tensor):
        raise TypeError(f'expected torch.Tensor, but got: {torch.typename(vec)}')
    # 参数所在设备的标志
    param_device = None

    # 用于切片向量的指针
    pointer = 0
    for param in parameters:
        # 确保参数位于同一设备上
        param_device = _check_param_device(param, param_device)

        # 参数的长度
        num_param = param.numel()
        # 切片向量，重塑它，并替换参数的旧数据
        param.data = vec[pointer:pointer + num_param].view_as(param).data

        # 增加指针
        pointer += num_param


def _check_param_device(param: torch.Tensor, old_param_device: Optional[int]) -> int:
    r"""检查参数是否位于同一设备上。

    目前，不支持在多个分配之间转换模型参数和单一向量形式，例如不同GPU/PrivateUse1中的参数，或CPU/GPU/PrivateUse1的混合。

    参数:
        param ([Tensor]): 模型的一个参数的Tensor
        old_param_device (int): 模型第一个参数分配的设备。

    返回:
        old_param_device (int): 首次报告设备
    """
    # 遇到第一个参数
    support_device_types = ["cuda", torch._C._get_privateuse1_backend_name()]
    if old_param_device is None:
        old_param_device = param.get_device() if param.device.type in support_device_types else -1
    else:
        warn = False
        if param.device.type in support_device_types:  # 检查是否在同一GPU/PrivateUse1中
            warn = (param.get_device() != old_param_device)
        else:  # 检查是否在CPU中
            warn = (old_param_device != -1)
        if warn:
            raise TypeError('Found two parameters on different devices, '
                            'this is currently not supported.')
    return old_param_device