|
|
|
|
|
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() |
|
|
|
