| from collections import OrderedDict |
| import torch |
| import torch.distributed as dist |
| from mmcv.runner import BaseModule |
| from typing import Dict, Tuple, List |
|
|
|
|
| def to_cpu(x: torch.Tensor) -> torch.Tensor: |
| """Move a tensor to CPU and detach it from the computation graph. |
| |
| Args: |
| x (torch.Tensor): The input tensor. |
| |
| Returns: |
| torch.Tensor: The tensor detached and moved to CPU. |
| """ |
| if isinstance(x, torch.Tensor): |
| return x.detach().cpu() |
| return x |
|
|
|
|
| class BaseArchitecture(BaseModule): |
| """Base class for mogen architecture. |
| |
| Args: |
| init_cfg (dict, optional): Initialization config for the module. |
| """ |
|
|
| def __init__(self, init_cfg: dict = None): |
| super(BaseArchitecture, self).__init__(init_cfg) |
|
|
| def forward_train(self, **kwargs): |
| """Forward computation during training.""" |
| pass |
|
|
| def forward_test(self, **kwargs): |
| """Forward computation during testing.""" |
| pass |
|
|
| def _parse_losses(self, losses: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, Dict[str, float]]: |
| """Parse the raw outputs (losses) of the network. |
| |
| Args: |
| losses (dict): Raw output of the network, which usually contains |
| losses and other necessary information. |
| |
| Returns: |
| tuple[Tensor, dict]: (loss, log_vars) |
| - loss is the loss tensor which may be a weighted sum of all losses, |
| - log_vars contains all the variables to be logged. |
| """ |
| log_vars = OrderedDict() |
| for loss_name, loss_value in losses.items(): |
| if isinstance(loss_value, torch.Tensor): |
| log_vars[loss_name] = loss_value.mean() |
| elif isinstance(loss_value, list): |
| log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value) |
| else: |
| raise TypeError(f'{loss_name} is not a tensor or list of tensors') |
|
|
| loss = sum(_value for _key, _value in log_vars.items() if 'loss' in _key) |
|
|
| log_vars['loss'] = loss |
| for loss_name, loss_value in log_vars.items(): |
| |
| if dist.is_available() and dist.is_initialized(): |
| loss_value = loss_value.data.clone() |
| dist.all_reduce(loss_value.div_(dist.get_world_size())) |
| log_vars[loss_name] = loss_value.item() |
|
|
| return loss, log_vars |
|
|
| def train_step(self, data: Dict, optimizer: torch.optim.Optimizer) -> Dict: |
| """The iteration step during training. |
| |
| This method defines an iteration step during training, excluding backpropagation |
| and optimizer updating, which are handled by an optimizer hook. |
| |
| Args: |
| data (dict): The output of the dataloader. |
| optimizer (torch.optim.Optimizer): The optimizer object (unused). |
| |
| Returns: |
| dict: A dictionary containing the loss, log_vars for logging, and the number of samples. |
| - ``loss``: A tensor for backpropagation, which may be a weighted sum of multiple losses. |
| - ``log_vars``: All the variables to be logged. |
| - ``num_samples``: The number of samples in the batch. |
| """ |
| losses = self(**data) |
| loss, log_vars = self._parse_losses(losses) |
|
|
| outputs = dict(loss=loss, log_vars=log_vars, num_samples=len(data['motion'])) |
| return outputs |
|
|
| def val_step(self, data: Dict, optimizer: torch.optim.Optimizer = None) -> Dict: |
| """The iteration step during validation. |
| |
| Args: |
| data (dict): The output of the dataloader. |
| optimizer (torch.optim.Optimizer, optional): The optimizer object (unused). |
| |
| Returns: |
| dict: A dictionary containing the loss, log_vars for logging, and the number of samples. |
| """ |
| losses = self(**data) |
| loss, log_vars = self._parse_losses(losses) |
|
|
| outputs = dict(loss=loss, log_vars=log_vars, num_samples=len(data['motion'])) |
| return outputs |
|
|
| def forward(self, **kwargs): |
| """Forward computation based on the training or testing mode.""" |
| if self.training: |
| return self.forward_train(**kwargs) |
| else: |
| return self.forward_test(**kwargs) |
|
|
| def split_results(self, results: Dict[str, torch.Tensor]) -> List[Dict]: |
| """Split batched results into individual outputs. |
| |
| Args: |
| results (dict): The batched results from the model containing 'motion', 'pred_motion', etc. |
| |
| Returns: |
| list: A list of dictionaries where each dictionary contains results for a single instance. |
| """ |
| B = results['motion'].shape[0] |
| output = [] |
| for i in range(B): |
| batch_output = dict() |
| batch_output['motion'] = to_cpu(results['motion'][i]) |
| batch_output['pred_motion'] = to_cpu(results['pred_motion'][i]) |
| batch_output['motion_length'] = to_cpu(results['motion_length'][i]) |
| batch_output['motion'][batch_output['motion_length']:, :] = 0 |
| batch_output['motion_mask'] = to_cpu(results['motion_mask'][i]) |
| if 'pred_motion_length' in results: |
| batch_output['pred_motion_length'] = to_cpu(results['pred_motion_length'][i]) |
| else: |
| batch_output['pred_motion_length'] = to_cpu(results['motion_length'][i]) |
| batch_output['pred_motion'][batch_output['pred_motion_length']:, :] = 0 |
| if 'pred_motion_mask' in results: |
| batch_output['pred_motion_mask'] = to_cpu(results['pred_motion_mask'][i]) |
| else: |
| batch_output['pred_motion_mask'] = to_cpu(results['motion_mask'][i]) |
| if 'motion_metas' in results: |
| motion_metas = results['motion_metas'][i] |
| if 'text' in motion_metas: |
| batch_output['text'] = motion_metas['text'] |
| if 'token' in motion_metas: |
| batch_output['token'] = motion_metas['token'] |
| if 'meta_data' in motion_metas and 'category_id' in motion_metas['meta_data']: |
| batch_output['category_id'] = motion_metas['meta_data']['category_id'] |
| batch_output['motion_metas'] = motion_metas |
| output.append(batch_output) |
| return output |
|
|
