| |
| import tempfile |
| from unittest import TestCase |
| from unittest.mock import Mock |
|
|
| import torch |
| import torch.nn as nn |
| from mmengine.evaluator import Evaluator |
| from mmengine.model import BaseModel |
| from mmengine.optim import OptimWrapper |
| from mmengine.runner import Runner |
| from torch.utils.data import Dataset |
|
|
| from mmdet.registry import DATASETS |
| from mmdet.utils import register_all_modules |
|
|
| register_all_modules() |
|
|
|
|
| class ToyModel(nn.Module): |
|
|
| def __init__(self): |
| super().__init__() |
| self.linear = nn.Linear(2, 1) |
|
|
| def forward(self, inputs, data_samples, mode='tensor'): |
| labels = torch.stack(data_samples) |
| inputs = torch.stack(inputs) |
| outputs = self.linear(inputs) |
| if mode == 'tensor': |
| return outputs |
| elif mode == 'loss': |
| loss = (labels - outputs).sum() |
| outputs = dict(loss=loss) |
| return outputs |
| else: |
| return outputs |
|
|
|
|
| class ToyModel1(BaseModel, ToyModel): |
|
|
| def __init__(self): |
| super().__init__() |
|
|
| def forward(self, *args, **kwargs): |
| return super(BaseModel, self).forward(*args, **kwargs) |
|
|
|
|
| class ToyModel2(BaseModel): |
|
|
| def __init__(self): |
| super().__init__() |
| self.teacher = ToyModel1() |
| self.student = ToyModel1() |
| self.semi_test_cfg = dict(predict_on='teacher') |
|
|
| def forward(self, *args, **kwargs): |
| return self.student(*args, **kwargs) |
|
|
|
|
| @DATASETS.register_module(force=True) |
| class DummyDataset(Dataset): |
| METAINFO = dict() |
| data = torch.randn(12, 2) |
| label = torch.ones(12) |
|
|
| @property |
| def metainfo(self): |
| return self.METAINFO |
|
|
| def __len__(self): |
| return self.data.size(0) |
|
|
| def __getitem__(self, index): |
| return dict(inputs=self.data[index], data_samples=self.label[index]) |
|
|
|
|
| class TestTeacherStudentValLoop(TestCase): |
|
|
| def setUp(self): |
| self.temp_dir = tempfile.TemporaryDirectory() |
|
|
| def tearDown(self): |
| self.temp_dir.cleanup() |
|
|
| def test_teacher_student_val_loop(self): |
| device = 'cuda:0' if torch.cuda.is_available() else 'cpu' |
| model = ToyModel2().to(device) |
| evaluator = Mock() |
| evaluator.evaluate = Mock(return_value=dict(acc=0.5)) |
| evaluator.__class__ = Evaluator |
| runner = Runner( |
| model=model, |
| train_dataloader=dict( |
| dataset=dict(type='DummyDataset'), |
| sampler=dict(type='DefaultSampler', shuffle=True), |
| batch_size=3, |
| num_workers=0), |
| val_dataloader=dict( |
| dataset=dict(type='DummyDataset'), |
| sampler=dict(type='DefaultSampler', shuffle=False), |
| batch_size=3, |
| num_workers=0), |
| val_evaluator=evaluator, |
| work_dir=self.temp_dir.name, |
| default_scope='mmdet', |
| optim_wrapper=OptimWrapper( |
| torch.optim.Adam(ToyModel().parameters())), |
| train_cfg=dict(by_epoch=True, max_epochs=2, val_interval=1), |
| val_cfg=dict(type='TeacherStudentValLoop'), |
| default_hooks=dict(logger=dict(type='LoggerHook', interval=1)), |
| experiment_name='test1') |
| runner.train() |
|
|
