timm加载模型

timm加载模型

输出模型列表

# 支持模糊查找
timm.list_models('*clip*')

直接加载

resnet34 = timm.create_model("resnet34", pretrained=True,)

下载之后加载

下载的可以是.pth, .bin, .safetensors

首先要查看去哪里下载

resnet34_2 = timm.create_model("resnet34", pretrained=False)

resnet34_2.default_cfg

url可以直接下载.pth文件，使用hf_hub_id可在timm模型页面下载.pth, .bin。其实这些东西都是保持的权重，使用对应的加载方法后，可以使用load_stat_dict()直接加载。但是timm提供了方便的接口。

resnet34_pth_path = '/public/slq/timm_hub/resnet34_a1_0-46f8f793.pth'
resnet34 = timm.create_model("resnet34", pretrained=True, pretrained_cfg_overlay=dict(file=resnet34_pth_path))


resnet18_st_path = '/public/slq/timm_hub/resnet18.a1_in1k/model.safetensors'
resnet18 = timm.create_model("resnet18", pretrained=True, pretrained_cfg_overlay=dict(file=resnet18_st_path))

resnet18_bin_path = '/public/slq/timm_hub/resnet18.a1_in1k/pytorch_model.bin'
resnet18_2 = timm.create_model("resnet18", pretrained=True, pretrained_cfg_overlay=dict(file=resnet18_bin_path))

---

{'url': 'https://github.com/huggingface/pytorch-image-models/releases/download/v0.1-rsb-weights/resnet34_a1_0-46f8f793.pth',
 'hf_hub_id': 'timm/resnet34.a1_in1k',
 'architecture': 'resnet34',
 'tag': 'a1_in1k',
 'custom_load': False,
 'input_size': (3, 224, 224),
 'test_input_size': (3, 288, 288),
 'fixed_input_size': False,
 'interpolation': 'bicubic',
 'crop_pct': 0.95,
 'test_crop_pct': 1.0,
 'crop_mode': 'center',
 'mean': (0.485, 0.456, 0.406),
 'std': (0.229, 0.224, 0.225),
 'num_classes': 1000,
 'pool_size': (7, 7),
 'first_conv': 'conv1',
 'classifier': 'fc',
 'origin_url': 'https://github.com/huggingface/pytorch-image-models',
 'paper_ids': 'arXiv:2110.00476'}

输出模型结构

from torchinfo import summary
summary(resnet34_2, input_size=(5, 3, 224, 224))

模型参数统计

from thop import profile 

flops, params = profile(resnet34, inputs=(input1, ))
print(flops, params)

total_megabytes = params / 1e6
total_params_G = flops / 1e9
print(f'total_megabytes: {total_megabytes} M, total_params: {total_params_G} G')