Pytorch如何判断两个模型的权重是否相同（比较权重差异/参数字典）

思念殇千寻

于 2024-08-15 07:48:00 发布

阅读量249

点赞数 1

文章标签： pytorch 人工智能 python 深度学习机器学习

本文链接：https://blog.csdn.net/weixin_43590796/article/details/141235226

版权

　　参考资料：

　　GPT-4o mini的回答

　　第一种方法是使用md5sum这个命令(Linux上)，但是由于环境的不同，哪怕是load之后转存似乎都会有差，所以效果不大。

　　第二种方法是使用代码比较，这段代码是我找GPT要的，感觉非常不错。

from safetensors.torch import load_file

A = load_file("/data1/chaishang.cs/Tryon/hf_DATA/almmtryon-refinpaintxl-v0-1/unet/diffusion_pytorch_model.safetensors")
B = load_file("/data1/chaishang.cs/Tryon/datasets_and_weights/exps/refinpaintXL_v1_36cate_batch1_itagged_task1fulltask2partly_bz2_a100/2024-08-14T21-26-40/22000/unet/diffusion_pytorch_model.safetensors")

weight_changes = {}

# 遍历模型的所有参数
for (name, initial_param), (_, trained_param) in zip(A.items(), B.items()):
    # 计算权重变化
    change = torch.norm(trained_param - initial_param)  # 计算 L2 范数
    weight_changes[name] = change.item()  # 存储变化量

# 打印每层权重的变化
for layer_name, change in weight_changes.items():
    print(f"Layer: {layer_name}, Change: {change:.6f}")

　　如果参数是内存中的state_dict，对A和B进行替换即可。输出结果如下：