File size: 1,633 Bytes
d670799 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmaction.models import RGBPoseHead
def test_rgbpose_head():
"""Test RGBPoseHead."""
rgbpose_head = RGBPoseHead(
num_classes=4,
in_channels=[2048, 512],
dropout=dict(rgb=0.51, pose=0.49))
rgbpose_head.init_weights()
assert rgbpose_head.num_classes == 4
assert rgbpose_head.dropout == dict(rgb=0.51, pose=0.49)
assert rgbpose_head.in_channels == [2048, 512]
assert rgbpose_head.init_std == 0.01
assert isinstance(rgbpose_head.dropout_rgb, nn.Dropout)
assert isinstance(rgbpose_head.dropout_pose, nn.Dropout)
assert rgbpose_head.dropout_rgb.p == rgbpose_head.dropout['rgb']
assert rgbpose_head.dropout_pose.p == rgbpose_head.dropout['pose']
assert isinstance(rgbpose_head.fc_rgb, nn.Linear)
assert isinstance(rgbpose_head.fc_pose, nn.Linear)
assert rgbpose_head.fc_rgb.in_features == rgbpose_head.in_channels[0]
assert rgbpose_head.fc_rgb.out_features == rgbpose_head.num_classes
assert rgbpose_head.fc_pose.in_features == rgbpose_head.in_channels[1]
assert rgbpose_head.fc_pose.out_features == rgbpose_head.num_classes
assert isinstance(rgbpose_head.avg_pool, nn.AdaptiveAvgPool3d)
assert rgbpose_head.avg_pool.output_size == (1, 1, 1)
feat_rgb = torch.rand((2, 2048, 8, 7, 7))
feat_pose = torch.rand((2, 512, 32, 7, 7))
cls_scores = rgbpose_head((feat_rgb, feat_pose))
assert cls_scores['rgb'].shape == torch.Size([2, 4])
assert cls_scores['pose'].shape == torch.Size([2, 4])
|