PyTorch load_state_dict()无法加载精确值问题咨询

我尝试把PyTorch模型的所有参数设置为常量72114982，但执行代码后发现所有参数实际变成了72114984，和预期值差了2。

模型定义 #

import torch.nn as nn

class Net(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(2, 2, 2)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(2, 2)

我使用的赋值代码 #

import torch

model = Net()
params = model.state_dict()

for k, v in params.items():
    params[k] = torch.full(v.shape, 72114982, dtype=torch.long) 

model.load_state_dict(params)
print(model.state_dict().values())

原因分析 #

问题出在模型参数的 dtype 和你赋值时指定的 dtype 不匹配：

• PyTorch的Conv2d、Linear等层默认参数 dtype 是float32；
• 你用torch.long（int64）创建了常量张量，调用load_state_dict时，PyTorch会自动将传入的张量转换为模型参数本身的 dtype；
• 72114982是8位整数，而float32的有效位数只有约7位，无法精确表示这个数，最接近的可表示值就是72114984，因此出现了数值偏移。

解决方法 #

方法一：匹配原参数的 dtype 赋值

创建常量张量时，直接使用模型参数原本的 dtype，避免类型转换带来的精度丢失：

model = Net()
params = model.state_dict()

for k, v in params.items():
    params[k] = torch.full(v.shape, 72114982, dtype=v.dtype) 

model.load_state_dict(params)
print(model.state_dict().values())

方法二：修改模型参数的 dtype 为支持精确整数的类型

如果不需要参数是浮点型，可以先把整个模型的参数 dtype 转换为torch.long，再赋值：

model = Net()
model = model.to(dtype=torch.long)  # 将所有参数转为int64类型
params = model.state_dict()

for k, v in params.items():
    params[k] = torch.full(v.shape, 72114982, dtype=torch.long) 

model.load_state_dict(params)
print(model.state_dict().values())

内容的提问来源于stack exchange，提问作者Random Seed