Модели

На этой странице представлена справочная документация API для создания и поиска моделей timm, их предобученных конфигураций и вспомогательных средств для чекпоинтов / Hub, а также обзор интерфейса, общего для моделей timm. Отдельные архитектуры описаны в сводках по моделям и на страницах отдельных моделей.

Использование

Модели создаются по имени с помощью create_model. Имя выбирает архитектуру, зарегистрированную в timm, при необходимости с тегом предобучения после . (resnet50.a1_in1k), либо модель из Hugging Face Hub или из локальной папки с помощью префикса источника. Именованные аргументы, которые create_model не использует, передаются через функцию-точку входа модели в конструктор класса модели, поэтому можно задать любой аргумент, специфичный для модели.

import torch
import timm

# Random init (pretrained=False is the default). The tag selects the pretrained cfg, and the weights loaded with
# pretrained=True. Without a tag the architecture's default tag is used.
model = timm.create_model('vit_tiny_patch16_224.augreg_in21k_ft_in1k')
print(model.pretrained_cfg['mean'], model.pretrained_cfg['crop_pct'])  # (0.5, 0.5, 0.5) 0.9

# Model arguments pass through to the model class, a 10 class head, single channel input and regularization.
model = timm.create_model('resnet50', num_classes=10, in_chans=1, drop_rate=0.2, drop_path_rate=0.05)

# Discover models, filters are fnmatch wildcards.
timm.list_models('convnext*')  # architecture names
timm.list_pretrained('resnet50.*')  # architecture.tag names w/ pretrained weights
timm.list_models(module='vision_transformer', pretrained=True)

# Image preprocessing to match the model's pretrained cfg (input size, interpolation, crop, mean / std).
model = timm.create_model('vit_tiny_patch16_224.augreg_in21k_ft_in1k')
data_config = timm.data.resolve_model_data_config(model)
transform = timm.data.create_transform(**data_config, is_training=False)

x = torch.randn(2, 3, 224, 224)
features = model.forward_features(x)  # unpooled features, (2, 197, 192)
pre_logits = model.forward_head(features, pre_logits=True)  # pooled features, (2, 192)
logits = model.forward_head(features)  # same as model(x), (2, 1000)

model.reset_classifier(5)  # replace the classifier w/ a new 5 class one
model.set_input_size(img_size=(384, 384))  # resample the position embedding for a new resolution
logits = model(torch.randn(2, 3, 384, 384))  # (2, 5)

Скрипт обучения (train.py) предоставляет общие аргументы через эти флаги:

• --model выбирает модель.
• --pretrained загружает предобученные веса.
• --num-classes задаёт размер классификатора (num_classes).
• --in-chans задаёт количество входных каналов (in_chans).
• --gp задаёт тип глобального пулинга (global_pool).
• --drop задаёт dropout классификатора (drop_rate).
• --drop-path задаёт стохастическую глубину (drop_path_rate).
• --pretrained-path загружает предобученные веса из локального файла с адаптацией классификатора / входных свёрток (pretrained_cfg_overlay=dict(file=...)).
• --initial-checkpoint загружает контрольную точку после создания модели, без адаптации (checkpoint_path).
• --model-kwargs key=value ... передаёт любые другие аргументы модели.

См. быстрый старт для введения.

Названия моделей

Hub и локальные источники читают конфигурацию даже при pretrained=False, веса загружаются только при pretrained=True. Имя без префикса источника, содержащее /, \ или :, вызывает ValueError вместо использования модели с тем же именем из реестра. hf_hub: — устаревший псевдоним для hf-hub:.

Общие аргументы

• pretrained загружает веса предобученной конфигурации из state_dict или file, если они заданы, иначе из Hub (hf_hub_id), если установлен huggingface_hub, или из url. Адаптация классификатора / входного свёрточного слоя ниже позволяет создать предобученную модель для дообучения одним вызовом.
• num_classes задаёт размер классификатора, 0 удаляет классификатор. Значение по умолчанию берётся из предобученного cfg. При загрузке предобученных весов, если количество классов отличается, веса классификатора отбрасываются (остаётся случайная инициализация).
• global_pool задаёт тип пулинга; варианты зависят от модели (например, avg, max, avgmax, catavgmax для большинства свёрточных сетей, token, avg, map для ViT), '' отключает пулинг.
• in_chans задаёт количество входных каналов. Значение по умолчанию берётся из предобученного cfg. При загрузке предобученных весов первые веса свёртки адаптируются (если это поддерживается), если указанное значение отличается от количества каналов в исходном чекпоинте (input_size[0] в предобученном cfg).
• drop_rate задаёт dropout для классификатора, поддерживается большинством моделей.
• drop_path_rate задаёт стохастическую глубину, поддерживаемую большинством моделей.
• img_size — это аргумент для моделей в стиле ViT и других моделей с фиксированным размером входа; по умолчанию они настраиваются под input_size из предобученной конфигурации, если только он не задан явно.
• dynamic_img_size=True в семействах ViT и EVA выполняет ресемплинг позиционного эмбеддинга под размер входа на каждом проходе forward.
• features_only (с out_indices, output_stride, feature_cls) оборачивает модель как фичевый бэкбон, возвращающий карты признаков, см. извлечение признаков.
• pretrained_cfg заменяет предобученный cfg (тег, словарь или PretrainedCfg).
• pretrained_cfg_overlay переопределяет некоторые из его полей, например, pretrained_cfg_overlay=dict(file='weights.pth') загружает веса с pretrained=True из локального файла, с адаптацией классификатора / входного свёрточного слоя, описанной выше (как это делает --pretrained-path в train.py).
• checkpoint_path загружает чекпоинт с помощью load_checkpoint после создания модели. Загрузка выполняется строго, без адаптации; если в чекпоинте есть веса EMA, используются они.
• cache_dir переопределяет кэш загрузок для предобученных весов и конфигураций Hub.
• device передаётся в класс модели, параметры создаются непосредственно на устройстве. device='meta' создаёт модель без выделения памяти для её весов.
• dtype передаётся в класс модели, параметры создаются сразу в этом типе данных.
• scriptable устанавливает флаги конфигурации слоя для скриптинга TorchScript.
• exportable устанавливает флаги конфигурации слоя для экспорта / трассировки.
• no_jit устанавливает флаги конфигурации слоёв для отключения слоёв, подготовленных с помощью jit.

