|
|
|
|
|
import pytest
|
|
|
import torch
|
|
|
|
|
|
from mmaction.models import SwinTransformer3D
|
|
|
from mmaction.testing import generate_backbone_demo_inputs
|
|
|
|
|
|
|
|
|
def test_swin_backbone():
|
|
|
"""Test swin backbone."""
|
|
|
with pytest.raises(AssertionError):
|
|
|
SwinTransformer3D(arch='-t')
|
|
|
|
|
|
with pytest.raises(AssertionError):
|
|
|
SwinTransformer3D(arch={'embed_dims': 96})
|
|
|
|
|
|
with pytest.raises(AssertionError):
|
|
|
SwinTransformer3D(arch={
|
|
|
'embed_dims': 96,
|
|
|
'depths': [2, 2, 6],
|
|
|
'num_heads': [3, 6, 12, 24]
|
|
|
})
|
|
|
|
|
|
with pytest.raises(AssertionError):
|
|
|
SwinTransformer3D(
|
|
|
arch={
|
|
|
'embed_dims': 96,
|
|
|
'depths': [2, 2, 6, 2, 2],
|
|
|
'num_heads': [3, 6, 12, 24, 48]
|
|
|
})
|
|
|
|
|
|
with pytest.raises(AssertionError):
|
|
|
SwinTransformer3D(arch='t', out_indices=(4, ))
|
|
|
|
|
|
with pytest.raises(TypeError):
|
|
|
swin_t = SwinTransformer3D(arch='t', pretrained=[0, 1, 1])
|
|
|
swin_t.init_weights()
|
|
|
|
|
|
with pytest.raises(TypeError):
|
|
|
swin_t = SwinTransformer3D(arch='t')
|
|
|
swin_t.init_weights(pretrained=[0, 1, 1])
|
|
|
|
|
|
swin_b = SwinTransformer3D(arch='b', pretrained=None, pretrained2d=False)
|
|
|
swin_b.init_weights()
|
|
|
swin_b.train()
|
|
|
|
|
|
pretrained_url = 'https://download.openmmlab.com/mmaction/v1.0/' \
|
|
|
'recognition/swin/swin_tiny_patch4_window7_224.pth'
|
|
|
|
|
|
swin_t_pre = SwinTransformer3D(
|
|
|
arch='t', pretrained=pretrained_url, pretrained2d=True)
|
|
|
swin_t_pre.init_weights()
|
|
|
swin_t_pre.train()
|
|
|
|
|
|
from mmengine.runner.checkpoint import _load_checkpoint
|
|
|
ckpt_2d = _load_checkpoint(pretrained_url, map_location='cpu')
|
|
|
state_dict = ckpt_2d['model']
|
|
|
|
|
|
patch_embed_weight2d = state_dict['patch_embed.proj.weight'].data
|
|
|
patch_embed_weight3d = swin_t_pre.patch_embed.proj.weight.data
|
|
|
assert torch.equal(
|
|
|
patch_embed_weight3d,
|
|
|
patch_embed_weight2d.unsqueeze(2).expand_as(patch_embed_weight3d) /
|
|
|
patch_embed_weight3d.shape[2])
|
|
|
|
|
|
norm = swin_t_pre.norm3
|
|
|
assert torch.equal(norm.weight.data, state_dict['norm.weight'])
|
|
|
assert torch.equal(norm.bias.data, state_dict['norm.bias'])
|
|
|
|
|
|
for name, param in swin_t_pre.named_parameters():
|
|
|
if 'relative_position_bias_table' in name:
|
|
|
bias2d = state_dict[name]
|
|
|
assert torch.equal(
|
|
|
param.data, bias2d.repeat(2 * swin_t_pre.window_size[0] - 1,
|
|
|
1))
|
|
|
|
|
|
frozen_stages = 1
|
|
|
swin_t_frozen = SwinTransformer3D(
|
|
|
arch='t',
|
|
|
pretrained=None,
|
|
|
pretrained2d=False,
|
|
|
frozen_stages=frozen_stages)
|
|
|
swin_t_frozen.init_weights()
|
|
|
swin_t_frozen.train()
|
|
|
for param in swin_t_frozen.patch_embed.parameters():
|
|
|
assert param.requires_grad is False
|
|
|
for i in range(frozen_stages):
|
|
|
layer = swin_t_frozen.layers[i]
|
|
|
for param in layer.parameters():
|
|
|
assert param.requires_grad is False
|
|
|
|
|
|
input_shape = (1, 3, 6, 64, 64)
|
|
|
imgs = generate_backbone_demo_inputs(input_shape)
|
|
|
|
|
|
feat = swin_t_frozen(imgs)
|
|
|
assert feat.shape == torch.Size([1, 768, 3, 2, 2])
|
|
|
|
|
|
input_shape = (1, 3, 5, 63, 63)
|
|
|
imgs = generate_backbone_demo_inputs(input_shape)
|
|
|
feat = swin_t_frozen(imgs)
|
|
|
assert feat.shape == torch.Size([1, 768, 3, 2, 2])
|
|
|
|
|
|
swin_t_all_stages = SwinTransformer3D(arch='t', out_indices=(0, 1, 2, 3))
|
|
|
feats = swin_t_all_stages(imgs)
|
|
|
assert feats[0].shape == torch.Size([1, 96, 3, 16, 16])
|
|
|
assert feats[1].shape == torch.Size([1, 192, 3, 8, 8])
|
|
|
assert feats[2].shape == torch.Size([1, 384, 3, 4, 4])
|
|
|
assert feats[3].shape == torch.Size([1, 768, 3, 2, 2])
|
|
|
|
|
|
swin_t_all_stages_after_ds = SwinTransformer3D(
|
|
|
arch='t', out_indices=(0, 1, 2, 3), out_after_downsample=True)
|
|
|
feats = swin_t_all_stages_after_ds(imgs)
|
|
|
assert feats[0].shape == torch.Size([1, 192, 3, 8, 8])
|
|
|
assert feats[1].shape == torch.Size([1, 384, 3, 4, 4])
|
|
|
assert feats[2].shape == torch.Size([1, 768, 3, 2, 2])
|
|
|
assert feats[3].shape == torch.Size([1, 768, 3, 2, 2])
|
|
|
|
