1. timm库
timm 是一个用于深度学习的开源库,全称是 "PyTorch Image Models"。该库由 Ross Wightman 创建并维护,旨在提供高效且易于使用的图像模型,包括大量预训练的模型和实用工具。timm 库基于 PyTorch 框架,主要特点包括:
- 预训练模型:timm 提供了大量预训练的图像分类模型,这些模型在各种数据集上进行过训练,可以直接用于迁移学习或作为研究的基线模型。
- 高效的模型实现:timm 的模型实现注重效率和性能,适用于训练和推理任务。
- 易于使用的接口:timm 提供了简洁的 API,使用户可以方便地加载、训练和评估模型。
- 持续更新:timm 库不断更新,引入新的模型架构和技术,以跟上深度学习领域的最新进展。
下面是一个使用 timm 库加载预训练模型并进行推理的简单示例:
python
import timm
import torch
from PIL import Image
from torchvision import transforms
# 加载预训练的 ResNet50 模型
model = timm.create_model('resnet50', pretrained=True)
model.eval()
# 图像预处理
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 加载并预处理图像
img = Image.open('path_to_your_image.jpg')
img_tensor = preprocess(img).unsqueeze(0)
# 推理
with torch.no_grad():
output = model(img_tensor)
# 输出结果
print(output)
2. 预训练模型
timm
(PyTorch Image Models)库包含了众多预训练的图像分类模型,这些模型在各种流行的数据集上进行了训练。以下是一些主要的预训练模型类别和具体模型名称:
主流模型架构
-
ResNet 系列
- ResNet-18
- ResNet-34
- ResNet-50
- ResNet-101
- ResNet-152
-
EfficientNet 系列
- EfficientNet-B0
- EfficientNet-B1
- EfficientNet-B2
- EfficientNet-B3
- EfficientNet-B4
- EfficientNet-B5
- EfficientNet-B6
- EfficientNet-B7
-
Vision Transformers (ViT)
- ViT-Base
- ViT-Large
- ViT-Huge
-
DeiT(Data-efficient Image Transformers)
- DeiT-Tiny
- DeiT-Small
- DeiT-Base
-
MobileNet 系列
- MobileNetV2
- MobileNetV3
-
RegNet 系列
- RegNetX-200MF
- RegNetX-600MF
- RegNetY-800MF
其他常见模型架构
-
DenseNet 系列
- DenseNet-121
- DenseNet-169
- DenseNet-201
- DenseNet-264
-
Inception 系列
- InceptionV3
- InceptionV4
- Inception-ResNet-V2
-
SENet 系列
- SENet-154
- SE-ResNet-50
- SE-ResNet-101
-
EfficientDet(用于目标检测)
- EfficientDet-D0
- EfficientDet-D1
- EfficientDet-D2
- EfficientDet-D3
- EfficientDet-D4
- EfficientDet-D5
- EfficientDet-D6
- EfficientDet-D7
最新和高性能模型
-
NFNet 系列
- NFNet-F0
- NFNet-F1
- NFNet-F2
- NFNet-F3
- NFNet-F4
- NFNet-F5
- NFNet-F6
-
ConvNeXt 系列
- ConvNeXt-Tiny
- ConvNeXt-Small
- ConvNeXt-Base
- ConvNeXt-Large
-
Swin Transformer
- Swin-Tiny
- Swin-Small
- Swin-Base
- Swin-Large
-
CaiT(Class-Attention in Image Transformers)
- CaiT-XXS
- CaiT-XS
- CaiT-S
- CaiT-M
- CaiT-L
使用方法
可以通过以下代码查看 timm
支持的所有预训练模型:
python
import timm
# 列出所有可用的模型
model_names = timm.list_models(pretrained=True)
print(model_names)
这些预训练模型已经在ImageNet等大型数据集上进行了训练,因此在迁移学习任务中通常表现良好。选择适合你任务的模型架构,可以加快训练过程,并提高模型的性能。