mirror of https://github.com/open-mmlab/mmcv.git
127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import logging
|
|
import os.path as osp
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import Dataset
|
|
|
|
from mmcv.runner import build_runner
|
|
from mmcv.utils import IS_IPU_AVAILABLE
|
|
|
|
if IS_IPU_AVAILABLE:
|
|
from mmcv.device.ipu import IPUDataLoader, runner
|
|
|
|
skip_no_ipu = pytest.mark.skipif(
|
|
not IS_IPU_AVAILABLE, reason='test case under ipu environment')
|
|
|
|
# Most of its functions are inherited from EpochBasedRunner and IterBasedRunner
|
|
# So only do incremental testing on overridden methods
|
|
# Comparing with base runner,
|
|
# Overridden functions are listed below:
|
|
# __init__, register_lr_hook, register_optimizer_hook
|
|
# register_lr_hook and register_optimizer_hook are tested in test_runner.py
|
|
|
|
|
|
class OldStyleModel(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(3, 3, 1)
|
|
|
|
|
|
class Model(OldStyleModel):
|
|
|
|
def train_step(self):
|
|
pass
|
|
|
|
def val_step(self):
|
|
pass
|
|
|
|
|
|
class ToyModel(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(3, 3, 1)
|
|
self.bn = nn.BatchNorm2d(3)
|
|
self.relu = nn.ReLU6()
|
|
self.fp16_enabled = False
|
|
|
|
def forward(self, img, return_loss=True, **kwargs):
|
|
x = self.conv(img)
|
|
x = self.bn(x)
|
|
x = self.relu(x)
|
|
if return_loss:
|
|
loss = ((x - kwargs['gt_label'])**2).sum()
|
|
return {'loss': loss, 'loss1': loss + 1}
|
|
return x
|
|
|
|
def _parse_losses(self, losses):
|
|
return losses['loss'], {'loss1': losses['loss']}
|
|
|
|
def train_step(self, data, optimizer=None, **kwargs):
|
|
losses = self(**data)
|
|
loss, log_vars = self._parse_losses(losses)
|
|
outputs = dict(
|
|
loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))
|
|
return outputs
|
|
|
|
|
|
class ToyDataset(Dataset):
|
|
|
|
def __getitem__(self, index):
|
|
return {
|
|
'img': torch.rand((3, 10, 10)),
|
|
'gt_label': torch.rand((3, 10, 10))
|
|
}
|
|
|
|
def __len__(self, ):
|
|
return 3
|
|
|
|
|
|
@skip_no_ipu
|
|
def test_build_runner(tmp_path):
|
|
# __init__
|
|
dir_name = 'a_tmp_dir'
|
|
|
|
default_args = dict(
|
|
model=Model(),
|
|
work_dir=osp.join(tmp_path, dir_name),
|
|
logger=logging.getLogger())
|
|
cfg = dict(type='IPUEpochBasedRunner', max_epochs=1)
|
|
ipu_runner = build_runner(cfg, default_args=default_args)
|
|
assert ipu_runner._max_epochs == 1
|
|
cfg = dict(type='IPUIterBasedRunner', max_iters=1)
|
|
ipu_runner = build_runner(cfg, default_args=default_args)
|
|
assert ipu_runner._max_iters == 1
|
|
|
|
runner.IS_IPU_AVAILABLE = False
|
|
cfg = dict(type='IPUIterBasedRunner', max_iters=1)
|
|
with pytest.raises(
|
|
NotImplementedError,
|
|
match='cpu mode on IPURunner is not supported'):
|
|
ipu_runner = build_runner(cfg, default_args=default_args)
|
|
|
|
runner.IS_IPU_AVAILABLE = True
|
|
with pytest.raises(ValueError, match='Only one of'):
|
|
cfg = dict(type='IPUIterBasedRunner', max_epochs=1, max_iters=1)
|
|
ipu_runner = build_runner(cfg, default_args=default_args)
|
|
|
|
model = ToyModel()
|
|
options_cfg = {'train_cfg': {}, 'eval_cfg': {}}
|
|
dataloader = IPUDataLoader(ToyDataset(), None, num_workers=1)
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
|
|
cfg = dict(type='IPUIterBasedRunner', max_iters=2, options_cfg=options_cfg)
|
|
default_args = dict(
|
|
model=model,
|
|
optimizer=optimizer,
|
|
work_dir=osp.join(tmp_path, dir_name),
|
|
logger=logging.getLogger())
|
|
ipu_runner = build_runner(cfg, default_args=default_args)
|
|
ipu_runner.run([dataloader], [('train', 2)])
|
|
ipu_runner.get_options('val')
|
|
with pytest.raises(ValueError, match='mode should be train or val'):
|
|
ipu_runner.get_options('666')
|