PyTorch保存再加载的模型与原模型==判定为False的原因是什么

核心原因 #

PyTorch的nn.Module类默认没有重载==运算符，你直接对两个模型实例执行==比较时，实际比较的是两个对象在内存中的存储地址，而非模型的结构、参数数值。
你实例化的model_cpu和加载得到的model_load是两个完全独立的对象，存储在不同的内存空间，哪怕二者的结构、参数完全一致，==判定也会返回False。你观察到的结构打印一致、训练效果相同才是模型真实状态的表现，保存加载过程本身没有改变模型的任何内容。

正确验证模型一致性的方法 #

如果要确认两个模型的参数、状态完全相同，应该比较二者的state_dict（包含所有可学习参数、以及BN层的运行统计量等全部状态信息）：

# 方式1：仅比较可学习参数
param_is_same = all(torch.allclose(p1, p2) for p1, p2 in zip(model_cpu.parameters(), model_load.parameters()))
print(param_is_same) # 输出为True

# 方式2：更严谨，比较全部状态（包含BN层的running_mean、running_var等非参数状态）
state_is_same = all(torch.allclose(model_cpu.state_dict()[k], model_load.state_dict()[k]) for k in model_cpu.state_dict().keys())
print(state_is_same) # 输出为True

自定义==比较逻辑 #

如果你需要直接用==对模型实例做等值判定，可以在你的模型类中重载__eq__方法：

class test_model(nn.Module):
    # 原有__init__、forward方法保持不变
    def __eq__(self, other):
        # 先判断类型是否匹配
        if not isinstance(other, test_model):
            return False
        # 比较全部参数是否一致
        return all(torch.allclose(p1, p2) for p1, p2 in zip(self.parameters(), other.parameters()))

重载后再执行model_load == model_cpu就会返回True。

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