add tests part

tests/test_models/test_heads.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+import os
+import random
+import tempfile
+from unittest import TestCase
+
+import numpy as np
+import torch
+from mmengine import is_seq_of
+
+from mmpretrain.registry import MODELS
+from mmpretrain.structures import DataSample, MultiTaskDataSample
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+class TestClsHead(TestCase):
+    DEFAULT_ARGS = dict(type='ClsHead')
+    FAKE_FEATS = (torch.rand(4, 10), )
+
+    def test_pre_logits(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # return the last item
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        pre_logits = head.pre_logits(feats)
+        self.assertIs(pre_logits, feats[-1])
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # return the last item (same as pre_logits)
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertIs(outs, feats[-1])
+
+    def test_loss(self):
+        feats = self.FAKE_FEATS
+        data_samples = [DataSample().set_gt_label(1) for _ in range(4)]
+
+        # with cal_acc = False
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # with cal_acc = True
+        cfg = {**self.DEFAULT_ARGS, 'topk': (1, 2), 'cal_acc': True}
+        head = MODELS.build(cfg)
+
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(losses.keys(),
+                         {'loss', 'accuracy_top-1', 'accuracy_top-2'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # test assertion when cal_acc but data is batch agumented.
+        data_samples = [
+            sample.set_gt_score(torch.rand(10)) for sample in data_samples
+        ]
+        cfg = {
+            **self.DEFAULT_ARGS, 'cal_acc': True,
+            'loss': dict(type='CrossEntropyLoss', use_soft=True)
+        }
+        head = MODELS.build(cfg)
+        with self.assertRaisesRegex(AssertionError, 'batch augmentation'):
+            head.loss(feats, data_samples)
+
+    def test_predict(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = [DataSample().set_gt_label(1) for _ in range(4)]
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # with without data_samples
+        predictions = head.predict(feats)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for pred in predictions:
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+        # with with data_samples
+        predictions = head.predict(feats, data_samples)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for sample, pred in zip(data_samples, predictions):
+            self.assertIs(sample, pred)
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+
+class TestLinearClsHead(TestCase):
+    DEFAULT_ARGS = dict(type='LinearClsHead', in_channels=10, num_classes=5)
+    FAKE_FEATS = (torch.rand(4, 10), )
+
+    def test_initialize(self):
+        with self.assertRaisesRegex(ValueError, 'num_classes=-5 must be'):
+            MODELS.build({**self.DEFAULT_ARGS, 'num_classes': -5})
+
+    def test_pre_logits(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # return the last item
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        pre_logits = head.pre_logits(feats)
+        self.assertIs(pre_logits, feats[-1])
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+
+class TestVisionTransformerClsHead(TestCase):
+    DEFAULT_ARGS = dict(
+        type='VisionTransformerClsHead', in_channels=10, num_classes=5)
+    fake_feats = ([torch.rand(4, 7, 7, 16), torch.rand(4, 10)], )
+
+    def test_initialize(self):
+        with self.assertRaisesRegex(ValueError, 'num_classes=-5 must be'):
+            MODELS.build({**self.DEFAULT_ARGS, 'num_classes': -5})
+
+        # test vit head default
+        head = MODELS.build(self.DEFAULT_ARGS)
+        assert not hasattr(head.layers, 'pre_logits')
+        assert not hasattr(head.layers, 'act')
+
+        # test vit head hidden_dim
+        head = MODELS.build({**self.DEFAULT_ARGS, 'hidden_dim': 30})
+        assert hasattr(head.layers, 'pre_logits')
+        assert hasattr(head.layers, 'act')
+
+        # test vit head init_weights
+        head = MODELS.build(self.DEFAULT_ARGS)
+        head.init_weights()
+
+        # test vit head init_weights with hidden_dim
+        head = MODELS.build({**self.DEFAULT_ARGS, 'hidden_dim': 30})
+        head.init_weights()
+        assert abs(head.layers.pre_logits.weight).sum() > 0
+
+    def test_pre_logits(self):
+        # test default
+        head = MODELS.build(self.DEFAULT_ARGS)
+        pre_logits = head.pre_logits(self.fake_feats)
+        self.assertIs(pre_logits, self.fake_feats[-1][1])
+
+        # test hidden_dim
+        head = MODELS.build({**self.DEFAULT_ARGS, 'hidden_dim': 30})
+        pre_logits = head.pre_logits(self.fake_feats)
+        self.assertEqual(pre_logits.shape, (4, 30))
+
+    def test_forward(self):
+        # test default
+        head = MODELS.build(self.DEFAULT_ARGS)
+        outs = head(self.fake_feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+        # test hidden_dim
+        head = MODELS.build({**self.DEFAULT_ARGS, 'hidden_dim': 30})
+        outs = head(self.fake_feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+
+class TestDeiTClsHead(TestVisionTransformerClsHead):
+    DEFAULT_ARGS = dict(type='DeiTClsHead', in_channels=10, num_classes=5)
+    fake_feats = ([
+        torch.rand(4, 7, 7, 16),
+        torch.rand(4, 10),
+        torch.rand(4, 10)
+    ], )
+
+    def test_pre_logits(self):
+        # test default
+        head = MODELS.build(self.DEFAULT_ARGS)
+        cls_token, dist_token = head.pre_logits(self.fake_feats)
+        self.assertIs(cls_token, self.fake_feats[-1][1])
+        self.assertIs(dist_token, self.fake_feats[-1][2])
+
+        # test hidden_dim
+        head = MODELS.build({**self.DEFAULT_ARGS, 'hidden_dim': 30})
+        cls_token, dist_token = head.pre_logits(self.fake_feats)
+        self.assertEqual(cls_token.shape, (4, 30))
+        self.assertEqual(dist_token.shape, (4, 30))
+
+
+class TestConformerHead(TestCase):
+    DEFAULT_ARGS = dict(
+        type='ConformerHead', in_channels=[64, 96], num_classes=5)
+    fake_feats = ([torch.rand(4, 64), torch.rand(4, 96)], )
+
+    def test_initialize(self):
+        with self.assertRaisesRegex(ValueError, 'num_classes=-5 must be'):
+            MODELS.build({**self.DEFAULT_ARGS, 'num_classes': -5})
+
+        # test default
+        head = MODELS.build(self.DEFAULT_ARGS)
+        assert hasattr(head, 'conv_cls_head')
+        assert hasattr(head, 'trans_cls_head')
+
+        # test init_weights
+        head = MODELS.build(self.DEFAULT_ARGS)
+        head.init_weights()
+        assert abs(head.conv_cls_head.weight).sum() > 0
+        assert abs(head.trans_cls_head.weight).sum() > 0
+
+    def test_pre_logits(self):
+        # test default
+        head = MODELS.build(self.DEFAULT_ARGS)
+        pre_logits = head.pre_logits(self.fake_feats)
+        self.assertIs(pre_logits, self.fake_feats[-1])
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+        outs = head(self.fake_feats)
+        self.assertEqual(outs[0].shape, (4, 5))
+        self.assertEqual(outs[1].shape, (4, 5))
+
+    def test_loss(self):
+        data_samples = [DataSample().set_gt_label(1) for _ in range(4)]
+
+        # with cal_acc = False
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        losses = head.loss(self.fake_feats, data_samples)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # with cal_acc = True
+        cfg = {**self.DEFAULT_ARGS, 'topk': (1, 2), 'cal_acc': True}
+        head = MODELS.build(cfg)
+
+        losses = head.loss(self.fake_feats, data_samples)
+        self.assertEqual(losses.keys(),
+                         {'loss', 'accuracy_top-1', 'accuracy_top-2'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # test assertion when cal_acc but data is batch agumented.
+        data_samples = [
+            sample.set_gt_score(torch.rand(5)) for sample in data_samples
+        ]
+        cfg = {
+            **self.DEFAULT_ARGS, 'cal_acc': True,
+            'loss': dict(type='CrossEntropyLoss', use_soft=True)
+        }
+        head = MODELS.build(cfg)
+        with self.assertRaisesRegex(AssertionError, 'batch augmentation'):
+            head.loss(self.fake_feats, data_samples)
+
+    def test_predict(self):
+        data_samples = [DataSample().set_gt_label(1) for _ in range(4)]
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # with without data_samples
+        predictions = head.predict(self.fake_feats)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for pred in predictions:
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+        # with with data_samples
+        predictions = head.predict(self.fake_feats, data_samples)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for sample, pred in zip(data_samples, predictions):
+            self.assertIs(sample, pred)
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+
+class TestStackedLinearClsHead(TestCase):
+    DEFAULT_ARGS = dict(
+        type='StackedLinearClsHead', in_channels=10, num_classes=5)
+    fake_feats = (torch.rand(4, 10), )
+
+    def test_initialize(self):
+        with self.assertRaisesRegex(ValueError, 'num_classes=-5 must be'):
+            MODELS.build({
+                **self.DEFAULT_ARGS, 'num_classes': -5,
+                'mid_channels': 10
+            })
+
+        # test mid_channels
+        with self.assertRaisesRegex(AssertionError, 'should be a sequence'):
+            MODELS.build({**self.DEFAULT_ARGS, 'mid_channels': 10})
+
+        # test default
+        head = MODELS.build({**self.DEFAULT_ARGS, 'mid_channels': [20]})
+        assert len(head.layers) == 2
+        head.init_weights()
+
+    def test_pre_logits(self):
+        # test default
+        head = MODELS.build({**self.DEFAULT_ARGS, 'mid_channels': [20, 30]})
+        pre_logits = head.pre_logits(self.fake_feats)
+        self.assertEqual(pre_logits.shape, (4, 30))
+
+    def test_forward(self):
+        # test default
+        head = MODELS.build({**self.DEFAULT_ARGS, 'mid_channels': [20, 30]})
+        outs = head(self.fake_feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+        head = MODELS.build({
+            **self.DEFAULT_ARGS, 'mid_channels': [8, 10],
+            'dropout_rate': 0.2,
+            'norm_cfg': dict(type='BN1d'),
+            'act_cfg': dict(type='HSwish')
+        })
+        outs = head(self.fake_feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+
+class TestMultiLabelClsHead(TestCase):
+    DEFAULT_ARGS = dict(type='MultiLabelClsHead')
+
+    def test_pre_logits(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # return the last item
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        pre_logits = head.pre_logits(feats)
+        self.assertIs(pre_logits, feats[-1])
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # return the last item (same as pre_logits)
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertIs(outs, feats[-1])
+
+    def test_loss(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = [DataSample().set_gt_label([0, 3]) for _ in range(4)]
+
+        # Test with thr and topk are all None
+        head = MODELS.build(self.DEFAULT_ARGS)
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(head.thr, 0.5)
+        self.assertEqual(head.topk, None)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # Test with topk
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        cfg['topk'] = 2
+        head = MODELS.build(cfg)
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(head.thr, None, cfg)
+        self.assertEqual(head.topk, 2)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # Test with thr
+        setup_seed(0)
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        cfg['thr'] = 0.1
+        head = MODELS.build(cfg)
+        thr_losses = head.loss(feats, data_samples)
+        self.assertEqual(head.thr, 0.1)
+        self.assertEqual(head.topk, None)
+        self.assertEqual(thr_losses.keys(), {'loss'})
+        self.assertGreater(thr_losses['loss'].item(), 0)
+
+        # Test with thr and topk are all not None
+        setup_seed(0)
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        cfg['thr'] = 0.1
+        cfg['topk'] = 2
+        head = MODELS.build(cfg)
+        thr_topk_losses = head.loss(feats, data_samples)
+        self.assertEqual(head.thr, 0.1)
+        self.assertEqual(head.topk, 2)
+        self.assertEqual(thr_topk_losses.keys(), {'loss'})
+        self.assertGreater(thr_topk_losses['loss'].item(), 0)
+
+        # Test with gt_lable with score
+        data_samples = [
+            DataSample().set_gt_score(torch.rand((10, ))) for _ in range(4)
+        ]
+
+        head = MODELS.build(self.DEFAULT_ARGS)
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(head.thr, 0.5)
+        self.assertEqual(head.topk, None)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+    def test_predict(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = [DataSample().set_gt_label([1, 2]) for _ in range(4)]
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # with without data_samples
+        predictions = head.predict(feats)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for pred in predictions:
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+        # with with data_samples
+        predictions = head.predict(feats, data_samples)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for sample, pred in zip(data_samples, predictions):
+            self.assertIs(sample, pred)
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+        # Test with topk
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        cfg['topk'] = 2
+        head = MODELS.build(cfg)
+        predictions = head.predict(feats, data_samples)
+        self.assertEqual(head.thr, None)
+        self.assertTrue(is_seq_of(predictions, DataSample))
+        for sample, pred in zip(data_samples, predictions):
+            self.assertIs(sample, pred)
+            self.assertIn('pred_label', pred)
+            self.assertIn('pred_score', pred)
+
+
+class EfficientFormerClsHead(TestClsHead):
+    DEFAULT_ARGS = dict(
+        type='EfficientFormerClsHead',
+        in_channels=10,
+        num_classes=10,
+        distillation=False)
+    FAKE_FEATS = (torch.rand(4, 10), )
+
+    def test_forward(self):
+        # test with distillation head
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        cfg['distillation'] = True
+        head = MODELS.build(cfg)
+        self.assertTrue(hasattr(head, 'dist_head'))
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertEqual(outs.shape, (4, 10))
+
+        # test without distillation head
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        head = MODELS.build(cfg)
+        self.assertFalse(hasattr(head, 'dist_head'))
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertEqual(outs.shape, (4, 10))
+
+    def test_loss(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = [DataSample().set_gt_label(1) for _ in range(4)]
+
+        # test with distillation head
+        cfg = copy.deepcopy(self.DEFAULT_ARGS)
+        cfg['distillation'] = True
+        head = MODELS.build(cfg)
+        with self.assertRaisesRegex(NotImplementedError, 'MMPretrain '):
+            head.loss(feats, data_samples)
+
+        # test without distillation head
+        super().test_loss()
+
+
+class TestMultiLabelLinearClsHead(TestMultiLabelClsHead):
+    DEFAULT_ARGS = dict(
+        type='MultiLabelLinearClsHead', num_classes=10, in_channels=10)
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+        self.assertTrue(hasattr(head, 'fc'))
+        self.assertTrue(isinstance(head.fc, torch.nn.Linear))
+
+        # return the last item (same as pre_logits)
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        head(feats)
+
+
+class TestMultiTaskHead(TestCase):
+    DEFAULT_ARGS = dict(
+        type='MultiTaskHead',  # <- Head config, depends on #675
+        task_heads={
+            'task0': dict(type='LinearClsHead', num_classes=3),
+            'task1': dict(type='LinearClsHead', num_classes=6),
+        },
+        in_channels=10,
+        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
+    )
+
+    DEFAULT_ARGS2 = dict(
+        type='MultiTaskHead',  # <- Head config, depends on #675
+        task_heads={
+            'task0':
+            dict(
+                type='MultiTaskHead',
+                task_heads={
+                    'task00': dict(type='LinearClsHead', num_classes=3),
+                    'task01': dict(type='LinearClsHead', num_classes=6),
+                }),
+            'task1':
+            dict(type='LinearClsHead', num_classes=6)
+        },
+        in_channels=10,
+        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
+    )
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+        # return the last item (same as pre_logits)
+        feats = (torch.rand(4, 10), )
+        outs = head(feats)
+        self.assertEqual(outs['task0'].shape, (4, 3))
+        self.assertEqual(outs['task1'].shape, (4, 6))
+        self.assertTrue(isinstance(outs, dict))
+
+    def test_loss(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = []
+
+        for _ in range(4):
+            data_sample = MultiTaskDataSample()
+            for task_name in self.DEFAULT_ARGS['task_heads']:
+                task_sample = DataSample().set_gt_label(1)
+                data_sample.set_field(task_sample, task_name)
+            data_samples.append(data_sample)
+        # with cal_acc = False
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(
+            losses.keys(),
+            {'task0_loss', 'task0_mask_size', 'task1_loss', 'task1_mask_size'})
+        self.assertGreater(losses['task0_loss'].item(), 0)
+        self.assertGreater(losses['task1_loss'].item(), 0)
+
+    def test_predict(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = []
+
+        for _ in range(4):
+            data_sample = MultiTaskDataSample()
+            for task_name in self.DEFAULT_ARGS['task_heads']:
+                task_sample = DataSample().set_gt_label(1)
+                data_sample.set_field(task_sample, task_name)
+            data_samples.append(data_sample)
+        head = MODELS.build(self.DEFAULT_ARGS)
+        # without data_samples
+        predictions = head.predict(feats)
+        self.assertTrue(is_seq_of(predictions, MultiTaskDataSample))
+        for pred in predictions:
+            self.assertIn('task0', pred)
+        task0_sample = predictions[0].task0
+        self.assertTrue(type(task0_sample.pred_score), 'torch.tensor')
+
+        # with with data_samples
+        predictions = head.predict(feats, data_samples)
+        self.assertTrue(is_seq_of(predictions, MultiTaskDataSample))
+        for sample, pred in zip(data_samples, predictions):
+            self.assertIs(sample, pred)
+            self.assertIn('task0', pred)
+
+        # with data samples and nested
+        head_nested = MODELS.build(self.DEFAULT_ARGS2)
+        # adding a None data sample at the beginning
+        data_samples_nested = [None]
+        for _ in range(3):
+            data_sample_nested = MultiTaskDataSample()
+            data_sample_nested0 = MultiTaskDataSample()
+            data_sample_nested0.set_field(DataSample().set_gt_label(1),
+                                          'task00')
+            data_sample_nested0.set_field(DataSample().set_gt_label(1),
+                                          'task01')
+            data_sample_nested.set_field(data_sample_nested0, 'task0')
+            data_sample_nested.set_field(DataSample().set_gt_label(1), 'task1')
+            data_samples_nested.append(data_sample_nested)
+
+        predictions = head_nested.predict(feats, data_samples_nested)
+        self.assertTrue(is_seq_of(predictions, MultiTaskDataSample))
+        for i in range(3):
+            sample = data_samples_nested[i + 1]
+            pred = predictions[i + 1]
+            self.assertIn('task0', pred)
+            self.assertIn('task1', pred)
+            self.assertIn('task01', pred.get('task0'))
+            self.assertEqual(
+                sample.get('task0').get('task01').gt_label.numpy()[0], 1)
+
+    def test_loss_empty_data_sample(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = []
+
+        for _ in range(4):
+            data_sample = MultiTaskDataSample()
+            data_samples.append(data_sample)
+        # with cal_acc = False
+        head = MODELS.build(self.DEFAULT_ARGS)
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(
+            losses.keys(),
+            {'task0_loss', 'task0_mask_size', 'task1_loss', 'task1_mask_size'})
+        self.assertEqual(losses['task0_loss'].item(), 0)
+        self.assertEqual(losses['task1_loss'].item(), 0)
+
+    def test_nested_multi_task_loss(self):
+
+        head = MODELS.build(self.DEFAULT_ARGS2)
+        # return the last item (same as pre_logits)
+        feats = (torch.rand(4, 10), )
+        outs = head(feats)
+        self.assertEqual(outs['task0']['task01'].shape, (4, 6))
+        self.assertTrue(isinstance(outs, dict))
+        self.assertTrue(isinstance(outs['task0'], dict))
+
+    def test_nested_invalid_sample(self):
+        feats = (torch.rand(4, 10), )
+        gt_label = {'task0': 1, 'task1': 1}
+        head = MODELS.build(self.DEFAULT_ARGS2)
+        data_sample = MultiTaskDataSample()
+        for task_name in gt_label:
+            task_sample = DataSample().set_gt_label(gt_label[task_name])
+            data_sample.set_field(task_sample, task_name)
+        with self.assertRaises(Exception):
+            head.loss(feats, data_sample)
+
+    def test_nested_invalid_sample2(self):
+        feats = (torch.rand(4, 10), )
+        gt_label = {'task0': {'task00': 1, 'task01': 1}, 'task1': 1}
+        head = MODELS.build(self.DEFAULT_ARGS)
+        data_sample = MultiTaskDataSample()
+        task_sample = DataSample().set_gt_label(gt_label['task1'])
+        data_sample.set_field(task_sample, 'task1')
+        data_sample.set_field(MultiTaskDataSample(), 'task0')
+        for task_name in gt_label['task0']:
+            task_sample = DataSample().set_gt_label(
+                gt_label['task0'][task_name])
+            data_sample.task0.set_field(task_sample, task_name)
+        with self.assertRaises(Exception):
+            head.loss(feats, data_sample)
+
+
+class TestArcFaceClsHead(TestCase):
+    DEFAULT_ARGS = dict(type='ArcFaceClsHead', in_channels=10, num_classes=5)
+
+    def test_initialize(self):
+        with self.assertRaises(AssertionError):
+            MODELS.build({**self.DEFAULT_ARGS, 'num_classes': -5})
+
+        with self.assertRaises(AssertionError):
+            MODELS.build({**self.DEFAULT_ARGS, 'num_subcenters': 0})
+
+        # Test margins
+        with self.assertRaises(AssertionError):
+            MODELS.build({**self.DEFAULT_ARGS, 'margins': dict()})
+
+        with self.assertRaises(AssertionError):
+            MODELS.build({**self.DEFAULT_ARGS, 'margins': [0.1] * 4})
+
+        with self.assertRaises(AssertionError):
+            MODELS.build({**self.DEFAULT_ARGS, 'margins': [0.1] * 4 + ['0.1']})
+
+        arcface = MODELS.build(self.DEFAULT_ARGS)
+        torch.allclose(arcface.margins, torch.tensor([0.5] * 5))
+
+        arcface = MODELS.build({**self.DEFAULT_ARGS, 'margins': [0.1] * 5})
+        torch.allclose(arcface.margins, torch.tensor([0.1] * 5))
+
+        margins = [0.1, 0.2, 0.3, 0.4, 5]
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            tmp_path = os.path.join(tmpdirname, 'margins.txt')
+            with open(tmp_path, 'w') as tmp_file:
+                for m in margins:
+                    tmp_file.write(f'{m}\n')
+            arcface = MODELS.build({**self.DEFAULT_ARGS, 'margins': tmp_path})
+            torch.allclose(arcface.margins, torch.tensor(margins))
+
+    def test_pre_logits(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+
+        # return the last item
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        pre_logits = head.pre_logits(feats)
+        self.assertIs(pre_logits, feats[-1])
+
+        # Test with SubCenterArcFace
+        head = MODELS.build({**self.DEFAULT_ARGS, 'num_subcenters': 3})
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        pre_logits = head.pre_logits(feats)
+        self.assertIs(pre_logits, feats[-1])
+
+    def test_forward(self):
+        head = MODELS.build(self.DEFAULT_ARGS)
+        # target is not None
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        target = torch.zeros(4).long()
+        outs = head(feats, target)
+        self.assertEqual(outs.shape, (4, 5))
+
+        # target is None
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+        # Test with SubCenterArcFace
+        head = MODELS.build({**self.DEFAULT_ARGS, 'num_subcenters': 3})
+        # target is not None
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        target = torch.zeros(4)
+        outs = head(feats, target)
+        self.assertEqual(outs.shape, (4, 5))
+
+        # target is None
+        feats = (torch.rand(4, 10), torch.rand(4, 10))
+        outs = head(feats)
+        self.assertEqual(outs.shape, (4, 5))
+
+    def test_loss(self):
+        feats = (torch.rand(4, 10), )
+        data_samples = [DataSample().set_gt_label(1) for _ in range(4)]
+
+        # test loss with used='before'
+        head = MODELS.build(self.DEFAULT_ARGS)
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # Test with SubCenterArcFace
+        head = MODELS.build({**self.DEFAULT_ARGS, 'num_subcenters': 3})
+        # test loss with used='before'
+        losses = head.loss(feats, data_samples)
+        self.assertEqual(losses.keys(), {'loss'})
+        self.assertGreater(losses['loss'].item(), 0)

tests/test_models/test_losses.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+import torch
+
+from mmpretrain.models import build_loss
+
+
+def test_asymmetric_loss():
+    # test asymmetric_loss
+    cls_score = torch.Tensor([[5, -5, 0], [5, -5, 0]])
+    label = torch.Tensor([[1, 0, 1], [0, 1, 0]])
+    weight = torch.tensor([0.5, 0.5])
+
+    loss_cfg = dict(
+        type='AsymmetricLoss',
+        gamma_pos=1.0,
+        gamma_neg=4.0,
+        clip=0.05,
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(3.80845 / 3))
+
+    # test asymmetric_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(3.80845 / 6))
+
+    # test asymmetric_loss without clip
+    loss_cfg = dict(
+        type='AsymmetricLoss',
+        gamma_pos=1.0,
+        gamma_neg=4.0,
+        clip=None,
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(5.1186 / 3))
+
+    # test asymmetric_loss with softmax for single label task
+    cls_score = torch.Tensor([[5, -5, 0], [5, -5, 0]])
+    label = torch.Tensor([0, 1])
+    weight = torch.tensor([0.5, 0.5])
+    loss_cfg = dict(
+        type='AsymmetricLoss',
+        gamma_pos=0.0,
+        gamma_neg=0.0,
+        clip=None,
+        reduction='mean',
+        loss_weight=1.0,
+        use_sigmoid=False,
+        eps=1e-8)
+    loss = build_loss(loss_cfg)
+    # test asymmetric_loss for single label task without weight
+    assert torch.allclose(loss(cls_score, label), torch.tensor(2.5045))
+    # test asymmetric_loss for single label task with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(2.5045 * 0.5))
+
+    # test soft asymmetric_loss with softmax
+    cls_score = torch.Tensor([[5, -5, 0], [5, -5, 0]])
+    label = torch.Tensor([[1, 0, 0], [0, 1, 0]])
+    weight = torch.tensor([0.5, 0.5])
+    loss_cfg = dict(
+        type='AsymmetricLoss',
+        gamma_pos=0.0,
+        gamma_neg=0.0,
+        clip=None,
+        reduction='mean',
+        loss_weight=1.0,
+        use_sigmoid=False,
+        eps=1e-8)
+    loss = build_loss(loss_cfg)
+    # test soft asymmetric_loss with softmax without weight
+    assert torch.allclose(loss(cls_score, label), torch.tensor(2.5045))
+    # test soft asymmetric_loss with softmax with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(2.5045 * 0.5))
+
+
+def test_cross_entropy_loss():
+    with pytest.raises(AssertionError):
+        # use_sigmoid and use_soft could not be set simultaneously
+        loss_cfg = dict(
+            type='CrossEntropyLoss', use_sigmoid=True, use_soft=True)
+        loss = build_loss(loss_cfg)
+
+    # test ce_loss
+    cls_score = torch.Tensor([[-1000, 1000], [100, -100]])
+    label = torch.Tensor([0, 1]).long()
+    class_weight = [0.3, 0.7]  # class 0 : 0.3, class 1 : 0.7
+    weight = torch.tensor([0.6, 0.4])
+
+    # test ce_loss without class weight
+    loss_cfg = dict(type='CrossEntropyLoss', reduction='mean', loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(1100.))
+    # test ce_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(640.))
+
+    # test ce_loss with class weight
+    loss_cfg = dict(
+        type='CrossEntropyLoss',
+        reduction='mean',
+        loss_weight=1.0,
+        class_weight=class_weight)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(370.))
+    # test ce_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(208.))
+
+    # test bce_loss
+    cls_score = torch.Tensor([[-200, 100], [500, -1000], [300, -300]])
+    label = torch.Tensor([[1, 0], [0, 1], [1, 0]])
+    weight = torch.Tensor([0.6, 0.4, 0.5])
+    class_weight = [0.1, 0.9]  # class 0: 0.1, class 1: 0.9
+    pos_weight = [0.1, 0.2]
+
+    # test bce_loss without class weight
+    loss_cfg = dict(
+        type='CrossEntropyLoss',
+        use_sigmoid=True,
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(300.))
+    # test ce_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(130.))
+
+    # test bce_loss with class weight
+    loss_cfg = dict(
+        type='CrossEntropyLoss',
+        use_sigmoid=True,
+        reduction='mean',
+        loss_weight=1.0,
+        class_weight=class_weight)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(176.667))
+    # test bce_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(74.333))
+
+    # test bce loss with pos_weight
+    loss_cfg = dict(
+        type='CrossEntropyLoss',
+        use_sigmoid=True,
+        reduction='mean',
+        loss_weight=1.0,
+        pos_weight=pos_weight)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(136.6667))
+
+    # test soft_ce_loss
+    cls_score = torch.Tensor([[-1000, 1000], [100, -100]])
+    label = torch.Tensor([[1.0, 0.0], [0.0, 1.0]])
+    class_weight = [0.3, 0.7]  # class 0 : 0.3, class 1 : 0.7
+    weight = torch.tensor([0.6, 0.4])
+
+    # test soft_ce_loss without class weight
+    loss_cfg = dict(
+        type='CrossEntropyLoss',
+        use_soft=True,
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(1100.))
+    # test soft_ce_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(640.))
+
+    # test soft_ce_loss with class weight
+    loss_cfg = dict(
+        type='CrossEntropyLoss',
+        use_soft=True,
+        reduction='mean',
+        loss_weight=1.0,
+        class_weight=class_weight)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(370.))
+    # test soft_ce_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(208.))
+
+
+def test_focal_loss():
+    # test focal_loss
+    cls_score = torch.Tensor([[5, -5, 0], [5, -5, 0]])
+    label = torch.Tensor([[1, 0, 1], [0, 1, 0]])
+    weight = torch.tensor([0.5, 0.5])
+
+    loss_cfg = dict(
+        type='FocalLoss',
+        gamma=2.0,
+        alpha=0.25,
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(loss(cls_score, label), torch.tensor(0.8522))
+    # test focal_loss with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(0.8522 / 2))
+    # test focal loss for single label task
+    cls_score = torch.Tensor([[5, -5, 0], [5, -5, 0]])
+    label = torch.Tensor([0, 1])
+    weight = torch.tensor([0.5, 0.5])
+    assert torch.allclose(loss(cls_score, label), torch.tensor(0.86664125))
+    # test focal_loss single label with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(0.86664125 / 2))
+
+
+def test_label_smooth_loss():
+    # test label_smooth_val assertion
+    with pytest.raises(AssertionError):
+        loss_cfg = dict(type='LabelSmoothLoss', label_smooth_val=1.0)
+        build_loss(loss_cfg)
+
+    with pytest.raises(AssertionError):
+        loss_cfg = dict(type='LabelSmoothLoss', label_smooth_val='str')
+        build_loss(loss_cfg)
+
+    # test reduction assertion
+    with pytest.raises(AssertionError):
+        loss_cfg = dict(
+            type='LabelSmoothLoss', label_smooth_val=0.1, reduction='unknown')
+        build_loss(loss_cfg)
+
+    # test mode assertion
+    with pytest.raises(AssertionError):
+        loss_cfg = dict(
+            type='LabelSmoothLoss', label_smooth_val=0.1, mode='unknown')
+        build_loss(loss_cfg)
+
+    # test original mode label smooth loss
+    cls_score = torch.tensor([[1., -1.]])
+    label = torch.tensor([0])
+
+    loss_cfg = dict(
+        type='LabelSmoothLoss',
+        label_smooth_val=0.1,
+        mode='original',
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    correct = 0.2269  # from timm
+    assert loss(cls_score, label) - correct <= 0.0001
+
+    loss_cfg = dict(
+        type='LabelSmoothLoss',
+        label_smooth_val=0.1,
+        mode='original',
+        use_sigmoid=True,
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    correct = 0.3633  # from timm
+    assert loss(cls_score, label) - correct <= 0.0001
+
+    # test classy_vision mode label smooth loss
+    loss_cfg = dict(
+        type='LabelSmoothLoss',
+        label_smooth_val=0.1,
+        mode='classy_vision',
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    correct = 0.2178  # from ClassyVision
+    assert loss(cls_score, label) - correct <= 0.0001
+
+    # test multi_label mode label smooth loss
+    cls_score = torch.tensor([[1., -1., 1]])
+    label = torch.tensor([[1, 0, 1]])
+
+    loss_cfg = dict(
+        type='LabelSmoothLoss',
+        label_smooth_val=0.1,
+        mode='multi_label',
+        reduction='mean',
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    smooth_label = torch.tensor([[0.9, 0.1, 0.9]])
+    correct = torch.binary_cross_entropy_with_logits(cls_score,
+                                                     smooth_label).mean()
+    assert torch.allclose(loss(cls_score, label), correct)
+
+    # test label linear combination smooth loss
+    cls_score = torch.tensor([[1., -1., 0.]])
+    label1 = torch.tensor([[1., 0., 0.]])
+    label2 = torch.tensor([[0., 0., 1.]])
+    label_mix = label1 * 0.6 + label2 * 0.4
+
+    loss_cfg = dict(
+        type='LabelSmoothLoss',
+        label_smooth_val=0.1,
+        mode='original',
+        reduction='mean',
+        num_classes=3,
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    smooth_label1 = loss.original_smooth_label(label1)
+    smooth_label2 = loss.original_smooth_label(label2)
+    label_smooth_mix = smooth_label1 * 0.6 + smooth_label2 * 0.4
+    correct = (-torch.log_softmax(cls_score, -1) * label_smooth_mix).sum()
+
+    assert loss(cls_score, label_mix) - correct <= 0.0001
+
+    # test label smooth loss with weight
+    cls_score = torch.tensor([[1., -1.], [1., -1.]])
+    label = torch.tensor([0, 1])
+    weight = torch.tensor([0.5, 0.5])
+
+    loss_cfg = dict(
+        type='LabelSmoothLoss',
+        reduction='mean',
+        label_smooth_val=0.1,
+        loss_weight=1.0)
+    loss = build_loss(loss_cfg)
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight),
+        loss(cls_score, label) / 2)
+
+
+# migrate from mmdetection with modifications
+def test_seesaw_loss():
+    # only softmax version of Seesaw Loss is implemented
+    with pytest.raises(AssertionError):
+        loss_cfg = dict(type='SeesawLoss', use_sigmoid=True, loss_weight=1.0)
+        build_loss(loss_cfg)
+
+    # test that cls_score.size(-1) == num_classes
+    loss_cls_cfg = dict(
+        type='SeesawLoss', p=0.0, q=0.0, loss_weight=1.0, num_classes=2)
+    loss_cls = build_loss(loss_cls_cfg)
+    # the length of fake_pred should be num_classe = 4
+    with pytest.raises(AssertionError):
+        fake_pred = torch.Tensor([[-100, 100, -100]])
+        fake_label = torch.Tensor([1]).long()
+        loss_cls(fake_pred, fake_label)
+    # the length of fake_pred should be num_classes + 2 = 4
+    with pytest.raises(AssertionError):
+        fake_pred = torch.Tensor([[-100, 100, -100, 100]])
+        fake_label = torch.Tensor([1]).long()
+        loss_cls(fake_pred, fake_label)
+
+    # test the calculation without p and q
+    loss_cls_cfg = dict(
+        type='SeesawLoss', p=0.0, q=0.0, loss_weight=1.0, num_classes=2)
+    loss_cls = build_loss(loss_cls_cfg)
+    fake_pred = torch.Tensor([[-100, 100]])
+    fake_label = torch.Tensor([1]).long()
+    loss = loss_cls(fake_pred, fake_label)
+    assert torch.allclose(loss, torch.tensor(0.))
+
+    # test the calculation with p and without q
+    loss_cls_cfg = dict(
+        type='SeesawLoss', p=1.0, q=0.0, loss_weight=1.0, num_classes=2)
+    loss_cls = build_loss(loss_cls_cfg)
+    fake_pred = torch.Tensor([[-100, 100]])
+    fake_label = torch.Tensor([0]).long()
+    loss_cls.cum_samples[0] = torch.exp(torch.Tensor([20]))
+    loss = loss_cls(fake_pred, fake_label)
+    assert torch.allclose(loss, torch.tensor(180.))
+
+    # test the calculation with q and without p
+    loss_cls_cfg = dict(
+        type='SeesawLoss', p=0.0, q=1.0, loss_weight=1.0, num_classes=2)
+    loss_cls = build_loss(loss_cls_cfg)
+    fake_pred = torch.Tensor([[-100, 100]])
+    fake_label = torch.Tensor([0]).long()
+    loss = loss_cls(fake_pred, fake_label)
+    assert torch.allclose(loss, torch.tensor(200.) + torch.tensor(100.).log())
+
+
+def test_reconstruction_loss():
+
+    # test L2 loss
+    loss_config = dict(type='PixelReconstructionLoss', criterion='L2')
+    loss = build_loss(loss_config)
+
+    fake_pred = torch.rand((2, 196, 768))
+    fake_target = torch.rand((2, 196, 768))
+    fake_mask = torch.ones((2, 196))
+    loss_value = loss(fake_pred, fake_target, fake_mask)
+
+    assert isinstance(loss_value.item(), float)
+
+    # test L1 loss
+    loss_config = dict(
+        type='PixelReconstructionLoss', criterion='L1', channel=3)
+    loss = build_loss(loss_config)
+
+    fake_pred = torch.rand((2, 3, 192, 192))
+    fake_target = torch.rand((2, 3, 192, 192))
+    fake_mask = torch.ones((2, 1, 192, 192))
+    loss_value = loss(fake_pred, fake_target, fake_mask)
+
+    assert isinstance(loss_value.item(), float)
+
+    with pytest.raises(NotImplementedError):
+        loss_config = dict(type='PixelReconstructionLoss', criterion='L3')
+        loss = build_loss(loss_config)

tests/test_models/test_models.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+import mmpretrain.models
+from mmpretrain.apis import ModelHub, get_model
+
+
+@dataclass
+class Cfg:
+    name: str
+    backbone: type
+    num_classes: int = 1000
+    build: bool = True
+    forward: bool = True
+    backward: bool = True
+    extract_feat: bool = True
+    input_shape: tuple = (1, 3, 224, 224)
+
+
+test_list = [
+    Cfg(name='xcit-small-12-p16_3rdparty_in1k',
+        backbone=mmpretrain.models.XCiT),
+    Cfg(name='xcit-nano-12-p8_3rdparty-dist_in1k-384px',
+        backbone=mmpretrain.models.XCiT,
+        input_shape=(1, 3, 384, 384)),
+    Cfg(name='vit-base-p16_sam-pre_3rdparty_sa1b-1024px',
+        backbone=mmpretrain.models.ViTSAM,
+        forward=False,
+        backward=False),
+    Cfg(name='vit-base-p14_dinov2-pre_3rdparty',
+        backbone=mmpretrain.models.VisionTransformer,
+        forward=False,
+        backward=False),
+    Cfg(name='hivit-tiny-p16_16xb64_in1k', backbone=mmpretrain.models.HiViT),
+]
+
+
+@pytest.mark.parametrize('cfg', test_list)
+def test_build(cfg: Cfg):
+    if not cfg.build:
+        return
+
+    model_name = cfg.name
+    ModelHub._register_mmpretrain_models()
+    assert ModelHub.has(model_name)
+
+    model = get_model(model_name)
+    backbone_class = cfg.backbone
+    assert isinstance(model.backbone, backbone_class)
+
+
+@pytest.mark.parametrize('cfg', test_list)
+def test_forward(cfg: Cfg):
+    if not cfg.forward:
+        return
+
+    model = get_model(cfg.name)
+    inputs = torch.rand(*cfg.input_shape)
+    outputs = model(inputs)
+    assert outputs.shape == (1, cfg.num_classes)
+
+
+@pytest.mark.parametrize('cfg', test_list)
+def test_extract_feat(cfg: Cfg):
+    if not cfg.extract_feat:
+        return
+
+    model = get_model(cfg.name)
+    inputs = torch.rand(*cfg.input_shape)
+    feats = model.extract_feat(inputs)
+    assert isinstance(feats, tuple)
+    assert len(feats) == 1
+
+
+@pytest.mark.parametrize('cfg', test_list)
+def test_backward(cfg: Cfg):
+    if not cfg.backward:
+        return
+
+    model = get_model(cfg.name)
+    inputs = torch.rand(*cfg.input_shape)
+    outputs = model(inputs)
+    outputs.mean().backward()
+
+    for n, x in model.named_parameters():
+        assert x.grad is not None, f'No gradient for {n}'
+    num_grad = sum(
+        [x.grad.numel() for x in model.parameters() if x.grad is not None])
+    assert outputs.shape[-1] == cfg.num_classes
+    num_params = sum([x.numel() for x in model.parameters()])
+    assert num_params == num_grad, 'Some parameters are missing gradients'
+    assert not torch.isnan(outputs).any(), 'Output included NaNs'

tests/test_models/test_necks.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+import torch
+
+from mmpretrain.models.necks import (GeneralizedMeanPooling,
+                                     GlobalAveragePooling, HRFuseScales,
+                                     LinearNeck)
+
+
+def test_gap_neck():
+
+    # test 1d gap_neck
+    neck = GlobalAveragePooling(dim=1)
+    # batch_size, num_features, feature_size
+    fake_input = torch.rand(1, 16, 24)
+
+    output = neck(fake_input)
+    # batch_size, num_features
+    assert output.shape == (1, 16)
+
+    # test 1d gap_neck
+    neck = GlobalAveragePooling(dim=2)
+    # batch_size, num_features, feature_size(2)
+    fake_input = torch.rand(1, 16, 24, 24)
+
+    output = neck(fake_input)
+    # batch_size, num_features
+    assert output.shape == (1, 16)
+
+    # test 1d gap_neck
+    neck = GlobalAveragePooling(dim=3)
+    # batch_size, num_features, feature_size(3)
+    fake_input = torch.rand(1, 16, 24, 24, 5)
+
+    output = neck(fake_input)
+    # batch_size, num_features
+    assert output.shape == (1, 16)
+
+    with pytest.raises(AssertionError):
+        # dim must in [1, 2, 3]
+        GlobalAveragePooling(dim='other')
+
+
+def test_gem_neck():
+
+    # test gem_neck
+    neck = GeneralizedMeanPooling()
+
+    # default p is trainable
+    assert neck.p.requires_grad
+
+    # batch_size, num_features, feature_size(2)
+    fake_input = torch.rand(1, 16, 24, 24)
+
+    output = neck(fake_input)
+    # batch_size, num_features
+    assert output.shape == (1, 16)
+
+    # test tuple input gem_neck
+    neck = GeneralizedMeanPooling()
+    # batch_size, num_features, feature_size(2)
+    fake_input = (torch.rand(1, 8, 24, 24), torch.rand(1, 16, 24, 24))
+
+    output = neck(fake_input)
+    # batch_size, num_features
+    assert output[0].shape == (1, 8)
+    assert output[1].shape == (1, 16)
+
+    # test gem_neck with p_trainable=False
+    neck = GeneralizedMeanPooling(p_trainable=False)
+
+    # p is not trainable
+    assert not neck.p.requires_grad
+
+    # batch_size, num_features, feature_size(2)
+    fake_input = torch.rand(1, 16, 24, 24)
+
+    output = neck(fake_input)
+    # batch_size, num_features
+    assert output.shape == (1, 16)
+
+    with pytest.raises(AssertionError):
+        # p must be a value greater then 1
+        GeneralizedMeanPooling(p=0.5)
+
+
+def test_hr_fuse_scales():
+
+    in_channels = (18, 32, 64, 128)
+    neck = HRFuseScales(in_channels=in_channels, out_channels=1024)
+
+    feat_size = 56
+    inputs = []
+    for in_channel in in_channels:
+        input_tensor = torch.rand(3, in_channel, feat_size, feat_size)
+        inputs.append(input_tensor)
+        feat_size = feat_size // 2
+
+    with pytest.raises(AssertionError):
+        neck(inputs)
+
+    outs = neck(tuple(inputs))
+    assert isinstance(outs, tuple)
+    assert len(outs) == 1
+    assert outs[0].shape == (3, 1024, 7, 7)
+
+
+def test_linear_reduction():
+    # test linear_reduction without `act_cfg` and `norm_cfg`
+    neck = LinearNeck(10, 5, 0, None, None)
+    neck.eval()
+    assert isinstance(neck.gap, torch.nn.Identity)
+    assert isinstance(neck.act, torch.nn.Identity)
+    assert isinstance(neck.norm, torch.nn.Identity)
+
+    # batch_size, in_channels, out_channels
+    fake_input = torch.rand(1, 10)
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+
+    # batch_size, in_features, feature_size(2)
+    fake_input = (torch.rand(1, 20), torch.rand(1, 10))
+
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+
+    # batch_size, in_channels, out_channels, gap_dim
+    neck = LinearNeck(10, 5, 1, None, None)
+    fake_input = torch.rand(1, 10, 10)
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+
+    # batch_size, in_channels, out_channels, gap_dim
+    neck = LinearNeck(10, 5, 2, None, None)
+    fake_input = torch.rand(1, 10, 10, 10)
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+
+    # batch_size, in_channels, out_channels, gap_dim
+    neck = LinearNeck(10, 5, 3, None, None)
+    fake_input = torch.rand(1, 10, 10, 10, 10)
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+
+    # batch_size, in_channels, out_channels, gap_dim
+    with pytest.raises(AssertionError):
+        neck = LinearNeck(10, 5, None, None, None)
+
+    # test linear_reduction with `init_cfg`
+    neck = LinearNeck(10, 5, init_cfg=dict(type='Xavier', layer=['Linear']))
+
+    # test linear_reduction with `act_cfg` and `norm_cfg`
+    neck = LinearNeck(
+        10, 5, act_cfg=dict(type='ReLU'), norm_cfg=dict(type='BN1d'))
+    neck.eval()
+
+    assert isinstance(neck.act, torch.nn.ReLU)
+    assert isinstance(neck.norm, torch.nn.BatchNorm1d)
+
+    # batch_size, in_channels, out_channels
+    fake_input = torch.rand(1, 10)
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+    #
+    # # batch_size, in_features, feature_size(2)
+    fake_input = (torch.rand(1, 20), torch.rand(1, 10))
+
+    output = neck(fake_input)
+    # batch_size, out_features
+    assert output[-1].shape == (1, 5)
+
+    with pytest.raises(AssertionError):
+        neck([])

tests/test_models/test_peft/test_lora.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import re
+
+import pytest
+from mmengine.utils import digit_version
+from mmengine.utils.dl_utils import TORCH_VERSION
+
+from mmpretrain.models.peft import LoRAModel
+
+
+@pytest.mark.skipif(
+    digit_version(TORCH_VERSION) < digit_version('1.9.0'),
+    reason='get_submodule requires torch >= 1.9.0')
+def test_lora_backbone():
+    module = dict(
+        type='VisionTransformer',
+        arch='base',
+        img_size=224,
+        patch_size=16,
+        drop_path_rate=0.1,
+        out_type='avg_featmap',
+        final_norm=False)
+
+    lora_cfg = dict(
+        module=module,
+        alpha=1,
+        rank=4,
+        drop_rate=0.1,
+        targets=[
+            dict(type='qkv'),
+            dict(type='.*proj', alpha=2, rank=2, drop_rate=0.2),
+        ])
+
+    lora_model = LoRAModel(**lora_cfg)
+
+    # test replace module
+    for name, module in lora_model.named_modules():
+        if name.endswith('qkv'):
+            assert module.scaling == 0.25
+        if re.fullmatch('.*proj', name):
+            assert module.scaling == 1
+
+    # test freeze module
+    for name, param in lora_model.named_parameters():
+        if 'lora_' in name:
+            assert param.requires_grad
+        else:
+            assert not param.requires_grad
+
+    # test get state dict
+    state_dict = lora_model.state_dict()
+    assert len(state_dict) != 0
+    for name, param in state_dict.items():
+        assert 'lora_' in name
+
+    # test load state dict
+    incompatible_keys = lora_model.load_state_dict(state_dict, strict=True)
+    assert str(incompatible_keys) == '<All keys matched successfully>'
+
+
+@pytest.mark.skipif(
+    digit_version(TORCH_VERSION) < digit_version('1.9.0'),
+    reason='get_submodule requires torch >= 1.9.0')
+def test_lora_model():
+    module = dict(
+        type='MAE',
+        backbone=dict(type='MAEViT', arch='b', patch_size=16, mask_ratio=0.75),
+        neck=dict(
+            type='MAEPretrainDecoder',
+            patch_size=16,
+            in_chans=3,
+            embed_dim=768,
+            decoder_embed_dim=512,
+            decoder_depth=8,
+            decoder_num_heads=16,
+            mlp_ratio=4.,
+        ),
+        head=dict(
+            type='MAEPretrainHead',
+            norm_pix=True,
+            patch_size=16,
+            loss=dict(type='PixelReconstructionLoss', criterion='L2')),
+        init_cfg=[
+            dict(type='Xavier', layer='Linear', distribution='uniform'),
+            dict(type='Constant', layer='LayerNorm', val=1.0, bias=0.0)
+        ])
+
+    lora_cfg = dict(
+        module=module,
+        alpha=1,
+        rank=4,
+        drop_rate=0.1,
+        targets=[
+            dict(type='qkv'),
+            dict(type='.*proj', alpha=2, rank=2, drop_rate=0.2),
+        ])
+
+    lora_model = LoRAModel(**lora_cfg)
+
+    # test replace module
+    for name, module in lora_model.named_modules():
+        if name.endswith('qkv'):
+            assert module.scaling == 0.25
+        if re.fullmatch('.*proj', name):
+            assert module.scaling == 1
+
+    # test freeze module
+    for name, param in lora_model.named_parameters():
+        if 'lora_' in name:
+            assert param.requires_grad
+        else:
+            assert not param.requires_grad
+
+    # test get state dict
+    state_dict = lora_model.state_dict()
+    assert len(state_dict) != 0
+    for name, param in state_dict.items():
+        assert 'lora_' in name
+
+    # test load state dict
+    incompatible_keys = lora_model.load_state_dict(state_dict, strict=True)
+    assert str(incompatible_keys) == '<All keys matched successfully>'

tests/test_models/test_retrievers.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import os
+import tempfile
+from typing import Callable
+from unittest import TestCase
+from unittest.mock import MagicMock
+
+import numpy as np
+import torch
+from mmengine import ConfigDict
+from mmengine.dataset.utils import default_collate
+from torch.utils.data import DataLoader, Dataset
+
+from mmpretrain.datasets.transforms import PackInputs
+from mmpretrain.registry import MODELS
+from mmpretrain.structures import DataSample
+
+
+class ExampleDataset(Dataset):
+
+    def __init__(self):
+        self.metainfo = None
+        self.pipe = PackInputs()
+
+    def __getitem__(self, idx):
+        results = dict(
+            img=np.random.random((64, 64, 3)), meta=dict(sampleidx=idx))
+
+        return self.pipe(results)
+
+    def __len__(self):
+        return 10
+
+
+class TestImageToImageRetriever(TestCase):
+    DEFAULT_ARGS = dict(
+        type='ImageToImageRetriever',
+        image_encoder=[
+            dict(type='ResNet', depth=18, out_indices=(3, )),
+            dict(type='GlobalAveragePooling'),
+        ],
+        head=dict(
+            type='LinearClsHead',
+            num_classes=10,
+            in_channels=512,
+            loss=dict(type='CrossEntropyLoss')),
+        prototype=torch.rand((10, 512)),
+    )
+
+    def test_initialize(self):
+        # test error prototype type
+        cfg = {**self.DEFAULT_ARGS, 'prototype': 5}
+        with self.assertRaises(AssertionError):
+            model = MODELS.build(cfg)
+
+        # test prototype is tensor
+        model = MODELS.build(self.DEFAULT_ARGS)
+        self.assertEqual(type(model.prototype), torch.Tensor)
+        self.assertFalse(model.prototype_inited)
+        self.assertIsInstance(model.similarity_fn, Callable)
+        self.assertEqual(model.topk, -1)
+
+        # test prototype is str
+        cfg = {**self.DEFAULT_ARGS, 'prototype': './proto.pth'}
+        model = MODELS.build(cfg)
+        self.assertEqual(type(model.prototype), str)
+
+        # test prototype is dict
+        lodaer = DataLoader(ExampleDataset())
+        cfg = {**self.DEFAULT_ARGS, 'prototype': lodaer}
+        model = MODELS.build(cfg)
+        self.assertEqual(type(model.prototype), DataLoader)
+
+        # test prototype is dataloader
+        loader_cfg = dict(
+            batch_size=16,
+            num_workers=2,
+            dataset=dict(
+                type='CIFAR100',
+                data_prefix='data/cifar100',
+                test_mode=False,
+                pipeline=[]),
+            sampler=dict(type='DefaultSampler', shuffle=True),
+            persistent_workers=True)
+        cfg = {**self.DEFAULT_ARGS, 'prototype': loader_cfg}
+        model = MODELS.build(cfg)
+        self.assertEqual(type(model.prototype), dict)
+
+        # test similarity function
+        self.assertEqual(model.similarity, 'cosine_similarity')
+
+        def fn(a, b):
+            return a * b
+
+        cfg = {**self.DEFAULT_ARGS, 'similarity_fn': fn}
+        model = MODELS.build(cfg)
+        self.assertEqual(model.similarity, fn)
+        self.assertIsInstance(model.similarity_fn, Callable)
+
+        # test set batch augmentation from train_cfg
+        cfg = {
+            **self.DEFAULT_ARGS, 'train_cfg':
+            dict(augments=dict(
+                type='Mixup',
+                alpha=1.,
+            ))
+        }
+        model = MODELS.build(cfg)
+
+        self.assertIsNotNone(model.data_preprocessor.batch_augments)
+
+        cfg = {**self.DEFAULT_ARGS, 'train_cfg': dict()}
+        model = MODELS.build(cfg)
+        self.assertIsNone(model.data_preprocessor.batch_augments)
+
+    def test_extract_feat(self):
+        inputs = torch.rand(1, 3, 64, 64)
+        cfg = ConfigDict(self.DEFAULT_ARGS)
+        model = MODELS.build(cfg)
+
+        # test extract_feat
+        feats = model.extract_feat(inputs)
+        self.assertEqual(len(feats), 1)
+        self.assertEqual(feats[0].shape, (1, 512))
+
+    def test_loss(self):
+        inputs = torch.rand(1, 3, 64, 64)
+        data_samples = [DataSample().set_gt_label(1)]
+
+        model = MODELS.build(self.DEFAULT_ARGS)
+        losses = model.loss(inputs, data_samples)
+        self.assertGreater(losses['loss'].item(), 0)
+
+    def test_prepare_prototype(self):
+        tmpdir = tempfile.TemporaryDirectory()
+        # tensor
+        cfg = {**self.DEFAULT_ARGS}
+        model = MODELS.build(cfg)
+        model.prepare_prototype()
+        self.assertEqual(type(model.prototype_vecs), torch.Tensor)
+        self.assertEqual(model.prototype_vecs.shape, (10, 512))
+        self.assertTrue(model.prototype_inited)
+
+        # test dump prototype
+        ori_proto_vecs = model.prototype_vecs
+        save_path = os.path.join(tmpdir.name, 'proto.pth')
+        model.dump_prototype(save_path)
+
+        # Check whether the saved feature exists
+        feat = torch.load(save_path)
+        self.assertEqual(feat.shape, (10, 512))
+
+        # str
+        cfg = {**self.DEFAULT_ARGS, 'prototype': save_path}
+        model = MODELS.build(cfg)
+        model.prepare_prototype()
+        self.assertEqual(type(model.prototype_vecs), torch.Tensor)
+        self.assertEqual(model.prototype_vecs.shape, (10, 512))
+        self.assertTrue(model.prototype_inited)
+        torch.allclose(ori_proto_vecs, model.prototype_vecs)
+
+        # dict
+        lodaer = DataLoader(ExampleDataset(), collate_fn=default_collate)
+        cfg = {**self.DEFAULT_ARGS, 'prototype': lodaer}
+        model = MODELS.build(cfg)
+        model.prepare_prototype()
+        self.assertEqual(type(model.prototype_vecs), torch.Tensor)
+        self.assertEqual(model.prototype_vecs.shape, (10, 512))
+        self.assertTrue(model.prototype_inited)
+
+        tmpdir.cleanup()
+
+    def test_predict(self):
+        inputs = torch.rand(1, 3, 64, 64)
+        data_samples = [DataSample().set_gt_label([1, 2, 6])]
+        # default
+        model = MODELS.build(self.DEFAULT_ARGS)
+        predictions = model.predict(inputs)
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))
+
+        predictions = model.predict(inputs, data_samples)
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))
+        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
+        torch.testing.assert_allclose(data_samples[0].pred_score,
+                                      predictions[0].pred_score)
+
+        # k is not -1
+        cfg = {**self.DEFAULT_ARGS, 'topk': 2}
+        model = MODELS.build(cfg)
+
+        predictions = model.predict(inputs)
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))
+
+        predictions = model.predict(inputs, data_samples)
+        assert predictions is data_samples
+        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
+
+    def test_forward(self):
+        inputs = torch.rand(1, 3, 64, 64)
+        data_samples = [DataSample().set_gt_label(1)]
+        model = MODELS.build(self.DEFAULT_ARGS)
+
+        # test pure forward
+        outs = model(inputs)
+        # assert False, type(outs)
+        self.assertIsInstance(outs, tuple)
+        self.assertEqual(len(outs), 1)
+        self.assertIsInstance(outs[0], torch.Tensor)
+
+        # test forward train
+        losses = model(inputs, data_samples, mode='loss')
+        self.assertGreater(losses['loss'].item(), 0)
+
+        # test forward test
+        predictions = model(inputs, mode='predict')
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))
+
+        predictions = model(inputs, data_samples, mode='predict')
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))
+        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
+        torch.testing.assert_allclose(data_samples[0].pred_score,
+                                      predictions[0].pred_score)
+
+        # test forward with invalid mode
+        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
+            model(inputs, mode='unknown')
+
+    def test_train_step(self):
+        cfg = {
+            **self.DEFAULT_ARGS, 'data_preprocessor':
+            dict(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5])
+        }
+        model = MODELS.build(cfg)
+
+        data = {
+            'inputs': torch.randint(0, 256, (1, 3, 64, 64)),
+            'data_samples': [DataSample().set_gt_label(1)]
+        }
+
+        optim_wrapper = MagicMock()
+        log_vars = model.train_step(data, optim_wrapper)
+        self.assertIn('loss', log_vars)
+        optim_wrapper.update_params.assert_called_once()
+
+    def test_val_step(self):
+        cfg = {
+            **self.DEFAULT_ARGS, 'data_preprocessor':
+            dict(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5])
+        }
+        model = MODELS.build(cfg)
+
+        data = {
+            'inputs': torch.randint(0, 256, (1, 3, 64, 64)),
+            'data_samples': [DataSample().set_gt_label(1)]
+        }
+
+        predictions = model.val_step(data)
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))
+
+    def test_test_step(self):
+        cfg = {
+            **self.DEFAULT_ARGS, 'data_preprocessor':
+            dict(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5])
+        }
+        model = MODELS.build(cfg)
+
+        data = {
+            'inputs': torch.randint(0, 256, (1, 3, 64, 64)),
+            'data_samples': [DataSample().set_gt_label(1)]
+        }
+
+        predictions = model.test_step(data)
+        self.assertEqual(predictions[0].pred_score.shape, (10, ))

tests/test_models/test_selfsup/test_barlowtwins.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import BarlowTwins
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_barlowtwins():
+    data_preprocessor = {
+        'mean': (123.675, 116.28, 103.53),
+        'std': (58.395, 57.12, 57.375),
+        'to_rgb': True
+    }
+    backbone = dict(type='ResNet', depth=18, norm_cfg=dict(type='BN'))
+    neck = dict(
+        type='NonLinearNeck',
+        in_channels=512,
+        hid_channels=2,
+        out_channels=2,
+        num_layers=3,
+        with_last_bn=False,
+        with_last_bn_affine=False,
+        with_avg_pool=True,
+        norm_cfg=dict(type='BN1d'))
+    head = dict(
+        type='LatentCrossCorrelationHead',
+        in_channels=2,
+        loss=dict(type='CrossCorrelationLoss'))
+
+    alg = BarlowTwins(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    fake_data = {
+        'inputs':
+        [torch.randn((2, 3, 224, 224)),
+         torch.randn((2, 3, 224, 224))],
+        'data_sample': [DataSample() for _ in range(2)]
+    }
+
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_loss = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_loss['loss'].item(), float)

tests/test_models/test_selfsup/test_beit.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+from unittest import TestCase
+
+import pytest
+import torch
+
+from mmpretrain.models import BEiT, BEiTPretrainViT
+from mmpretrain.structures import DataSample
+
+
+class TestBEiT(TestCase):
+
+    @pytest.mark.skipif(
+        platform.system() == 'Windows', reason='Windows mem limit')
+    def test_beit_pretrain_vit(self):
+        backbone = dict(
+            arch='base',
+            patch_size=16,
+            drop_path_rate=0.1,
+            final_norm=True,
+            layer_scale_init_value=0.1,
+        )
+
+        beit_backbone = BEiTPretrainViT(**backbone)
+        beit_backbone.init_weights()
+
+        fake_inputs = torch.randn((2, 3, 224, 224))
+        fake_mask = torch.zeros((2, 196))
+        fake_mask[:, 75:150] = 1
+
+        # test with mask
+        fake_outputs = beit_backbone(fake_inputs, fake_mask)
+        assert fake_outputs[0].shape == torch.Size([2, 197, 768])
+
+        # test without mask
+        fake_outputs = beit_backbone(fake_inputs, None)
+        assert fake_outputs[0].shape == torch.Size([2, 197, 768])
+
+    @pytest.mark.skipif(
+        platform.system() == 'Windows', reason='Windows mem limit')
+    def test_beitv1(self):
+        data_preprocessor = dict(
+            type='TwoNormDataPreprocessor',
+            mean=[123.675, 116.28, 103.53],
+            std=[58.395, 57.12, 57.375],
+            second_mean=[-31.875, -31.875, -31.875],
+            second_std=[318.75, 318.75, 318.75],
+            to_rgb=True)
+
+        # model settings
+        backbone = dict(
+            type='BEiTPretrainViT',
+            arch='base',
+            patch_size=16,
+            drop_path_rate=0.1,
+            final_norm=True,
+            layer_scale_init_value=0.1)
+        neck = None
+        head = dict(
+            type='BEiTV1Head',
+            embed_dims=768,
+            num_embed=8192,
+            loss=dict(type='CrossEntropyLoss'))
+        target_generator = dict(type='DALL-E')
+
+        # build model
+        model = BEiT(
+            backbone=backbone,
+            neck=neck,
+            head=head,
+            target_generator=target_generator,
+            data_preprocessor=data_preprocessor)
+
+        fake_img = torch.rand((1, 3, 224, 224))
+        fake_target_img = torch.rand((1, 3, 112, 112))
+        fake_mask = torch.zeros((196)).bool()
+        fake_mask[75:150] = 1
+        fake_data_sample = DataSample()
+        fake_data_sample.set_mask(fake_mask)
+        fake_data = {
+            'inputs': [fake_img, fake_target_img],
+            'data_samples': [fake_data_sample]
+        }
+
+        fake_inputs = model.data_preprocessor(fake_data)
+        fake_outputs = model(**fake_inputs, mode='loss')
+        assert isinstance(fake_outputs['loss'].item(), float)
+
+    @pytest.mark.skipif(
+        platform.system() == 'Windows', reason='Windows mem limit')
+    def test_beitv2(self):
+        data_preprocessor = dict(
+            type='TwoNormDataPreprocessor',
+            mean=(123.675, 116.28, 103.53),
+            std=(58.395, 57.12, 57.375),
+            second_mean=(127.5, 127.5, 127.5),
+            second_std=(127.5, 127.5, 127.5),
+            to_rgb=True)
+
+        # model settings
+        vqkd_encoder = dict(
+            arch='base',
+            img_size=224,
+            patch_size=16,
+            in_channels=3,
+            out_indices=-1,
+            drop_rate=0.,
+            drop_path_rate=0.,
+            norm_cfg=dict(type='LN', eps=1e-6),
+            final_norm=True,
+            out_type='featmap',
+            with_cls_token=True,
+            frozen_stages=-1,
+            use_abs_pos_emb=True,
+            use_rel_pos_bias=False,
+            use_shared_rel_pos_bias=False,
+            layer_scale_init_value=0.,
+            interpolate_mode='bicubic',
+            patch_cfg=dict(),
+            layer_cfgs=dict(),
+            init_cfg=None)
+
+        layer_scale_init_value = 0.1
+        drop_path_rate = 0.  # 0. for 300 epochs and 0.1 for 1600 epochs.
+        backbone = dict(
+            type='BEiTPretrainViT',
+            arch='base',
+            patch_size=16,
+            out_indices=[-4, -1],
+            drop_path_rate=drop_path_rate,
+            final_norm=False,
+            layer_scale_init_value=layer_scale_init_value)
+        neck = dict(
+            type='BEiTV2Neck',
+            num_layers=1,
+            early_layers=9,
+            backbone_arch='base',
+            drop_path_rate=drop_path_rate,
+            layer_scale_init_value=layer_scale_init_value)
+        head = dict(
+            type='BEiTV2Head',
+            embed_dims=768,
+            num_embed=8192,
+            loss=dict(type='CrossEntropyLoss'))
+        target_generator = dict(type='VQKD', encoder_config=vqkd_encoder)
+
+        model = BEiT(
+            backbone=backbone,
+            neck=neck,
+            head=head,
+            target_generator=target_generator,
+            data_preprocessor=data_preprocessor)
+
+        fake_img = torch.rand((1, 3, 224, 224))
+        fake_target_img = torch.rand((1, 3, 224, 224))
+        fake_mask = torch.zeros((196)).bool()
+        fake_mask[75:150] = 1
+        fake_data_sample = DataSample()
+        fake_data_sample.set_mask(fake_mask)
+        fake_data = {
+            'inputs': [fake_img, fake_target_img],
+            'data_samples': [fake_data_sample]
+        }
+
+        fake_inputs = model.data_preprocessor(fake_data)
+        fake_outputs = model(**fake_inputs, mode='loss')
+        assert isinstance(fake_outputs['loss_1'].item(), float)
+        assert isinstance(fake_outputs['loss_2'].item(), float)

tests/test_models/test_selfsup/test_byol.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import BYOL
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_byol():
+    data_preprocessor = dict(
+        mean=(123.675, 116.28, 103.53),
+        std=(58.395, 57.12, 57.375),
+        to_rgb=True)
+    backbone = dict(type='ResNet', depth=18, norm_cfg=dict(type='BN'))
+    neck = dict(
+        type='NonLinearNeck',
+        in_channels=512,
+        hid_channels=2,
+        out_channels=2,
+        with_bias=True,
+        with_last_bn=False,
+        with_avg_pool=True,
+        norm_cfg=dict(type='BN1d'))
+    head = dict(
+        type='LatentPredictHead',
+        loss=dict(type='CosineSimilarityLoss'),
+        predictor=dict(
+            type='NonLinearNeck',
+            in_channels=2,
+            hid_channels=2,
+            out_channels=2,
+            with_bias=True,
+            with_last_bn=False,
+            with_avg_pool=False,
+            norm_cfg=dict(type='BN1d')))
+
+    alg = BYOL(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    fake_data = {
+        'inputs':
+        [torch.randn((2, 3, 224, 224)),
+         torch.randn((2, 3, 224, 224))],
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+
+    fake_loss = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_loss['loss'].item(), float)
+    assert fake_loss['loss'].item() > -4
+
+    fake_feats = alg(fake_inputs['inputs'][0], mode='tensor')
+    assert list(fake_feats[0].shape) == [2, 512, 7, 7]

tests/test_models/test_selfsup/test_cae.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import CAE, CAEPretrainViT
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_cae_vit():
+    backbone = dict(
+        arch='deit-tiny', patch_size=16, layer_scale_init_value=0.1)
+
+    cae_backbone = CAEPretrainViT(**backbone)
+    cae_backbone.init_weights()
+    fake_inputs = torch.randn((1, 3, 224, 224))
+    fake_mask = torch.zeros((1, 196)).bool()
+    fake_mask[:, 75:150] = 1
+
+    # test with mask
+    fake_outputs = cae_backbone(fake_inputs, fake_mask)
+    assert list(fake_outputs.shape) == [1, 122, 192]
+
+    # test without mask
+    fake_outputs = cae_backbone(fake_inputs, None)
+    assert fake_outputs[0].shape == torch.Size([1, 197, 192])
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_cae():
+    data_preprocessor = dict(
+        type='TwoNormDataPreprocessor',
+        mean=[123.675, 116.28, 103.53],
+        std=[58.395, 57.12, 57.375],
+        second_mean=[-31.875, -31.875, -31.875],
+        second_std=[318.75, 318.75, 318.75],
+        to_rgb=True)
+
+    # model settings
+    backbone = dict(
+        type='CAEPretrainViT',
+        arch='deit-tiny',
+        patch_size=16,
+        layer_scale_init_value=0.1)
+    neck = dict(
+        type='CAENeck',
+        embed_dims=192,
+        num_heads=12,
+        regressor_depth=4,
+        decoder_depth=4,
+        mlp_ratio=4,
+        layer_scale_init_value=0.1)
+    head = dict(type='CAEHead', loss=dict(type='CAELoss', lambd=2))
+    target_generator = dict(type='DALL-E')
+
+    model = CAE(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        target_generator=target_generator,
+        data_preprocessor=data_preprocessor)
+
+    fake_img = torch.rand((1, 3, 224, 224))
+    fake_target_img = torch.rand((1, 3, 112, 112))
+    fake_mask = torch.zeros((196)).bool()
+    fake_mask[75:150] = 1
+    fake_data_sample = DataSample()
+    fake_data_sample.set_mask(fake_mask)
+    fake_data = {
+        'inputs': [fake_img, fake_target_img],
+        'data_samples': [fake_data_sample]
+    }
+
+    fake_inputs = model.data_preprocessor(fake_data)
+    fake_outputs = model(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_densecl.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import DenseCL
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_densecl():
+    data_preprocessor = {
+        'mean': (123.675, 116.28, 103.53),
+        'std': (58.395, 57.12, 57.375),
+        'to_rgb': True
+    }
+    queue_len = 32
+    feat_dim = 2
+    momentum = 0.001
+    loss_lambda = 0.5
+    backbone = dict(type='ResNet', depth=18, norm_cfg=dict(type='BN'))
+    neck = dict(
+        type='DenseCLNeck',
+        in_channels=512,
+        hid_channels=2,
+        out_channels=2,
+        num_grid=None)
+    head = dict(
+        type='ContrastiveHead',
+        loss=dict(type='CrossEntropyLoss'),
+        temperature=0.2)
+
+    alg = DenseCL(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        queue_len=queue_len,
+        feat_dim=feat_dim,
+        momentum=momentum,
+        loss_lambda=loss_lambda,
+        data_preprocessor=data_preprocessor)
+
+    # test init
+    assert alg.queue.size() == torch.Size([feat_dim, queue_len])
+    assert alg.queue2.size() == torch.Size([feat_dim, queue_len])
+
+    # test loss
+    fake_data = {
+        'inputs':
+        [torch.randn((2, 3, 224, 224)),
+         torch.randn((2, 3, 224, 224))],
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_loss = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_loss['loss_single'].item(), float)
+    assert isinstance(fake_loss['loss_dense'].item(), float)
+    assert fake_loss['loss_single'].item() > 0
+    assert fake_loss['loss_dense'].item() > 0
+    assert alg.queue_ptr.item() == 2
+    assert alg.queue2_ptr.item() == 2

tests/test_models/test_selfsup/test_eva.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from mmpretrain.models import EVA
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_eva():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+    backbone = dict(type='MAEViT', arch='b', patch_size=16, mask_ratio=0.75)
+    neck = dict(
+        type='MAEPretrainDecoder',
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        predict_feature_dim=512,
+        mlp_ratio=4.)
+    head = dict(
+        type='MIMHead',
+        loss=dict(
+            type='CosineSimilarityLoss', shift_factor=1.0, scale_factor=1.0))
+
+    alg = EVA(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    target_generator = MagicMock(
+        return_value=(torch.ones(2, 197, 512), torch.ones(2, 197, 197)))
+    alg.target_generator = target_generator
+
+    fake_data = {
+        'inputs': torch.randn((2, 3, 224, 224)),
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_outputs = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_itpn.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import iTPN
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_itpn():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+    backbone = dict(
+        type='iTPNHiViT',
+        arch='base',
+        reconstruction_type='pixel',
+        mask_ratio=0.75)
+    neck = dict(
+        type='iTPNPretrainDecoder',
+        num_patches=196,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=512,
+        decoder_embed_dim=512,
+        decoder_depth=6,
+        decoder_num_heads=16,
+        mlp_ratio=4.,
+        reconstruction_type='pixel',
+        #  transformer pyramid
+        fpn_dim=256,
+        fpn_depth=2,
+        num_outs=3,
+    )
+    head = dict(
+        type='MAEPretrainHead',
+        norm_pix=True,
+        patch_size=16,
+        loss=dict(type='PixelReconstructionLoss', criterion='L2'))
+
+    alg = iTPN(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    fake_data = {
+        'inputs': torch.randn((2, 3, 224, 224)),
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_outputs = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_mae.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import MAE, MAEViT
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_mae_vit():
+    backbone = dict(arch='b', patch_size=16, mask_ratio=0.75)
+    mae_backbone = MAEViT(**backbone)
+    mae_backbone.init_weights()
+    fake_inputs = torch.randn((2, 3, 224, 224))
+
+    # test with mask
+    fake_outputs = mae_backbone(fake_inputs)[0]
+    assert list(fake_outputs.shape) == [2, 50, 768]
+
+    # test without mask
+    fake_outputs = mae_backbone(fake_inputs, None)
+    assert fake_outputs[0].shape == torch.Size([2, 197, 768])
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_mae():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+    backbone = dict(type='MAEViT', arch='b', patch_size=16, mask_ratio=0.75)
+    neck = dict(
+        type='MAEPretrainDecoder',
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        mlp_ratio=4.,
+    )
+    loss = dict(type='PixelReconstructionLoss', criterion='L2')
+    head = dict(
+        type='MAEPretrainHead', norm_pix=False, patch_size=16, loss=loss)
+
+    alg = MAE(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    fake_data = {
+        'inputs': torch.randn((2, 3, 224, 224)),
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_outputs = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_maskfeat.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+from mmengine.utils import digit_version
+
+from mmpretrain.models import MaskFeat, MaskFeatViT
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_maskfeat_vit():
+    maskfeat_backbone = MaskFeatViT()
+    maskfeat_backbone.init_weights()
+    fake_inputs = torch.randn((2, 3, 224, 224))
+    fake_mask = torch.randn((2, 14, 14)).flatten(1).bool()
+
+    # test with mask
+    fake_outputs = maskfeat_backbone(fake_inputs, fake_mask)
+    assert list(fake_outputs.shape) == [2, 197, 768]
+
+    # test without mask
+    fake_outputs = maskfeat_backbone(fake_inputs, None)
+    assert fake_outputs[0].shape == torch.Size([2, 197, 768])
+
+
+@pytest.mark.skipif(
+    digit_version(torch.__version__) < digit_version('1.7.0'),
+    reason='torch version')
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_maskfeat():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+
+    backbone = dict(type='MaskFeatViT', arch='b', patch_size=16)
+    neck = dict(
+        type='LinearNeck', in_channels=768, out_channels=108, gap_dim=0)
+    head = dict(
+        type='MIMHead',
+        loss=dict(type='PixelReconstructionLoss', criterion='L2'))
+    target_generator = dict(
+        type='HOGGenerator', nbins=9, pool=8, gaussian_window=16)
+
+    alg = MaskFeat(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        target_generator=target_generator,
+        data_preprocessor=data_preprocessor)
+
+    # test forward_train
+    fake_data_sample = DataSample()
+    fake_mask = torch.rand((14, 14)).bool()
+    fake_data_sample.set_mask(fake_mask)
+    fake_data = {
+        'inputs': torch.randn((1, 3, 224, 224)),
+        'data_samples': [fake_data_sample]
+    }
+
+    fake_input = alg.data_preprocessor(fake_data)
+    fake_outputs = alg(**fake_input, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_mff.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import MFF, MFFViT
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_mae_vit():
+    backbone = dict(
+        arch='b', patch_size=16, mask_ratio=0.75, out_indices=[1, 11])
+    mae_backbone = MFFViT(**backbone)
+    mae_backbone.init_weights()
+    fake_inputs = torch.randn((2, 3, 224, 224))
+
+    # test with mask
+    fake_outputs = mae_backbone(fake_inputs)[0]
+    assert list(fake_outputs.shape) == [2, 50, 768]
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_mae():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+    backbone = dict(
+        type='MFFViT',
+        arch='b',
+        patch_size=16,
+        mask_ratio=0.75,
+        out_indices=[1, 11])
+    neck = dict(
+        type='MAEPretrainDecoder',
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        mlp_ratio=4.,
+    )
+    loss = dict(type='PixelReconstructionLoss', criterion='L2')
+    head = dict(
+        type='MAEPretrainHead', norm_pix=False, patch_size=16, loss=loss)
+
+    alg = MFF(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    fake_data = {
+        'inputs': torch.randn((2, 3, 224, 224)),
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_outputs = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_milan.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+import platform
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from mmpretrain.models import MILAN, MILANViT
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_milan_vit():
+    backbone = dict(arch='b', patch_size=16, mask_ratio=0.75)
+    milan_backbone = MILANViT(**backbone)
+    milan_backbone.init_weights()
+    fake_inputs = torch.randn((2, 3, 224, 224))
+
+    # test with mask
+    fake_outputs = milan_backbone(fake_inputs,
+                                  torch.ones(2, 197, 197)[:, 0, 1:])[0]
+    assert list(fake_outputs.shape) == [2, 50, 768]
+
+    # test without mask
+    fake_outputs = milan_backbone(fake_inputs, None)
+    assert fake_outputs[0].shape == torch.Size([2, 197, 768])
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_milan():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+
+    backbone = dict(type='MILANViT', arch='b', patch_size=16, mask_ratio=0.75)
+    neck = dict(
+        type='MILANPretrainDecoder',
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        mlp_ratio=4.)
+    head = dict(
+        type='MIMHead',
+        loss=dict(
+            type='CosineSimilarityLoss', shift_factor=2.0, scale_factor=2.0))
+
+    alg = MILAN(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=copy.deepcopy(data_preprocessor))
+
+    target_generator = MagicMock(
+        return_value=(torch.ones(2, 197, 512), torch.ones(2, 197, 197)))
+    alg.target_generator = target_generator
+
+    fake_data = {
+        'inputs': torch.randn((2, 3, 224, 224)),
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_outputs = alg(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_mixmim.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import MixMIM, MixMIMPretrainTransformer
+from mmpretrain.structures import DataSample
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_mixmmim_backbone():
+    mixmmim_backbone = MixMIMPretrainTransformer(
+        arch=dict(embed_dims=128, depths=[2, 2, 4, 2], num_heads=[4, 4, 4, 4]))
+    mixmmim_backbone.init_weights()
+    fake_inputs = torch.randn((1, 3, 224, 224))
+
+    # test with mask
+    fake_outputs, fake_mask_s4 = mixmmim_backbone(fake_inputs)
+    assert fake_outputs.shape == torch.Size([1, 49, 1024])
+    assert fake_mask_s4.shape == torch.Size([1, 49, 1])
+
+    # test without mask
+    fake_outputs = mixmmim_backbone(fake_inputs, None)
+    assert len(fake_outputs) == 1
+    assert fake_outputs[0].shape == torch.Size([1, 1024])
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_simmim():
+    data_preprocessor = {
+        'mean': [0.5, 0.5, 0.5],
+        'std': [0.5, 0.5, 0.5],
+        'to_rgb': True
+    }
+
+    # model config
+    backbone = dict(
+        type='MixMIMPretrainTransformer',
+        arch='B',
+        drop_rate=0.0,
+        drop_path_rate=0.0)
+    neck = dict(
+        type='MixMIMPretrainDecoder',
+        num_patches=49,
+        encoder_stride=32,
+        embed_dim=1024,
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16)
+    head = dict(
+        type='MixMIMPretrainHead',
+        norm_pix=True,
+        loss=dict(type='PixelReconstructionLoss', criterion='L2'))
+
+    model = MixMIM(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        data_preprocessor=data_preprocessor)
+
+    # test forward_train
+    fake_data_sample = DataSample()
+    fake_data = {
+        'inputs': torch.randn((2, 3, 224, 224)),
+        'data_samples': [fake_data_sample for _ in range(2)]
+    }
+
+    fake_inputs = model.data_preprocessor(fake_data)
+    fake_outputs = model(**fake_inputs, mode='loss')
+    assert isinstance(fake_outputs['loss'].item(), float)

tests/test_models/test_selfsup/test_moco.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+
+import pytest
+import torch
+
+from mmpretrain.models import MoCo
+from mmpretrain.structures import DataSample
+
+queue_len = 32
+feat_dim = 2
+momentum = 0.001
+backbone = dict(type='ResNet', depth=18, norm_cfg=dict(type='BN'))
+neck = dict(
+    type='MoCoV2Neck',
+    in_channels=512,
+    hid_channels=2,
+    out_channels=2,
+    with_avg_pool=True)
+head = dict(
+    type='ContrastiveHead',
+    loss=dict(type='CrossEntropyLoss'),
+    temperature=0.2)
+
+
+@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
+def test_moco():
+    data_preprocessor = {
+        'mean': (123.675, 116.28, 103.53),
+        'std': (58.395, 57.12, 57.375),
+        'to_rgb': True
+    }
+
+    alg = MoCo(
+        backbone=backbone,
+        neck=neck,
+        head=head,
+        queue_len=queue_len,
+        feat_dim=feat_dim,
+        momentum=momentum,
+        data_preprocessor=data_preprocessor)
+    assert alg.queue.size() == torch.Size([feat_dim, queue_len])
+
+    fake_data = {
+        'inputs':
+        [torch.randn((2, 3, 224, 224)),
+         torch.randn((2, 3, 224, 224))],
+        'data_samples': [DataSample() for _ in range(2)]
+    }
+
+    fake_inputs = alg.data_preprocessor(fake_data)
+    fake_loss = alg(**fake_inputs, mode='loss')
+    assert fake_loss['loss'] > 0
+    assert alg.queue_ptr.item() == 2
+
+    # test extract
+    fake_feats = alg(fake_inputs['inputs'][0], mode='tensor')
+    assert fake_feats[0].size() == torch.Size([2, 512, 7, 7])

tests/test_models/test_selfsup/test_mocov3.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import platform
+from unittest import TestCase
+
+import pytest
+import torch
+
+from mmpretrain.models import MoCoV3, MoCoV3ViT
+from mmpretrain.structures import DataSample
+
+
+class TestMoCoV3(TestCase):
+
+    backbone = dict(
+        type='MoCoV3ViT',
+        arch='mocov3-small',  # embed_dim = 384
+        patch_size=16,
+        frozen_stages=12,
+        stop_grad_conv1=True,
+        norm_eval=True)
+    neck = dict(
+        type='NonLinearNeck',
+        in_channels=384,
+        hid_channels=2,
+        out_channels=2,
+        num_layers=2,
+        with_bias=False,
+        with_last_bn=True,
+        with_last_bn_affine=False,
+        with_last_bias=False,
+        with_avg_pool=False,
+        norm_cfg=dict(type='BN1d'))
+    head = dict(
+        type='MoCoV3Head',
+        predictor=dict(
+            type='NonLinearNeck',
+            in_channels=2,
+            hid_channels=2,
+            out_channels=2,
+            num_layers=2,
+            with_bias=False,
+            with_last_bn=True,
+            with_last_bn_affine=False,
+            with_last_bias=False,
+            with_avg_pool=False,
+            norm_cfg=dict(type='BN1d')),
+        loss=dict(type='CrossEntropyLoss', loss_weight=2 * 0.2),
+        temperature=0.2)
+
+    @pytest.mark.skipif(
+        platform.system() == 'Windows', reason='Windows mem limit')
+    def test_vit(self):
+        vit = MoCoV3ViT(
+            arch='mocov3-small',
+            patch_size=16,
+            frozen_stages=12,
+            stop_grad_conv1=True,
+            norm_eval=True)
+        vit.init_weights()
+        vit.train()
+
+        for p in vit.parameters():
+            assert p.requires_grad is False
+
+    @pytest.mark.skipif(
+        platform.system() == 'Windows', reason='Windows mem limit')
+    def test_mocov3(self):
+        data_preprocessor = dict(
+            mean=(123.675, 116.28, 103.53),
+            std=(58.395, 57.12, 57.375),
+            to_rgb=True)
+        alg = MoCoV3(
+            backbone=self.backbone,
+            neck=self.neck,
+            head=self.head,
+            data_preprocessor=data_preprocessor)
+
+        fake_data = {
+            'inputs':
+            [torch.randn((2, 3, 224, 224)),
+             torch.randn((2, 3, 224, 224))],
+            'data_samples': [DataSample() for _ in range(2)]
+        }
+
+        fake_inputs = alg.data_preprocessor(fake_data)
+        fake_loss = alg(**fake_inputs, mode='loss')
+        self.assertGreater(fake_loss['loss'], 0)
+
+        # test extract
+        fake_feats = alg(fake_inputs['inputs'][0], mode='tensor')
+        self.assertEqual(fake_feats[0].size(), torch.Size([2, 384]))