Аргументы, установленные в None, отбрасываются до передачи в модель, поэтому применяются значения по умолчанию модели. Аргументы, которые класс модели не принимает, вызывают TypeError в большинстве моделей; некоторые игнорируют неизвестные аргументы.

Конфигурация предобученной модели

Архитектуры имеют один или несколько предобученных cfg, по одному на каждый тег, зарегистрированных с моделью. Cfg описывает веса (местоположение, лицензия, происхождение) и то, как модель была обучена и как следует оценивать её с их использованием. Модели, созданные с помощью create_model, имеют свой cfg в виде словаря в model.pretrained_cfg (model.default_cfg является псевдонимом), версии из реестра возвращаются в виде экземпляров PretrainedCfg с помощью get_pretrained_cfg.

Поля ввода — это то, к чему resolve_data_config / resolve_model_data_config приводит конфигурацию данных для набора данных, трансформаций и загрузчика. Параметр cfg описывает исходную контрольную точку и не обновляется при создании модели или её изменении (num_classes, in_chans, img_size, set_input_size()). Текущие входные каналы, размер и ограничение фиксированного размера модели отслеживаются отдельно: get_model_input_config возвращает их, а resolve_input_data_config приводит конфигурацию данных к этим значениям (как это делают скрипты обучения, валидации и инференса). get_model_args возвращает аргументы для пересоздания модели в её текущем виде, которые save_for_hf записывает в экспортированную конфигурацию. Подробности см. в разделе адаптация и экспорт моделей.

Интерфейс модели

Модели timm имеют общий интерфейс, выходящий за рамки forward(): разделение на forward_features() / forward_head(), сброс классификатора, промежуточные признаки, изменение размера входных данных и хуки для обучения, такие как градиентный чекпоинтинг и группировка параметров. Это описано в API моделей timm, включая информацию о том, какие семейства моделей поддерживают какие части интерфейса.

Чекпоинты и Hub

• load_checkpoint загружает файл контрольной точки (PyTorch или .safetensors) в модель.
• load_state_dict читает файл контрольной точки и возвращает её state dict.
• resume_checkpoint восстанавливает состояние модели, оптимизатора и масштабирующего коэффициента потерь из контрольной точки train.py и возвращает номер эпохи, с которой следует продолжить обучение.

load_checkpoint и load_state_dict по умолчанию (use_ema) выбирают веса EMA из контрольной точки train.py и удаляют префиксы ключей module. (DDP) и _orig_mod. (torch.compile).

• save_for_hf записывает веса и config.json (архитектура, pretrained cfg, model_args, label_names) в папку.
• push_to_hf_hub загружает веса и конфигурацию в репозиторий Hub, создавая сгенерированную карточку модели, если её нет в репозитории.

Затем модель создаётся с помощью hf-hub: или local-dir:, см. раздел о шаринге и загрузке моделей.

import timm
from timm.models import save_for_hf

model = timm.create_model('resnet18', num_classes=4)
# ... fine-tune
save_for_hf(model, 'my_resnet18', model_config=dict(label_names=['a', 'b', 'c', 'd']), safe_serialization=True)
model = timm.create_model('local-dir:my_resnet18', pretrained=True)  # 4 classes, weights and labels restored
print(model.pretrained_cfg['label_names'])

Пользовательские модели

Функция-входная точка модели, декорированная с помощью register_model, добавляется в реестр, поэтому её можно создать с помощью create_model и найти с помощью list_models. Входные точки в timm строят свою модель с помощью build_model_with_cfg, который обрабатывает конфигурацию предобученной модели, загрузку предобученных весов и features_only.

API

Создание модели

>>> from timm import create_model

>>> # Create a MobileNetV3-Large model with no pretrained weights.
>>> model = create_model('mobilenetv3_large_100')

>>> # Create a MobileNetV3-Large model with pretrained weights.
>>> model = create_model('mobilenetv3_large_100', pretrained=True)
>>> model.num_classes
1000

>>> # Create a MobileNetV3-Large model with pretrained weights and a new head with 10 classes.
>>> model = create_model('mobilenetv3_large_100', pretrained=True, num_classes=10)
>>> model.num_classes
10

>>> # Create a Dinov2 small model with pretrained weights and save weights in a custom directory.
>>> model = create_model('vit_small_patch14_dinov2.lvd142m', pretrained=True, cache_dir="/data/my-models")
>>> # Data will be stored at */data/my-models/models--timm--vit_small_patch14_dinov2.lvd142m/*

Поиск моделей

Конфигурация предобученной модели

Текущий вход модели и экспорт

Чекпоинты и Hub

Пользовательские модели