49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
|
# Copyright (c) OpenMMLab. All rights reserved.
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from mmseg.models.backbones import ICNet
|
||
|
|
|
||
|
|
|
||
|
|
def test_icnet_backbone():
|
||
|
|
with pytest.raises(TypeError):
|
||
|
|
# Must give backbone dict in config file.
|
||
|
|
ICNet(
|
||
|
|
in_channels=3,
|
||
|
|
layer_channels=(512, 2048),
|
||
|
|
light_branch_middle_channels=32,
|
||
|
|
psp_out_channels=512,
|
||
|
|
out_channels=(64, 256, 256),
|
||
|
|
backbone_cfg=None)
|
||
|
|
|
||
|
|
# Test ICNet Standard Forward
|
||
|
|
model = ICNet(
|
||
|
|
backbone_cfg=dict(
|
||
|
|
type='ResNetV1c',
|
||
|
|
in_channels=3,
|
||
|
|
depth=50,
|
||
|
|
num_stages=4,
|
||
|
|
out_indices=(0, 1, 2, 3),
|
||
|
|
dilations=(1, 1, 2, 4),
|
||
|
|
strides=(1, 2, 1, 1),
|
||
|
|
norm_cfg=dict(type='BN', requires_grad=True),
|
||
|
|
norm_eval=False,
|
||
|
|
style='pytorch',
|
||
|
|
contract_dilation=True), )
|
||
|
|
assert hasattr(model.backbone,
|
||
|
|
'maxpool') and model.backbone.maxpool.ceil_mode is True
|
||
|
|
model.init_weights()
|
||
|
|
model.train()
|
||
|
|
batch_size = 2
|
||
|
|
imgs = torch.randn(batch_size, 3, 512, 1024)
|
||
|
|
feat = model(imgs)
|
||
|
|
|
||
|
|
assert model.psp_modules[0][0].output_size == 1
|
||
|
|
assert model.psp_modules[1][0].output_size == 2
|
||
|
|
assert model.psp_modules[2][0].output_size == 3
|
||
|
|
assert model.psp_bottleneck.padding == 1
|
||
|
|
assert model.conv_sub1[0].padding == 1
|
||
|
|
|
||
|
|
assert len(feat) == 3
|
||
|
|
assert feat[0].shape == torch.Size([batch_size, 64, 64, 128])
|