Dear community,

We are excited to announce the release of a new and upgraded deep learning pre-trained models library, MMPreTrain. We have integrated the original MMClassification, image classification algorithm library, and MMSelfSup, self-supervised learning algorithm to launch the deep learning pre-training algorithm library MMPreTrain.

🤔 Compatibility with MMClassification

Fully compatible with MMClassification's directory structure, supported algorithms as well as usage. All the code and projects, which are based on the ordinary mmcls, can be migrated by simply changing the library name.

For example:

1. Importing Components

2. Launch a train/test experiment

For more details about migrating from 0.x to pretrain, you can refer to the migration doc

👍 Major Upgrades

With the release of mmpretrain, we have made several major upgrades to our library.

1. Integrate Self-supervised Algorithms

we have integrated the self-supervised task, which enables users to easily get pre-trained models for various tasks. Users could find that in our directory mmpretrain/models, where a new folder selfsup was made, which support 18 recent self-supervised learning algorithms.

2. Provide convenient higher-level APIs

Secondly, we have provided a more convenient higher-level API, making it easier for users to interact with our library.

1. list_models

list_models supports fuzzy matching, you can use * to match any character.

>>> from mmpretrain import list_models
>>> list_models("*clip-openai")
['vit-base-p16_clip-openai-in12k-pre_3rdparty_in1k',
 'vit-base-p16_clip-openai-in12k-pre_3rdparty_in1k-384px',
 'vit-base-p16_clip-openai-pre_3rdparty_in1k',
 'vit-base-p16_clip-openai-pre_3rdparty_in1k-384px',
 'vit-base-p32_clip-openai-in12k-pre_3rdparty_in1k-384px',
 'vit-base-p32_clip-openai-pre_3rdparty_in1k']

>>> list_models("*convnext-b*21k")
['convnext-base_3rdparty_in21k',
 'convnext-base_in21k-pre-3rdparty_in1k-384px',
 'convnext-base_in21k-pre_3rdparty_in1k']

2. get_model

get_model can get the model from model names

>>> from mmpretrain import get_model
>>> init_model = get_model("convnext-base_in21k-pre_3rdparty_in1k")
>>> pretrained_model = get_model("convnext-base_in21k-pre_3rdparty_in1k", pretrained=True)
# Do the froward
>>> import torch
>>> x = torch.rand((1, 3, 224, 224))
>>> y = pretrained_model(x)
>>> print(type(y), y.shape)
<class 'torch.Tensor'> torch.Size([1, 1000])

3. ImageClassificationInferencer

To use the ImageClassificationInferencer

>>> from mmpretrain import ImageClassificationInferencer
>>> inferencer = ImageClassificationInferencer('resnet50_8xb32_in1k')
>>> results = inferencer('demo/demo.JPEG')
>>> print(results[0]['pred_class'])
sea snake

To inference multiple images by batch on CUDA

>>> from mmpretrain import ImageClassificationInferencer
>>> inferencer = ImageClassificationInferencer('resnet50_8xb32_in1k', device='cuda')
>>> imgs = ['demo/demo.JPEG'] * 100
>>> results = inferencer(imgs, batch_size=16)
Inference ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 24.5 it/s
>>> print(results[99]['pred_class'])
sea snake

4. FeatureExtractor

Compared with model.extract_feat, it's used to extract features from the image files directly, instead of a batch of tensors.

>>> from mmpretrain import FeatureExtractor, get_model
>>> model = get_model('resnet50_8xb32_in1k', backbone=dict(out_indices=(0, 1, 2, 3)))
>>> extractor = FeatureExtractor(model)
>>> features = extractor('demo/bird.JPEG')[0]
>>> for feature in features:
...     print(feature.shape)
torch.Size([256])
torch.Size([512])
torch.Size([1024])
torch.Size([2048])

3. Based on the new training engine MMEngine

Based on MMEngine can support more aspects of upstream chip, training framework updates, and also more aspects of downstream calls to mmpretrain pre-trained models.

1. Support Torch2.0 to accelerate your training

We have fully supported torch2.0, ensuring that our library is compatible with the latest version of PyTorch.
Add the following to your config. You can also refer to MMEngine DOC for help.

compile=True

This is the speed boosting effect

2. Powerful Visualizer

To visualize the image classification result.

from mmpretrain.visualization import UniversalVisualizer
from mmpretrain.structures import DataSample
import mmcv

visualizer = UniversalVisualizer()
image = mmcv.imread('demo/bird.JPEG')[..., ::-1]  # The visualization methods accept RGB image.
data_sample = DataSample().set_gt_label(1).set_pred_label(2).set_pred_score([0., 0.8, 0.2])
visualizer.visualize_cls(image, data_sample, show=True)

For more detail, You can refer to this PR.

↪ Feedbacks

We would like to invite the community to try it out and provide valuable feedback or suggestions. We are committed to improving our library and hope that you will join us on this journey.

The MMPreTrain team