mirror of https://github.com/open-mmlab/mmcv.git
59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from mmcv.cnn.bricks import ContextBlock
|
|
|
|
|
|
def test_context_block():
|
|
with pytest.raises(AssertionError):
|
|
# pooling_type should be in ['att', 'avg']
|
|
ContextBlock(16, 1. / 4, pooling_type='unsupport_type')
|
|
|
|
with pytest.raises(AssertionError):
|
|
# fusion_types should be of type list or tuple
|
|
ContextBlock(16, 1. / 4, fusion_types='unsupport_type')
|
|
|
|
with pytest.raises(AssertionError):
|
|
# fusion_types should be in ['channel_add', 'channel_mul']
|
|
ContextBlock(16, 1. / 4, fusion_types=('unsupport_type', ))
|
|
|
|
# test pooling_type='att'
|
|
imgs = torch.randn(2, 16, 20, 20)
|
|
context_block = ContextBlock(16, 1. / 4, pooling_type='att')
|
|
out = context_block(imgs)
|
|
assert context_block.conv_mask.in_channels == 16
|
|
assert context_block.conv_mask.out_channels == 1
|
|
assert out.shape == imgs.shape
|
|
|
|
# test pooling_type='avg'
|
|
imgs = torch.randn(2, 16, 20, 20)
|
|
context_block = ContextBlock(16, 1. / 4, pooling_type='avg')
|
|
out = context_block(imgs)
|
|
assert hasattr(context_block, 'avg_pool')
|
|
assert out.shape == imgs.shape
|
|
|
|
# test fusion_types=('channel_add',)
|
|
imgs = torch.randn(2, 16, 20, 20)
|
|
context_block = ContextBlock(16, 1. / 4, fusion_types=('channel_add', ))
|
|
out = context_block(imgs)
|
|
assert context_block.channel_add_conv is not None
|
|
assert context_block.channel_mul_conv is None
|
|
assert out.shape == imgs.shape
|
|
|
|
# test fusion_types=('channel_mul',)
|
|
imgs = torch.randn(2, 16, 20, 20)
|
|
context_block = ContextBlock(16, 1. / 4, fusion_types=('channel_mul', ))
|
|
out = context_block(imgs)
|
|
assert context_block.channel_add_conv is None
|
|
assert context_block.channel_mul_conv is not None
|
|
assert out.shape == imgs.shape
|
|
|
|
# test fusion_types=('channel_add', 'channel_mul')
|
|
imgs = torch.randn(2, 16, 20, 20)
|
|
context_block = ContextBlock(
|
|
16, 1. / 4, fusion_types=('channel_add', 'channel_mul'))
|
|
out = context_block(imgs)
|
|
assert context_block.channel_add_conv is not None
|
|
assert context_block.channel_mul_conv is not None
|
|
assert out.shape == imgs.shape
|