Spaces:
Running
on
L40S
Running
on
L40S
File size: 2,204 Bytes
d7e58f0 |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# Copyright (c) OpenMMLab. All rights reserved.
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
|