快速入门

本快速入门面向准备深入代码、查看如何将 timm 集成到模型训练工作流中的开发者。

首先，您需要安装 timm。有关安装的更多信息，请参阅安装。

pip install timm

加载预训练模型

可以使用 create_model() 加载预训练模型。

在此，我们加载预训练的 mobilenetv3_large_100 模型。

>>> import timm

>>> m = timm.create_model('mobilenetv3_large_100', pretrained=True)
>>> m.eval()

Note: The returned PyTorch model is set to train mode by default, so you must call .eval() on it if you plan to use it for inference.

列出具有预训练权重的模型

要列出 timm 中打包的模型，可以使用 list_models()。如果指定 pretrained=True，该函数将仅返回具有可用预训练权重的模型名称。

>>> import timm
>>> from pprint import pprint
>>> model_names = timm.list_models(pretrained=True)
>>> pprint(model_names)
[
    'adv_inception_v3',
    'cspdarknet53',
    'cspresnext50',
    'densenet121',
    'densenet161',
    'densenet169',
    'densenet201',
    'densenetblur121d',
    'dla34',
    'dla46_c',
]

您还可以通过名称中包含特定模式来列出模型。

>>> import timm
>>> from pprint import pprint
>>> model_names = timm.list_models('*resne*t*')
>>> pprint(model_names)
[
    'cspresnet50',
    'cspresnet50d',
    'cspresnet50w',
    'cspresnext50',
    ...
]

微调预训练模型

你只需更改分类器（最后一层），即可对任何预训练模型进行微调。

>>> model = timm.create_model('mobilenetv3_large_100', pretrained=True, num_classes=NUM_FINETUNE_CLASSES)

要在自己的数据集上进行微调，你需要编写一个 PyTorch 训练循环，或者将 timm 的训练脚本适配为使用你的数据集。

对于不同的输入通道，请将 in_chans 设置为模型的目标通道数；预训练配置必须描述源权重。有关通道转换以及使用当前配置保存适配后的模型，请参阅 Adapt and Export Models。

使用预训练模型进行特征提取

在不修改网络的情况下，可以调用任意模型的 model.forward_features(input) 来替代通常的 model(input)。这将绕过网络的分类头和全局池化。

如需了解使用 timm 进行特征提取的更详细指南，请参阅特征提取。

>>> import timm
>>> import torch
>>> x = torch.randn(1, 3, 224, 224)
>>> model = timm.create_model('mobilenetv3_large_100', pretrained=True)
>>> features = model.forward_features(x)
>>> print(features.shape)
torch.Size([1, 960, 7, 7])

图像增强

要将图像转换为模型的有效输入，您可以使用 timm.data.create_transform()，并提供模型所需的 input_size。

这将返回一个使用合理默认值的通用转换。

>>> timm.data.create_transform((3, 224, 224))
Compose(
    Resize(size=256, interpolation=bilinear, max_size=None, antialias=True)
    CenterCrop(size=(224, 224))
    MaybeToTensor()
    Normalize(mean=tensor([0.4850, 0.4560, 0.4060]), std=tensor([0.2290, 0.2240, 0.2250]))
)

预训练模型在训练时，输入给它们的图像会经过特定的变换处理。如果你对图像使用了错误的变换，模型将无法理解它所看到的内容！

要确定某个预训练模型使用了哪些变换，我们可以先查看其 pretrained_cfg

>>> model.pretrained_cfg
{'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_large_100_ra-f55367f5.pth',
 'hf_hub_id': 'timm/mobilenetv3_large_100.ra_in1k',
 'architecture': 'mobilenetv3_large_100',
 'tag': 'ra_in1k',
 'custom_load': False,
 'input_size': (3, 224, 224),
 'fixed_input_size': False,
 'interpolation': 'bicubic',
 'crop_pct': 0.875,
 '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': 'conv_stem',
 'classifier': 'classifier',
 'license': 'apache-2.0'}

然后，我们可以使用 timm.data.resolve_model_data_config() 仅解析与数据相关的配置。

>>> timm.data.resolve_model_data_config(model)
{'input_size': (3, 224, 224),
 'interpolation': 'bicubic',
 'mean': (0.485, 0.456, 0.406),
 'std': (0.229, 0.224, 0.225),
 'crop_pct': 0.875,
 'crop_mode': 'center'}

我们可以将此数据配置传递给 timm.data.create_transform() 来初始化模型关联的 transform。

>>> data_cfg = timm.data.resolve_model_data_config(model)
>>> transform = timm.data.create_transform(**data_cfg)
>>> transform
Compose(
    Resize(size=256, interpolation=bicubic, max_size=None, antialias=True)
    CenterCrop(size=(224, 224))
    MaybeToTensor()
    Normalize(mean=tensor([0.4850, 0.4560, 0.4060]), std=tensor([0.2290, 0.2240, 0.2250]))
)

Note: Here, the pretrained model's config happens to be the same as the generic config we made earlier. This is not always the case. So, it's safer to use the data config to create the transform as we did here instead of using the generic transform.

使用预训练模型进行推理

在这里，我们将整合上述各部分，并使用预训练模型进行推理。

首先，我们需要一张用于推理的图像。这里我们从网络上加载了一张叶子的图片：

>>> import requests
>>> from PIL import Image
>>> from io import BytesIO
>>> url = 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/timm/cat.jpg'
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> image

这是我们加载的图像：

现在，我们将再次创建模型和转换。这次，我们确保将模型设置为评估模式。

>>> model = timm.create_model('mobilenetv3_large_100', pretrained=True).eval()
>>> transform = timm.data.create_transform(
    **timm.data.resolve_model_data_config(model)
)

我们可以通过将图像传递给 transform 来为模型准备该图像。

>>> image_tensor = transform(image)
>>> image_tensor.shape
torch.Size([3, 224, 224])

现在我们可以将该图像传入模型以获取预测结果。在此情况下，我们使用 unsqueeze(0)，因为模型期望输入包含一个批次维度。

>>> output = model(image_tensor.unsqueeze(0))
>>> output.shape
torch.Size([1, 1000])

要获取预测概率，我们对输出应用 softmax。这将得到一个形状为 (num_classes,) 的张量。

>>> probabilities = torch.nn.functional.softmax(output[0], dim=0)
>>> probabilities.shape
torch.Size([1000])

现在我们将使用 torch.topk 找到预测类别索引和值的前 5 名。

>>> values, indices = torch.topk(probabilities, 5)
>>> indices
tensor([281, 282, 285, 673, 670])

如果我们查看 top index 对应的 imagenet 标签，就能看到模型的预测结果……

>>> IMAGENET_1k_URL = 'https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt'
>>> IMAGENET_1k_LABELS = requests.get(IMAGENET_1k_URL).text.strip().split('\n')
>>> [{'label': IMAGENET_1k_LABELS[idx], 'value': val.item()} for val, idx in zip(values, indices)]
[{'label': 'tabby, tabby_cat', 'value': 0.5101025700569153},
 {'label': 'tiger_cat', 'value': 0.22490699589252472},
 {'label': 'Egyptian_cat', 'value': 0.1835290789604187},
 {'label': 'mouse, computer_mouse', 'value': 0.006752475164830685},
 {'label': 'motor_scooter, scooter', 'value': 0.004942195490002632}]

Update on GitHub