|
|
|
|
|
import gc
|
|
|
from typing import Dict, List, Union
|
|
|
|
|
|
from mmengine.runner import EpochBasedTrainLoop
|
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
from mmaction.registry import LOOPS
|
|
|
|
|
|
|
|
|
class EpochMultiLoader:
|
|
|
"""Multi loaders based on epoch."""
|
|
|
|
|
|
def __init__(self, dataloaders: List[DataLoader]):
|
|
|
self._dataloaders = dataloaders
|
|
|
self.iter_loaders = [iter(loader) for loader in self._dataloaders]
|
|
|
|
|
|
@property
|
|
|
def num_loaders(self):
|
|
|
"""The number of dataloaders."""
|
|
|
return len(self._dataloaders)
|
|
|
|
|
|
def __iter__(self):
|
|
|
"""Return self when executing __iter__."""
|
|
|
return self
|
|
|
|
|
|
def __next__(self):
|
|
|
"""Get the next iter's data of multiple loaders."""
|
|
|
data = tuple([next(loader) for loader in self.iter_loaders])
|
|
|
return data
|
|
|
|
|
|
def __len__(self):
|
|
|
"""Get the length of loader."""
|
|
|
return min([len(loader) for loader in self._dataloaders])
|
|
|
|
|
|
|
|
|
@LOOPS.register_module()
|
|
|
class MultiLoaderEpochBasedTrainLoop(EpochBasedTrainLoop):
|
|
|
"""EpochBasedTrainLoop with multiple dataloaders.
|
|
|
|
|
|
Args:
|
|
|
runner (Runner): A reference of runner.
|
|
|
dataloader (Dataloader or Dict): A dataloader object or a dict to
|
|
|
build a dataloader for training the model.
|
|
|
other_loaders (List of Dataloader or Dict): A list of other loaders.
|
|
|
Each item in the list is a dataloader object or a dict to build
|
|
|
a dataloader.
|
|
|
max_epochs (int): Total training epochs.
|
|
|
val_begin (int): The epoch that begins validating. Defaults to 1.
|
|
|
val_interval (int): Validation interval. Defaults to 1.
|
|
|
"""
|
|
|
|
|
|
def __init__(self,
|
|
|
runner,
|
|
|
dataloader: Union[Dict, DataLoader],
|
|
|
other_loaders: List[Union[Dict, DataLoader]],
|
|
|
max_epochs: int,
|
|
|
val_begin: int = 1,
|
|
|
val_interval: int = 1) -> None:
|
|
|
super().__init__(runner, dataloader, max_epochs, val_begin,
|
|
|
val_interval)
|
|
|
multi_loaders = [self.dataloader]
|
|
|
for loader in other_loaders:
|
|
|
if isinstance(loader, dict):
|
|
|
loader = runner.build_dataloader(loader, seed=runner.seed)
|
|
|
multi_loaders.append(loader)
|
|
|
|
|
|
self.multi_loaders = multi_loaders
|
|
|
|
|
|
def run_epoch(self) -> None:
|
|
|
"""Iterate one epoch."""
|
|
|
self.runner.call_hook('before_train_epoch')
|
|
|
self.runner.model.train()
|
|
|
|
|
|
gc.collect()
|
|
|
for loader in self.multi_loaders:
|
|
|
if hasattr(loader, 'sampler') and hasattr(loader.sampler,
|
|
|
'set_epoch'):
|
|
|
loader.sampler.set_epoch(self._epoch)
|
|
|
|
|
|
for idx, data_batch in enumerate(EpochMultiLoader(self.multi_loaders)):
|
|
|
self.run_iter(idx, data_batch)
|
|
|
|
|
|
self.runner.call_hook('after_train_epoch')
|
|
|
self._epoch += 1
|
|
|
|
