|
|
|
|
|
from unittest import TestCase
|
|
|
|
|
|
import pytest
|
|
|
import torch
|
|
|
|
|
|
from mmaction.models import FeatureHead
|
|
|
from mmaction.registry import MODELS
|
|
|
from mmaction.structures import ActionDataSample
|
|
|
from mmaction.testing import get_recognizer_cfg
|
|
|
from mmaction.utils import register_all_modules
|
|
|
|
|
|
|
|
|
class TestFeatureHead(TestCase):
|
|
|
|
|
|
def test_2d_recognizer(self):
|
|
|
register_all_modules()
|
|
|
config = get_recognizer_cfg(
|
|
|
'tsn/tsn_imagenet-pretrained-r50_8xb32-1x1x3-100e_kinetics400-rgb.py'
|
|
|
)
|
|
|
config.model['backbone']['pretrained'] = None
|
|
|
config.model['cls_head'] = dict(
|
|
|
type='FeatureHead', average_clips='score')
|
|
|
|
|
|
recognizer = MODELS.build(config.model)
|
|
|
|
|
|
input_shape = [3, 3, 32, 32]
|
|
|
data_batch = {
|
|
|
'inputs': [torch.randint(0, 256, input_shape)],
|
|
|
'data_samples': [ActionDataSample().set_gt_label(2)]
|
|
|
}
|
|
|
feat = recognizer.test_step(data_batch)
|
|
|
assert isinstance(feat, torch.Tensor)
|
|
|
assert feat.shape == torch.Size([1, 2048])
|
|
|
|
|
|
def test_3d_recognizer(self):
|
|
|
register_all_modules()
|
|
|
config = get_recognizer_cfg(
|
|
|
'slowonly/slowonly_r50_8xb16-4x16x1-256e_kinetics400-rgb.py')
|
|
|
config.model['backbone']['pretrained'] = None
|
|
|
config.model['backbone']['pretrained2d'] = False
|
|
|
config.model['cls_head'] = dict(
|
|
|
type='FeatureHead', average_clips='score')
|
|
|
|
|
|
recognizer = MODELS.build(config.model)
|
|
|
input_shape = [1, 3, 4, 32, 32]
|
|
|
data_batch = {
|
|
|
'inputs': [torch.randint(0, 256, input_shape)],
|
|
|
'data_samples': [ActionDataSample().set_gt_label(2)]
|
|
|
}
|
|
|
feat = recognizer.test_step(data_batch)
|
|
|
assert isinstance(feat, torch.Tensor)
|
|
|
assert feat.shape == torch.Size([1, 2048])
|
|
|
|
|
|
def test_3d_backbone(self):
|
|
|
with pytest.raises(NotImplementedError):
|
|
|
head = FeatureHead(spatial_type='test')
|
|
|
|
|
|
head = FeatureHead(average_clips='score')
|
|
|
x = torch.rand(1, 64, 2, 7, 7)
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 64])
|
|
|
|
|
|
head = FeatureHead(spatial_type=None, average_clips='score')
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 64, 7, 7])
|
|
|
|
|
|
head = FeatureHead(temporal_type=None, average_clips='score')
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 64, 2])
|
|
|
|
|
|
head = FeatureHead(
|
|
|
spatial_type=None, temporal_type=None, average_clips='score')
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 64, 2, 7, 7])
|
|
|
|
|
|
def test_slowfast_backbone(self):
|
|
|
head = FeatureHead(backbone_name='slowfast', average_clips='score')
|
|
|
x_slow = torch.rand(1, 64, 2, 7, 7)
|
|
|
x_fast = torch.rand(1, 32, 6, 7, 7)
|
|
|
x = (x_slow, x_fast)
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 96])
|
|
|
|
|
|
head = FeatureHead(
|
|
|
backbone_name='slowfast', spatial_type=None, average_clips='score')
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 96, 7, 7])
|
|
|
|
|
|
with pytest.raises(AssertionError):
|
|
|
head = FeatureHead(
|
|
|
backbone_name='slowfast',
|
|
|
temporal_type=None,
|
|
|
average_clips='score')
|
|
|
feat = head(x)
|
|
|
|
|
|
def test_2d_backbone(self):
|
|
|
head = FeatureHead(average_clips='score')
|
|
|
x = torch.rand(2, 64, 7, 7)
|
|
|
with pytest.raises(AssertionError):
|
|
|
feat = head(x)
|
|
|
|
|
|
feat = head(x, num_segs=2)
|
|
|
assert feat.shape == torch.Size([1, 64])
|
|
|
|
|
|
x = torch.rand(2, 64, 7, 7)
|
|
|
head = FeatureHead(spatial_type=None, average_clips='score')
|
|
|
feat = head(x, num_segs=2)
|
|
|
assert feat.shape == torch.Size([1, 64, 7, 7])
|
|
|
|
|
|
head = FeatureHead(temporal_type=None, average_clips='score')
|
|
|
feat = head(x, num_segs=2)
|
|
|
assert feat.shape == torch.Size([1, 2, 64])
|
|
|
|
|
|
def test_tsm_backbone(self):
|
|
|
head = FeatureHead(backbone_name='tsm', average_clips='score')
|
|
|
x = torch.rand(2, 64, 7, 7)
|
|
|
with pytest.raises(AssertionError):
|
|
|
feat = head(x)
|
|
|
with pytest.raises(AssertionError):
|
|
|
feat = head(x, num_segs=2)
|
|
|
|
|
|
head = FeatureHead(num_segments=2, average_clips='score')
|
|
|
feat = head(x, num_segs=2)
|
|
|
assert feat.shape == torch.Size([1, 64])
|
|
|
|
|
|
x = torch.rand(2, 64, 7, 7)
|
|
|
head = FeatureHead(
|
|
|
num_segments=2, spatial_type=None, average_clips='score')
|
|
|
feat = head(x, num_segs=2)
|
|
|
assert feat.shape == torch.Size([1, 64, 7, 7])
|
|
|
|
|
|
def test_gcn_backbone(self):
|
|
|
|
|
|
head = FeatureHead(backbone_name='gcn', average_clips='score')
|
|
|
x = torch.rand(1, 5, 64, 2, 7)
|
|
|
feat = head(x)
|
|
|
assert feat.shape == torch.Size([1, 64])
|
|
|
|
