timm 和 torchvision 中的 resnet50

最新推荐文章于 2024-09-16 10:24:41 发布

wujpbb7

最新推荐文章于 2024-09-16 10:24:41 发布

阅读量1.9k

点赞数 1

CC 4.0 BY-SA版权

分类专栏：模式识别文章标签：深度学习 resnet50

本文链接：https://blog.youkuaiyun.com/blueblood7/article/details/127850489

模式识别专栏收录该内容

86 篇文章

订阅专栏

该博客展示了如何使用timm和torchvision库分别加载预训练的ResNet50模型，并将它们导出为ONNX格式。通过比较导出的ONNX模型，作者发现权重是相同的。这为在不同框架间转移模型提供了一种验证方法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

从 timm 和 torchvision 分别加载 resnet50 预训练模型，

import torch
def export_onnx(model_saved, onnx_save_name, input_name='img', output_name='logits'):
    dummy_input = torch.randn(1, 3, 224, 224)
    dynamic_axes = dict()
    dynamic_axes[input_name] = {0:"batch_size"}
    dynamic_axes[output_name] = {0:"batch_size"}
    torch.onnx.export(model_saved, dummy_input, onnx_save_name,
        input_names=[input_name], output_names=[output_name],
        export_params=True, verbose=False, opset_version=12,
        dynamic_axes=dynamic_axes)
    
if __name__ == '__main__':
    import torchvision
    net = torchvision.models.resnet50(pretrained=True)
    export_onnx(net, './resnet50_torchvision.onnx')
    
    import timm
    net = timm.create_model('resnet50', pretrained=True)
    export_onnx(net, './resnet50_timm.onnx')

从 onnx 看，权重是一样的。