mmpose-0.29.0/tests/test_backbones/test_hrformer.py

# Copyright (c) OpenMMLab. All rights reserved.
import pytest
import torch

from mmpose.models.backbones.hrformer import (HRFomerModule, HRFormer,
 HRFormerBlock)


def test_hrformer_module():
 norm_cfg = dict(type='BN')
 block = HRFormerBlock
 # Test multiscale forward
 num_channles = (32, 64)
 num_inchannels = [c * block.expansion for c in num_channles]
 hrmodule = HRFomerModule(
 num_branches=2,
 block=block,
 num_blocks=(2, 2),
 num_inchannels=num_inchannels,
 num_channels=num_channles,
 num_heads=(1, 2),
 num_window_sizes=(7, 7),
 num_mlp_ratios=(4, 4),
 drop_paths=(0., 0.),
 norm_cfg=norm_cfg)

 feats = [
 torch.randn(1, num_inchannels[0], 64, 64),
 torch.randn(1, num_inchannels[1], 32, 32)
 ]
 feats = hrmodule(feats)

 assert len(str(hrmodule)) > 0
 assert len(feats) == 2
 assert feats[0].shape == torch.Size([1, num_inchannels[0], 64, 64])
 assert feats[1].shape == torch.Size([1, num_inchannels[1], 32, 32])

 # Test single scale forward
 num_channles = (32, 64)
 in_channels = [c * block.expansion for c in num_channles]
 hrmodule = HRFomerModule(
 num_branches=2,
 block=block,
 num_blocks=(2, 2),
 num_inchannels=num_inchannels,
 num_channels=num_channles,
 num_heads=(1, 2),
 num_window_sizes=(7, 7),
 num_mlp_ratios=(4, 4),
 drop_paths=(0., 0.),
 norm_cfg=norm_cfg,
 multiscale_output=False,
 )

 feats = [
 torch.randn(1, in_channels[0], 64, 64),
 torch.randn(1, in_channels[1], 32, 32)
 ]
 feats = hrmodule(feats)

 assert len(feats) == 1
 assert feats[0].shape == torch.Size([1, in_channels[0], 64, 64])

 # Test single branch HRFormer module
 hrmodule = HRFomerModule(
 num_branches=1,
 block=block,
 num_blocks=(1, ),
 num_inchannels=[num_inchannels[0]],
 num_channels=[num_channles[0]],
 num_heads=(1, ),
 num_window_sizes=(7, ),
 num_mlp_ratios=(4, ),
 drop_paths=(0.1, ),
 norm_cfg=norm_cfg,
 )

 feats = [
 torch.randn(1, in_channels[0], 64, 64),
 ]
 feats = hrmodule(feats)

 assert len(feats) == 1
 assert feats[0].shape == torch.Size([1, in_channels[0], 64, 64])

 # Value tests
 kwargs = dict(
 num_branches=2,
 block=block,
 num_blocks=(2, 2),
 num_inchannels=num_inchannels,
 num_channels=num_channles,
 num_heads=(1, 2),
 num_window_sizes=(7, 7),
 num_mlp_ratios=(4, 4),
 drop_paths=(0.1, 0.1),
 norm_cfg=norm_cfg,
 )

 with pytest.raises(ValueError):
 # len(num_blocks) should equal num_branches
 kwargs['num_blocks'] = [2, 2, 2]
 HRFomerModule(**kwargs)
 kwargs['num_blocks'] = [2, 2]

 with pytest.raises(ValueError):
 # len(num_blocks) should equal num_branches
 kwargs['num_channels'] = [2]
 HRFomerModule(**kwargs)
 kwargs['num_channels'] = [2, 2]

 with pytest.raises(ValueError):
 # len(num_blocks) should equal num_branches
 kwargs['num_inchannels'] = [2]
 HRFomerModule(**kwargs)
 kwargs['num_inchannels'] = [2, 2]


def test_hrformer_backbone():
 norm_cfg = dict(type='BN')
 # only have 3 stages
 extra = dict(
 drop_path_rate=0.2,
 stage1=dict(
 num_modules=1,
 num_branches=1,
 block='BOTTLENECK',
 num_blocks=(2, ),
 num_channels=(64, )),
 stage2=dict(
 num_modules=1,
 num_branches=2,
 block='HRFORMERBLOCK',
 window_sizes=(7, 7),
 num_heads=(1, 2),
 mlp_ratios=(4, 4),
 num_blocks=(2, 2),
 num_channels=(32, 64)),
 stage3=dict(
 num_modules=4,
 num_branches=3,
 block='HRFORMERBLOCK',
 window_sizes=(7, 7, 7),
 num_heads=(1, 2, 4),
 mlp_ratios=(4, 4, 4),
 num_blocks=(2, 2, 2),
 num_channels=(32, 64, 128)),
 stage4=dict(
 num_modules=3,
 num_branches=4,
 block='HRFORMERBLOCK',
 window_sizes=(7, 7, 7, 7),
 num_heads=(1, 2, 4, 8),
 mlp_ratios=(4, 4, 4, 4),
 num_blocks=(2, 2, 2, 2),
 num_channels=(32, 64, 128, 256),
 multiscale_output=True))

 with pytest.raises(ValueError):
 # len(num_blocks) should equal num_branches
 extra['stage4']['num_branches'] = 3
 HRFormer(extra=extra)
 extra['stage4']['num_branches'] = 4

 # Test HRFormer-S
 model = HRFormer(extra=extra, norm_cfg=norm_cfg)
 model.init_weights()
 model.train()

 imgs = torch.randn(1, 3, 64, 64)
 feats = model(imgs)
 assert len(feats) == 4
 assert feats[0].shape == torch.Size([1, 32, 16, 16])
 assert feats[3].shape == torch.Size([1, 256, 2, 2])

 # Test single scale output and model
 # without relative position bias
 extra['stage4']['multiscale_output'] = False
 extra['with_rpe'] = False
 model = HRFormer(extra=extra, norm_cfg=norm_cfg)
 model.init_weights()
 model.train()

 imgs = torch.randn(1, 3, 64, 64)
 feats = model(imgs)
 assert len(feats) == 1
 assert feats[0].shape == torch.Size([1, 32, 16, 16])