1,介绍
torchvision是大家最常用的预训练模型来源,但是其包含的预训练模型种类很少,往往并不能满足研究者们的需求。
而timm库提供了一个更强大的替代选项。
利用如下代码查询
import timm
print(len(timm.list_models()))
输出
1032
可知,timm内置了1032个预训练模型权重。
例如ghostnet、caformer、crossvit等模型都是不在torchvision内的。
2,利用timm创建带预训练权重的模型
model = timm.create_model('ghostnetv2_100.in1k', pretrained=True)
其中权重的名称可以到hugging face hub上查询。
https://huggingface.co/
比如说输入ghostnet,就会出现以下可选的权重,对应不同版本和超参数的ghostnet模型,选择自己需要的就可以。
3,不加载分类头,输出特征图
输出的卷积层特征
model = timm.create_model(
'ghostnet_100.in1k',
pretrained=True,
features_only=True,
)
4,获取图像的嵌入
所谓图像嵌入,也就是不包含最后一层分类层,但是包含了之前的其他线性变换层
model = timm.create_model(
'ghostnet_100.in1k',
pretrained=True,
num_classes=0, # remove classifier nn.Linear
)
另一种等效的方式是:
在创建模型时不需要移除分类层,但是用forward_features()方法来获取图像嵌入。
model = timm.create_model('ghostnetv2_100.in1k', pretrained=True)
#获取的是(b, c, h, w)形状的tensor
output = model.forward_features(transforms(img).unsqueeze(0))
#获取的是(b,c x h x w)形状的tensor
output = model.forward_head(output, pre_logits=True)
forward_head方法一般来说返回的是最后一个分类层之前的输出,如果你的模型在分类头里包含了全局池化层,那么forward_head方法返回的就是分类层之前最后一个全局池化层的结果。
但是具体返回结果需要结合你使用的模型来确认。
5,使用预下载到本地的权重创建模型
由于某些众所周知的原因,timm的模型下载经常无法连接。所以,我们可以事先从hugging face hub上把模型权重文件下载到本地。然后用如下代码创建
model = timm.create_model('ghostnet_100.in1k', pretrained=True,pretrained_cfg_overlay=dict(file='D:/pytorch_model.bin'))
只需要把pretrained_cfg_overlay=dict(file='D:/pytorch_model.bin')中的路径替换成你自己保存模型的路径即可。
6,秒杀torchvision,对非三通道图像的神奇处理能力
如果是使用的是torchvision内的预训练模型来处理非三通道图像,那么会报错。只能通过将自己的数据集预处理,转换成三通道图才可以运行。
但是利用timm创建的模型可以直接输入非三通道图进行训练或者预测。
例如,单通道图
m = timm.create_model('resnet34', pretrained=True, in_chans=1)
或者25通道图
m = timm.create_model('resnet34', pretrained=True, in_chans=25)
我们只需要在创建模型时制定in_chans这个参数就可以了。
剩下的操作,timm都会自动帮你完成。其内部已经实现了对通道数的变换。
7,查看预训练模型的配置
from timm.models.resnet import default_cfgs
resnet34_default_cfg = default_cfgs['resnet34']
print(resnet34_default_cfg)
输出结果太长就不贴了,自行查看。