Spaces:
Sleeping
Sleeping
| # 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 | |
| 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') | |