mmpose-0.29.0/tests/test_backbones/test_tcn.py

# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import pytest
import torch
import torch.nn as nn

from mmpose.models.backbones import TCN
from mmpose.models.backbones.tcn import BasicTemporalBlock


def test_basic_temporal_block():
 with pytest.raises(AssertionError):
 # padding( + shift) should not be larger than x.shape[2]
 block = BasicTemporalBlock(1024, 1024, dilation=81)
 x = torch.rand(2, 1024, 150)
 x_out = block(x)

 with pytest.raises(AssertionError):
 # when use_stride_conv is True, shift + kernel_size // 2 should
 # not be larger than x.shape[2]
 block = BasicTemporalBlock(
 1024, 1024, kernel_size=5, causal=True, use_stride_conv=True)
 x = torch.rand(2, 1024, 3)
 x_out = block(x)

 # BasicTemporalBlock with causal == False
 block = BasicTemporalBlock(1024, 1024)
 x = torch.rand(2, 1024, 241)
 x_out = block(x)
 assert x_out.shape == torch.Size([2, 1024, 235])

 # BasicTemporalBlock with causal == True
 block = BasicTemporalBlock(1024, 1024, causal=True)
 x = torch.rand(2, 1024, 241)
 x_out = block(x)
 assert x_out.shape == torch.Size([2, 1024, 235])

 # BasicTemporalBlock with residual == False
 block = BasicTemporalBlock(1024, 1024, residual=False)
 x = torch.rand(2, 1024, 241)
 x_out = block(x)
 assert x_out.shape == torch.Size([2, 1024, 235])

 # BasicTemporalBlock, use_stride_conv == True
 block = BasicTemporalBlock(1024, 1024, use_stride_conv=True)
 x = torch.rand(2, 1024, 81)
 x_out = block(x)
 assert x_out.shape == torch.Size([2, 1024, 27])

 # BasicTemporalBlock with use_stride_conv == True and causal == True
 block = BasicTemporalBlock(1024, 1024, use_stride_conv=True, causal=True)
 x = torch.rand(2, 1024, 81)
 x_out = block(x)
 assert x_out.shape == torch.Size([2, 1024, 27])


def test_tcn_backbone():
 with pytest.raises(AssertionError):
 # num_blocks should equal len(kernel_sizes) - 1
 TCN(in_channels=34, num_blocks=3, kernel_sizes=(3, 3, 3))

 with pytest.raises(AssertionError):
 # kernel size should be odd
 TCN(in_channels=34, kernel_sizes=(3, 4, 3))

 # Test TCN with 2 blocks (use_stride_conv == False)
 model = TCN(in_channels=34, num_blocks=2, kernel_sizes=(3, 3, 3))
 pose2d = torch.rand((2, 34, 243))
 feat = model(pose2d)
 assert len(feat) == 2
 assert feat[0].shape == (2, 1024, 235)
 assert feat[1].shape == (2, 1024, 217)

 # Test TCN with 4 blocks and weight norm clip
 max_norm = 0.1
 model = TCN(
 in_channels=34,
 num_blocks=4,
 kernel_sizes=(3, 3, 3, 3, 3),
 max_norm=max_norm)
 pose2d = torch.rand((2, 34, 243))
 feat = model(pose2d)
 assert len(feat) == 4
 assert feat[0].shape == (2, 1024, 235)
 assert feat[1].shape == (2, 1024, 217)
 assert feat[2].shape == (2, 1024, 163)
 assert feat[3].shape == (2, 1024, 1)

 for module in model.modules():
 if isinstance(module, torch.nn.modules.conv._ConvNd):
 norm = module.weight.norm().item()
 np.testing.assert_allclose(
 np.maximum(norm, max_norm), max_norm, rtol=1e-4)

 # Test TCN with 4 blocks (use_stride_conv == True)
 model = TCN(
 in_channels=34,
 num_blocks=4,
 kernel_sizes=(3, 3, 3, 3, 3),
 use_stride_conv=True)
 pose2d = torch.rand((2, 34, 243))
 feat = model(pose2d)
 assert len(feat) == 4
 assert feat[0].shape == (2, 1024, 27)
 assert feat[1].shape == (2, 1024, 9)
 assert feat[2].shape == (2, 1024, 3)
 assert feat[3].shape == (2, 1024, 1)

 # Check that the model w. or w/o use_stride_conv will have the same
 # output and gradient after a forward+backward pass
 model1 = TCN(
 in_channels=34,
 stem_channels=4,
 num_blocks=1,
 kernel_sizes=(3, 3),
 dropout=0,
 residual=False,
 norm_cfg=None)
 model2 = TCN(
 in_channels=34,
 stem_channels=4,
 num_blocks=1,
 kernel_sizes=(3, 3),
 dropout=0,
 residual=False,
 norm_cfg=None,
 use_stride_conv=True)
 for m in model1.modules():
 if isinstance(m, nn.Conv1d):
 nn.init.constant_(m.weight, 0.5)
 if m.bias is not None:
 nn.init.constant_(m.bias, 0)
 for m in model2.modules():
 if isinstance(m, nn.Conv1d):
 nn.init.constant_(m.weight, 0.5)
 if m.bias is not None:
 nn.init.constant_(m.bias, 0)
 input1 = torch.rand((1, 34, 9))
 input2 = input1.clone()
 outputs1 = model1(input1)
 outputs2 = model2(input2)
 for output1, output2 in zip(outputs1, outputs2):
 assert torch.isclose(output1, output2).all()

 criterion = nn.MSELoss()
 target = torch.rand(output1.shape)
 loss1 = criterion(output1, target)
 loss2 = criterion(output2, target)
 loss1.backward()
 loss2.backward()
 for m1, m2 in zip(model1.modules(), model2.modules()):
 if isinstance(m1, nn.Conv1d):
 assert torch.isclose(m1.weight.grad, m2.weight.grad).all()