File size: 1,030 Bytes
3bbb319 |
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 43 44 |
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import pytest
import torch
from mmpose.models.necks import GlobalAveragePooling
def test_gap():
"""Test GlobalAveragePooling neck."""
gap = GlobalAveragePooling()
with pytest.raises(TypeError):
gap(1)
x0_shape = (32, 1024, 4, 4)
x1_shape = (32, 2048, 2, 2)
x0 = _demo_inputs(x0_shape)
x1 = _demo_inputs(x1_shape)
y = gap(x0)
assert y.shape == torch.Size([32, 1024])
y = gap([x0, x1])
assert y[0].shape == torch.Size([32, 1024])
assert y[1].shape == torch.Size([32, 2048])
y = gap((x0, x1))
assert y[0].shape == torch.Size([32, 1024])
assert y[1].shape == torch.Size([32, 2048])
def _demo_inputs(input_shape=(1, 3, 64, 64)):
"""Create a superset of inputs needed to run backbone.
Args:
input_shape (tuple): input batch dimensions.
Default: (1, 3, 64, 64).
"""
imgs = np.random.random(input_shape)
imgs = torch.FloatTensor(imgs)
return imgs
|