Быстрый старт

Этот быстрый старт предназначен для разработчиков, готовых погрузиться в код и посмотреть пример того, как интегрировать 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 в целевое количество каналов модели; предобученная конфигурация должна описывать исходные веса. См. Адаптация и экспорт моделей для преобразования каналов и сохранения адаптированных моделей с их текущей конфигурацией.

Использование предобученной модели для извлечения признаков

Без изменения сети можно вызвать 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(), чтобы инициализировать связанное с моделью преобразование.

>>> 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 мы найдём индексы и значения пяти наиболее вероятных классов.

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

Если мы проверим метки 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