|
|
|
|
|
import os
|
|
|
import tempfile
|
|
|
|
|
|
import torch
|
|
|
from mmengine.runner import load_checkpoint, save_checkpoint
|
|
|
from mmengine.runner.checkpoint import _load_checkpoint_with_prefix
|
|
|
|
|
|
from mmaction.models.backbones.mobileone_tsm import MobileOneTSM
|
|
|
from mmaction.testing import generate_backbone_demo_inputs
|
|
|
|
|
|
|
|
|
def test_mobileone_tsm_backbone():
|
|
|
"""Test MobileOne TSM backbone."""
|
|
|
|
|
|
from mmpretrain.models.backbones.mobileone import MobileOneBlock
|
|
|
|
|
|
from mmaction.models.backbones.resnet_tsm import TemporalShift
|
|
|
|
|
|
model = MobileOneTSM('s0', pretrained2d=False)
|
|
|
model.init_weights()
|
|
|
for cur_module in model.modules():
|
|
|
if isinstance(cur_module, TemporalShift):
|
|
|
|
|
|
assert isinstance(cur_module.net, MobileOneBlock)
|
|
|
assert cur_module.num_segments == model.num_segments
|
|
|
assert cur_module.shift_div == model.shift_div
|
|
|
|
|
|
inputs = generate_backbone_demo_inputs((8, 3, 64, 64))
|
|
|
|
|
|
feat = model(inputs)
|
|
|
assert feat.shape == torch.Size([8, 1024, 2, 2])
|
|
|
|
|
|
model = MobileOneTSM('s1', pretrained2d=False)
|
|
|
feat = model(inputs)
|
|
|
assert feat.shape == torch.Size([8, 1280, 2, 2])
|
|
|
|
|
|
model = MobileOneTSM('s2', pretrained2d=False)
|
|
|
feat = model(inputs)
|
|
|
assert feat.shape == torch.Size([8, 2048, 2, 2])
|
|
|
|
|
|
model = MobileOneTSM('s3', pretrained2d=False)
|
|
|
feat = model(inputs)
|
|
|
assert feat.shape == torch.Size([8, 2048, 2, 2])
|
|
|
|
|
|
model = MobileOneTSM('s4', pretrained2d=False)
|
|
|
feat = model(inputs)
|
|
|
assert feat.shape == torch.Size([8, 2048, 2, 2])
|
|
|
|
|
|
|
|
|
def test_mobileone_init_weight():
|
|
|
checkpoint = ('https://download.openmmlab.com/mmclassification/v0'
|
|
|
'/mobileone/mobileone-s0_8xb32_in1k_20221110-0bc94952.pth')
|
|
|
|
|
|
model = MobileOneTSM(
|
|
|
arch='s0',
|
|
|
init_cfg=dict(
|
|
|
type='Pretrained', checkpoint=checkpoint, prefix='backbone'))
|
|
|
model.init_weights()
|
|
|
ori_ckpt = _load_checkpoint_with_prefix(
|
|
|
'backbone', model.init_cfg['checkpoint'], map_location='cpu')
|
|
|
for name, param in model.named_parameters():
|
|
|
ori_name = name.replace('.net', '')
|
|
|
assert torch.allclose(param, ori_ckpt[ori_name]), \
|
|
|
f'layer {name} fail to load from pretrained checkpoint'
|
|
|
|
|
|
|
|
|
def test_load_deploy_mobileone():
|
|
|
|
|
|
model = MobileOneTSM('s0', pretrained2d=False)
|
|
|
inputs = generate_backbone_demo_inputs((8, 3, 64, 64))
|
|
|
tmpdir = tempfile.gettempdir()
|
|
|
ckpt_path = os.path.join(tmpdir, 'ckpt.pth')
|
|
|
model.switch_to_deploy()
|
|
|
model.eval()
|
|
|
outputs = model(inputs)
|
|
|
|
|
|
model_deploy = MobileOneTSM('s0', pretrained2d=False, deploy=True)
|
|
|
save_checkpoint(model.state_dict(), ckpt_path)
|
|
|
load_checkpoint(model_deploy, ckpt_path)
|
|
|
|
|
|
outputs_load = model_deploy(inputs)
|
|
|
for feat, feat_load in zip(outputs, outputs_load):
|
|
|
assert torch.allclose(feat, feat_load)
|
|
|
os.remove(ckpt_path)
|
|
|
|
