今天给大家推荐一个硬核干货:一个基于 PyTorch 的图像模型库(PyTorch Image Models,TIMM),用于最新图像分类。
这个库从 330+ 种预训练的最新图像分类模型中进行选择,方便我们使用提供的脚本在 ImageNet 等研究数据集上重新训练模型。而且,可以在自己的数据集上微调预训练的模型,包括最新的前沿模型。
话不多说,直接放上 TIMM 的开源地址:
https://paperswithcode.com/lib/timm
TIMM 包含的模型很丰富,例如:
- Vision Transformer
- RexNet
- CSP DarkNet
- ResNeSt
- TResNet
- RegNetX
- RegNetY
- EfficientNet Pruned
- Big Transfer
- CSP ResNet
下面以 Vision Transformer 为例,看看这个库包含了哪些内容!
Vision Transformer 将输入图片拆分成 16x16个patches,每个 patch 做一次线性变换降维同时嵌入位置信息,然后送入 Transformer,避免了像素级 attention 的运算。
https://paperswithcode.com/lib/timm/vision-transformer
Vision Transformer 支持的模型有 vit_base_patch16_224、vit_base_patch16_384、vit_base_resnet50_384 等等。vit_base_patch16_224 包含的参数为 87 million,FLOPs 为 67 billion,文件大小为 330.25 Mb,训练集为 ImageNet,训练资源为 TPUv3。
详细的训练参数如下:
作者给出了 Paper 地址:
https://arxiv.org/abs/2010.11929v1
完整代码:
权重文件:
如何使用?
导入预训练模型:
import timm m = timm.create_model('vit_large_patch16_224', pretrained=True) m.eval()
用你要使用的型号替换型号名称,例如 vit_large_patch16_224。
如何训练模型?
你可以使用 TIMM 的脚本来重新训练新模型,链接:
https://rwightman.github.io/pytorch-image-models/scripts/
结果比较
不同模型在 ImageNet 上测试的结果如下:
更多精彩的功能还等着大家去挖掘和发现!最后再次放上该库的地址:
https://paperswithcode.com/lib/timm
更多硬核干货,请见阅读原文!