236 lines
6.4 KiB
Python
236 lines
6.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
"""pytest tests/test_forward.py."""
|
|
import copy
|
|
from os.path import dirname, exists, join
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from mmcv.cnn.utils import revert_sync_batchnorm
|
|
|
|
|
|
def _demo_mm_inputs(input_shape=(2, 3, 8, 16), num_classes=10):
|
|
"""Create a superset of inputs needed to run test or train batches.
|
|
|
|
Args:
|
|
input_shape (tuple):
|
|
input batch dimensions
|
|
|
|
num_classes (int):
|
|
number of semantic classes
|
|
"""
|
|
(N, C, H, W) = input_shape
|
|
|
|
rng = np.random.RandomState(0)
|
|
|
|
imgs = rng.rand(*input_shape)
|
|
segs = rng.randint(
|
|
low=0, high=num_classes - 1, size=(N, 1, H, W)).astype(np.uint8)
|
|
|
|
img_metas = [{
|
|
'img_shape': (H, W, C),
|
|
'ori_shape': (H, W, C),
|
|
'pad_shape': (H, W, C),
|
|
'filename': '<demo>.png',
|
|
'scale_factor': 1.0,
|
|
'flip': False,
|
|
'flip_direction': 'horizontal'
|
|
} for _ in range(N)]
|
|
|
|
mm_inputs = {
|
|
'imgs': torch.FloatTensor(imgs),
|
|
'img_metas': img_metas,
|
|
'gt_semantic_seg': torch.LongTensor(segs)
|
|
}
|
|
return mm_inputs
|
|
|
|
|
|
def _get_config_directory():
|
|
"""Find the predefined segmentor config directory."""
|
|
try:
|
|
# Assume we are running in the source mmsegmentation repo
|
|
repo_dpath = dirname(dirname(dirname(__file__)))
|
|
except NameError:
|
|
# For IPython development when this __file__ is not defined
|
|
import mmseg
|
|
repo_dpath = dirname(dirname(dirname(mmseg.__file__)))
|
|
config_dpath = join(repo_dpath, 'configs')
|
|
if not exists(config_dpath):
|
|
raise Exception('Cannot find config path')
|
|
return config_dpath
|
|
|
|
|
|
def _get_config_module(fname):
|
|
"""Load a configuration as a python module."""
|
|
from mmcv import Config
|
|
config_dpath = _get_config_directory()
|
|
config_fpath = join(config_dpath, fname)
|
|
config_mod = Config.fromfile(config_fpath)
|
|
return config_mod
|
|
|
|
|
|
def _get_segmentor_cfg(fname):
|
|
"""Grab configs necessary to create a segmentor.
|
|
|
|
These are deep copied to allow for safe modification of parameters without
|
|
influencing other tests.
|
|
"""
|
|
config = _get_config_module(fname)
|
|
model = copy.deepcopy(config.model)
|
|
return model
|
|
|
|
|
|
def test_pspnet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_fcn_forward():
|
|
_test_encoder_decoder_forward('fcn/fcn_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_deeplabv3_forward():
|
|
_test_encoder_decoder_forward(
|
|
'deeplabv3/deeplabv3_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_deeplabv3plus_forward():
|
|
_test_encoder_decoder_forward(
|
|
'deeplabv3plus/deeplabv3plus_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_gcnet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'gcnet/gcnet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_ann_forward():
|
|
_test_encoder_decoder_forward('ann/ann_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_ccnet_forward():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip('CCNet requires CUDA')
|
|
_test_encoder_decoder_forward(
|
|
'ccnet/ccnet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_danet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'danet/danet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_nonlocal_net_forward():
|
|
_test_encoder_decoder_forward(
|
|
'nonlocal_net/nonlocal_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_upernet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'upernet/upernet_r50_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_hrnet_forward():
|
|
_test_encoder_decoder_forward('hrnet/fcn_hr18s_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_ocrnet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'ocrnet/ocrnet_hr18s_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_psanet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'psanet/psanet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_encnet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'encnet/encnet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_sem_fpn_forward():
|
|
_test_encoder_decoder_forward('sem_fpn/fpn_r50_512x1024_80k_cityscapes.py')
|
|
|
|
|
|
def test_point_rend_forward():
|
|
_test_encoder_decoder_forward(
|
|
'point_rend/pointrend_r50_512x1024_80k_cityscapes.py')
|
|
|
|
|
|
def test_mobilenet_v2_forward():
|
|
_test_encoder_decoder_forward(
|
|
'mobilenet_v2/pspnet_m-v2-d8_512x1024_80k_cityscapes.py')
|
|
|
|
|
|
def test_dnlnet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'dnlnet/dnl_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def test_emanet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'emanet/emanet_r50-d8_512x1024_80k_cityscapes.py')
|
|
|
|
|
|
def test_isanet_forward():
|
|
_test_encoder_decoder_forward(
|
|
'isanet/isanet_r50-d8_512x1024_40k_cityscapes.py')
|
|
|
|
|
|
def get_world_size(process_group):
|
|
|
|
return 1
|
|
|
|
|
|
def _check_input_dim(self, inputs):
|
|
pass
|
|
|
|
|
|
@patch('torch.nn.modules.batchnorm._BatchNorm._check_input_dim',
|
|
_check_input_dim)
|
|
@patch('torch.distributed.get_world_size', get_world_size)
|
|
def _test_encoder_decoder_forward(cfg_file):
|
|
model = _get_segmentor_cfg(cfg_file)
|
|
model['pretrained'] = None
|
|
model['test_cfg']['mode'] = 'whole'
|
|
|
|
from mmseg.models import build_segmentor
|
|
segmentor = build_segmentor(model)
|
|
segmentor.init_weights()
|
|
|
|
if isinstance(segmentor.decode_head, nn.ModuleList):
|
|
num_classes = segmentor.decode_head[-1].num_classes
|
|
else:
|
|
num_classes = segmentor.decode_head.num_classes
|
|
# batch_size=2 for BatchNorm
|
|
input_shape = (2, 3, 32, 32)
|
|
mm_inputs = _demo_mm_inputs(input_shape, num_classes=num_classes)
|
|
|
|
imgs = mm_inputs.pop('imgs')
|
|
img_metas = mm_inputs.pop('img_metas')
|
|
gt_semantic_seg = mm_inputs['gt_semantic_seg']
|
|
|
|
# convert to cuda Tensor if applicable
|
|
if torch.cuda.is_available():
|
|
segmentor = segmentor.cuda()
|
|
imgs = imgs.cuda()
|
|
gt_semantic_seg = gt_semantic_seg.cuda()
|
|
else:
|
|
segmentor = revert_sync_batchnorm(segmentor)
|
|
|
|
# Test forward train
|
|
losses = segmentor.forward(
|
|
imgs, img_metas, gt_semantic_seg=gt_semantic_seg, return_loss=True)
|
|
assert isinstance(losses, dict)
|
|
|
|
# Test forward test
|
|
with torch.no_grad():
|
|
segmentor.eval()
|
|
# pack into lists
|
|
img_list = [img[None, :] for img in imgs]
|
|
img_meta_list = [[img_meta] for img_meta in img_metas]
|
|
segmentor.forward(img_list, img_meta_list, return_loss=False)
|