mmpose/tests/test_models/test_backbones/test_shufflenet_v1.py

# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase

import torch
from torch.nn.modules import GroupNorm
from torch.nn.modules.batchnorm import _BatchNorm

from mmpose.models.backbones import ShuffleNetV1
from mmpose.models.backbones.shufflenet_v1 import ShuffleUnit


class TestShufflenetV1(TestCase):

    @staticmethod
    def is_block(modules):
        """Check if is ResNet building block."""
        if isinstance(modules, (ShuffleUnit, )):
            return True
        return False

    @staticmethod
    def is_norm(modules):
        """Check if is one of the norms."""
        if isinstance(modules, (GroupNorm, _BatchNorm)):
            return True
        return False

    @staticmethod
    def check_norm_state(modules, train_state):
        """Check if norm layer is in correct train state."""
        for mod in modules:
            if isinstance(mod, _BatchNorm):
                if mod.training != train_state:
                    return False
        return True

    def test_shufflenetv1_shuffleuint(self):

        with self.assertRaises(ValueError):
            # combine must be in ['add', 'concat']
            ShuffleUnit(24, 16, groups=3, first_block=True, combine='test')

        with self.assertRaises(AssertionError):
            # inplanes must be equal tp = outplanes when combine='add'
            ShuffleUnit(64, 24, groups=4, first_block=True, combine='add')

        # Test ShuffleUnit with combine='add'
        block = ShuffleUnit(24, 24, groups=3, first_block=True, combine='add')
        x = torch.randn(1, 24, 56, 56)
        x_out = block(x)
        self.assertEqual(x_out.shape, torch.Size((1, 24, 56, 56)))

        # Test ShuffleUnit with combine='concat'
        block = ShuffleUnit(
            24, 240, groups=3, first_block=True, combine='concat')
        x = torch.randn(1, 24, 56, 56)
        x_out = block(x)
        self.assertEqual(x_out.shape, torch.Size((1, 240, 28, 28)))

        # Test ShuffleUnit with checkpoint forward
        block = ShuffleUnit(
            24, 24, groups=3, first_block=True, combine='add', with_cp=True)
        self.assertTrue(block.with_cp)
        x = torch.randn(1, 24, 56, 56)
        x.requires_grad = True
        x_out = block(x)
        self.assertEqual(x_out.shape, torch.Size((1, 24, 56, 56)))

    def test_shufflenetv1_backbone(self):

        with self.assertRaises(ValueError):
            # frozen_stages must be in  range(-1, 4)
            ShuffleNetV1(frozen_stages=10)

        with self.assertRaises(ValueError):
            # the item in out_indices must be in  range(0, 4)
            ShuffleNetV1(out_indices=[5])

        with self.assertRaises(ValueError):
            # groups must be in  [1, 2, 3, 4, 8]
            ShuffleNetV1(groups=10)

        # Test ShuffleNetV1 norm state
        model = ShuffleNetV1()
        model.init_weights()
        model.train()
        self.assertTrue(self.check_norm_state(model.modules(), True))

        # Test ShuffleNetV1 with first stage frozen
        frozen_stages = 1
        model = ShuffleNetV1(
            frozen_stages=frozen_stages, out_indices=(0, 1, 2))
        model.init_weights()
        model.train()
        for param in model.conv1.parameters():
            self.assertFalse(param.requires_grad)
        for i in range(frozen_stages):
            layer = model.layers[i]
            for mod in layer.modules():
                if isinstance(mod, _BatchNorm):
                    self.assertFalse(mod.training)
            for param in layer.parameters():
                self.assertFalse(param.requires_grad)

        # Test ShuffleNetV1 forward with groups=1
        model = ShuffleNetV1(groups=1, out_indices=(0, 1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 3)
        self.assertEqual(feat[0].shape, torch.Size((1, 144, 28, 28)))
        self.assertEqual(feat[1].shape, torch.Size((1, 288, 14, 14)))
        self.assertEqual(feat[2].shape, torch.Size((1, 576, 7, 7)))

        # Test ShuffleNetV1 forward with groups=2
        model = ShuffleNetV1(groups=2, out_indices=(0, 1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 3)
        self.assertEqual(feat[0].shape, torch.Size((1, 200, 28, 28)))
        self.assertEqual(feat[1].shape, torch.Size((1, 400, 14, 14)))
        self.assertEqual(feat[2].shape, torch.Size((1, 800, 7, 7)))

        # Test ShuffleNetV1 forward with groups=3
        model = ShuffleNetV1(groups=3, out_indices=(0, 1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 3)
        self.assertEqual(feat[0].shape, torch.Size((1, 240, 28, 28)))
        self.assertEqual(feat[1].shape, torch.Size((1, 480, 14, 14)))
        self.assertEqual(feat[2].shape, torch.Size((1, 960, 7, 7)))

        # Test ShuffleNetV1 forward with groups=4
        model = ShuffleNetV1(groups=4, out_indices=(0, 1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 3)
        self.assertEqual(feat[0].shape, torch.Size((1, 272, 28, 28)))
        self.assertEqual(feat[1].shape, torch.Size((1, 544, 14, 14)))
        self.assertEqual(feat[2].shape, torch.Size((1, 1088, 7, 7)))

        # Test ShuffleNetV1 forward with groups=8
        model = ShuffleNetV1(groups=8, out_indices=(0, 1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 3)
        self.assertEqual(feat[0].shape, torch.Size((1, 384, 28, 28)))
        self.assertEqual(feat[1].shape, torch.Size((1, 768, 14, 14)))
        self.assertEqual(feat[2].shape, torch.Size((1, 1536, 7, 7)))

        # Test ShuffleNetV1 forward with GroupNorm forward
        model = ShuffleNetV1(
            groups=3,
            norm_cfg=dict(type='GN', num_groups=2, requires_grad=True),
            out_indices=(0, 1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, GroupNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 3)
        self.assertEqual(feat[0].shape, torch.Size((1, 240, 28, 28)))
        self.assertEqual(feat[1].shape, torch.Size((1, 480, 14, 14)))
        self.assertEqual(feat[2].shape, torch.Size((1, 960, 7, 7)))

        # Test ShuffleNetV1 forward with layers 1, 2 forward
        model = ShuffleNetV1(groups=3, out_indices=(1, 2))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertEqual(len(feat), 2)
        self.assertEqual(feat[0].shape, torch.Size((1, 480, 14, 14)))
        self.assertEqual(feat[1].shape, torch.Size((1, 960, 7, 7)))

        # Test ShuffleNetV1 forward with layers 2 forward
        model = ShuffleNetV1(groups=3, out_indices=(2, ))
        model.init_weights()
        model.train()

        for m in model.modules():
            if self.is_norm(m):
                self.assertIsInstance(m, _BatchNorm)

        imgs = torch.randn(1, 3, 224, 224)
        feat = model(imgs)
        self.assertIsInstance(feat, tuple)
        self.assertEqual(feat[-1].shape, torch.Size((1, 960, 7, 7)))

        # Test ShuffleNetV1 forward with checkpoint forward
        model = ShuffleNetV1(groups=3, with_cp=True)
        for m in model.modules():
            if self.is_block(m):
                self.assertTrue(m.with_cp)

        # Test ShuffleNetV1 with norm_eval
        model = ShuffleNetV1(norm_eval=True)
        model.init_weights()
        model.train()

        self.assertTrue(self.check_norm_state(model.modules(), False))