mirror of
https://github.com/open-mmlab/mmclassification.git
synced 2025-06-03 21:53:55 +08:00
1a7cebe4b9
* Refactor unit tests folder structure. * Remove label smooth and Vit test in `test_classifiers.py` * Rename test_utils in dataset to test_dataset_utils * Split test_models/test_utils/test_utils.py to multiple sub files. * Add unit tests of classifiers and heads * Use patch context manager. * Add unit test of `is_tracing`, and add warning in `is_tracing` if torch verison is smaller than 1.6.0
39 lines
967 B
Python
39 lines
967 B
Python
import pytest
|
|
import torch
|
|
|
|
from mmcls.models.necks import GlobalAveragePooling
|
|
|
|
|
|
def test_gap_neck():
|
|
|
|
# test 1d gap_neck
|
|
neck = GlobalAveragePooling(dim=1)
|
|
# batch_size, num_features, feature_size
|
|
fake_input = torch.rand(1, 16, 24)
|
|
|
|
output = neck(fake_input)
|
|
# batch_size, num_features
|
|
assert output.shape == (1, 16)
|
|
|
|
# test 1d gap_neck
|
|
neck = GlobalAveragePooling(dim=2)
|
|
# batch_size, num_features, feature_size(2)
|
|
fake_input = torch.rand(1, 16, 24, 24)
|
|
|
|
output = neck(fake_input)
|
|
# batch_size, num_features
|
|
assert output.shape == (1, 16)
|
|
|
|
# test 1d gap_neck
|
|
neck = GlobalAveragePooling(dim=3)
|
|
# batch_size, num_features, feature_size(3)
|
|
fake_input = torch.rand(1, 16, 24, 24, 5)
|
|
|
|
output = neck(fake_input)
|
|
# batch_size, num_features
|
|
assert output.shape == (1, 16)
|
|
|
|
with pytest.raises(AssertionError):
|
|
# dim must in [1, 2, 3]
|
|
GlobalAveragePooling(dim='other')
|