Spaces:
Sleeping
Sleeping
File size: 3,826 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
# 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')
|