今天给大家推荐一个硬核干货：一个基于 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

完整代码：

https://github.com/rwightman/pytorch-image-models/blob/5f9aff395c224492e9e44248b15f44b5cc095d9c/timm/models/vision_transformer.py#L503

权重文件：

https://github-releases.githubusercontent.com/168799526/65360900-1a09-11eb-8b86-f0a014a6f156?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210301%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210301T142533Z&X-Amz-Expires=300&X-Amz-Signature=dd5f6097d17aa535123ad040855f737354442b69683756b0dcdf8f5fda8da6d8&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=168799526&response-content-disposition=attachment%3B%20filename%3Djx_vit_base_p16_224-80ecf9dd.pth&response-content-type=application%2Foctet-stream

如何使用？

导入预训练模型：

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

更多硬核干货，请见阅读原文！