|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
from mmengine.testing import assert_params_all_zeros
|
|
|
|
|
|
from mmaction.models.common import (DividedSpatialAttentionWithNorm,
|
|
|
DividedTemporalAttentionWithNorm,
|
|
|
FFNWithNorm)
|
|
|
|
|
|
|
|
|
def test_divided_temporal_attention_with_norm():
|
|
|
_cfg = dict(embed_dims=768, num_heads=12, num_frames=8)
|
|
|
divided_temporal_attention = DividedTemporalAttentionWithNorm(**_cfg)
|
|
|
assert isinstance(divided_temporal_attention.norm, nn.LayerNorm)
|
|
|
assert assert_params_all_zeros(divided_temporal_attention.temporal_fc)
|
|
|
|
|
|
x = torch.rand(1, 1 + 8 * 14 * 14, 768)
|
|
|
output = divided_temporal_attention(x)
|
|
|
assert output.shape == torch.Size([1, 1 + 8 * 14 * 14, 768])
|
|
|
|
|
|
|
|
|
def test_divided_spatial_attention_with_norm():
|
|
|
_cfg = dict(embed_dims=512, num_heads=8, num_frames=4, dropout_layer=None)
|
|
|
divided_spatial_attention = DividedSpatialAttentionWithNorm(**_cfg)
|
|
|
assert isinstance(divided_spatial_attention.dropout_layer, nn.Identity)
|
|
|
assert isinstance(divided_spatial_attention.norm, nn.LayerNorm)
|
|
|
|
|
|
x = torch.rand(1, 1 + 4 * 14 * 14, 512)
|
|
|
output = divided_spatial_attention(x)
|
|
|
assert output.shape == torch.Size([1, 1 + 4 * 14 * 14, 512])
|
|
|
|
|
|
|
|
|
def test_ffn_with_norm():
|
|
|
_cfg = dict(
|
|
|
embed_dims=256, feedforward_channels=256 * 2, norm_cfg=dict(type='LN'))
|
|
|
ffn_with_norm = FFNWithNorm(**_cfg)
|
|
|
assert isinstance(ffn_with_norm.norm, nn.LayerNorm)
|
|
|
|
|
|
x = torch.rand(1, 1 + 4 * 14 * 14, 256)
|
|
|
output = ffn_with_norm(x)
|
|
|
assert output.shape == torch.Size([1, 1 + 4 * 14 * 14, 256])
|
|
|
|
