|
- # Copyright (c) OpenMMLab. All rights reserved.
- import logging
- import tempfile
- from unittest.mock import MagicMock, patch
-
- import mmcv.runner
- import pytest
- import torch
- import torch.nn as nn
- from mmcv.runner import obj_from_dict
- from torch.utils.data import DataLoader, Dataset
-
- from mmseg.apis import single_gpu_test
- from mmseg.core import DistEvalHook, EvalHook
-
-
- class ExampleDataset(Dataset):
-
- def __getitem__(self, idx):
- results = dict(img=torch.tensor([1]), img_metas=dict())
- return results
-
- def __len__(self):
- return 1
-
-
- class ExampleModel(nn.Module):
-
- def __init__(self):
- super(ExampleModel, self).__init__()
- self.test_cfg = None
- self.conv = nn.Conv2d(3, 3, 3)
-
- def forward(self, img, img_metas, test_mode=False, **kwargs):
- return img
-
- def train_step(self, data_batch, optimizer):
- loss = self.forward(**data_batch)
- return dict(loss=loss)
-
-
- def test_iter_eval_hook():
- with pytest.raises(TypeError):
- test_dataset = ExampleModel()
- data_loader = [
- DataLoader(
- test_dataset,
- batch_size=1,
- sampler=None,
- num_worker=0,
- shuffle=False)
- ]
- EvalHook(data_loader)
-
- test_dataset = ExampleDataset()
- test_dataset.pre_eval = MagicMock(return_value=[torch.tensor([1])])
- test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
- loader = DataLoader(test_dataset, batch_size=1)
- model = ExampleModel()
- data_loader = DataLoader(
- test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
- optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
- optimizer = obj_from_dict(optim_cfg, torch.optim,
- dict(params=model.parameters()))
-
- # test EvalHook
- with tempfile.TemporaryDirectory() as tmpdir:
- eval_hook = EvalHook(data_loader, by_epoch=False, efficient_test=True)
- runner = mmcv.runner.IterBasedRunner(
- model=model,
- optimizer=optimizer,
- work_dir=tmpdir,
- logger=logging.getLogger())
- runner.register_hook(eval_hook)
- runner.run([loader], [('train', 1)], 1)
- test_dataset.evaluate.assert_called_with([torch.tensor([1])],
- logger=runner.logger)
-
-
- def test_epoch_eval_hook():
- with pytest.raises(TypeError):
- test_dataset = ExampleModel()
- data_loader = [
- DataLoader(
- test_dataset,
- batch_size=1,
- sampler=None,
- num_worker=0,
- shuffle=False)
- ]
- EvalHook(data_loader, by_epoch=True)
-
- test_dataset = ExampleDataset()
- test_dataset.pre_eval = MagicMock(return_value=[torch.tensor([1])])
- test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
- loader = DataLoader(test_dataset, batch_size=1)
- model = ExampleModel()
- data_loader = DataLoader(
- test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
- optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
- optimizer = obj_from_dict(optim_cfg, torch.optim,
- dict(params=model.parameters()))
-
- # test EvalHook with interval
- with tempfile.TemporaryDirectory() as tmpdir:
- eval_hook = EvalHook(data_loader, by_epoch=True, interval=2)
- runner = mmcv.runner.EpochBasedRunner(
- model=model,
- optimizer=optimizer,
- work_dir=tmpdir,
- logger=logging.getLogger())
- runner.register_hook(eval_hook)
- runner.run([loader], [('train', 1)], 2)
- test_dataset.evaluate.assert_called_once_with([torch.tensor([1])],
- logger=runner.logger)
-
-
- def multi_gpu_test(model,
- data_loader,
- tmpdir=None,
- gpu_collect=False,
- pre_eval=False):
- # Pre eval is set by default when training.
- results = single_gpu_test(model, data_loader, pre_eval=True)
- return results
-
-
- @patch('mmseg.apis.multi_gpu_test', multi_gpu_test)
- def test_dist_eval_hook():
- with pytest.raises(TypeError):
- test_dataset = ExampleModel()
- data_loader = [
- DataLoader(
- test_dataset,
- batch_size=1,
- sampler=None,
- num_worker=0,
- shuffle=False)
- ]
- DistEvalHook(data_loader)
-
- test_dataset = ExampleDataset()
- test_dataset.pre_eval = MagicMock(return_value=[torch.tensor([1])])
- test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
- loader = DataLoader(test_dataset, batch_size=1)
- model = ExampleModel()
- data_loader = DataLoader(
- test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
- optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
- optimizer = obj_from_dict(optim_cfg, torch.optim,
- dict(params=model.parameters()))
-
- # test DistEvalHook
- with tempfile.TemporaryDirectory() as tmpdir:
- eval_hook = DistEvalHook(
- data_loader, by_epoch=False, efficient_test=True)
- runner = mmcv.runner.IterBasedRunner(
- model=model,
- optimizer=optimizer,
- work_dir=tmpdir,
- logger=logging.getLogger())
- runner.register_hook(eval_hook)
- runner.run([loader], [('train', 1)], 1)
- test_dataset.evaluate.assert_called_with([torch.tensor([1])],
- logger=runner.logger)
-
-
- @patch('mmseg.apis.multi_gpu_test', multi_gpu_test)
- def test_dist_eval_hook_epoch():
- with pytest.raises(TypeError):
- test_dataset = ExampleModel()
- data_loader = [
- DataLoader(
- test_dataset,
- batch_size=1,
- sampler=None,
- num_worker=0,
- shuffle=False)
- ]
- DistEvalHook(data_loader)
-
- test_dataset = ExampleDataset()
- test_dataset.pre_eval = MagicMock(return_value=[torch.tensor([1])])
- test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
- loader = DataLoader(test_dataset, batch_size=1)
- model = ExampleModel()
- data_loader = DataLoader(
- test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
- optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
- optimizer = obj_from_dict(optim_cfg, torch.optim,
- dict(params=model.parameters()))
-
- # test DistEvalHook
- with tempfile.TemporaryDirectory() as tmpdir:
- eval_hook = DistEvalHook(data_loader, by_epoch=True, interval=2)
- runner = mmcv.runner.EpochBasedRunner(
- model=model,
- optimizer=optimizer,
- work_dir=tmpdir,
- logger=logging.getLogger())
- runner.register_hook(eval_hook)
- runner.run([loader], [('train', 1)], 2)
- test_dataset.evaluate.assert_called_with([torch.tensor([1])],
- logger=runner.logger)
|