对比两个模型权重是否一致 torch.equal

compare_models.py

import torch
model1 = torch.load('output/exp-step00000100/unet/diffusion_pytorch_model.bin')
print('load model1 success')
model1 = torch.load('output/exp-step00000100/unet/diffusion_pytorch_model.bin')
print('load model2 success')
def compare_models(dict1, dict2):
    models_differ = False
    for (key1, item1), (key2, item2) in zip(dict1.items(), dict2.items()):
        if key1 == key2 and torch.equal(item1, item2):
            pass
        else:
            models_differ = True
            print('Mismatch found in item:', key1)
    return not models_differ
is_same = compare_models(model1, model2)
print('Are the models the same?', is_same)