jone/Music_Source_Separation

1from typing import Any, Callable, Dict2 3import pytorch_lightning as pl4import torch5import torch.nn as nn6import torch.optim as optim7from torch.optim.lr_scheduler import LambdaLR8 9 10class LitSourceSeparation(pl.LightningModule):11    def __init__(12        self,13        batch_data_preprocessor,14        model: nn.Module,15        loss_function: Callable,16        optimizer_type: str,17        learning_rate: float,18        lr_lambda: Callable,19    ):20        r"""Pytorch Lightning wrapper of PyTorch model, including forward,21        optimization of model, etc.22 23        Args:24            batch_data_preprocessor: object, used for preparing inputs and25                targets for training. E.g., BasicBatchDataPreprocessor is used26                for preparing data in dictionary into tensor.27            model: nn.Module28            loss_function: function29            learning_rate: float30            lr_lambda: function31        """32        super().__init__()33 34        self.batch_data_preprocessor = batch_data_preprocessor35        self.model = model36        self.optimizer_type = optimizer_type37        self.loss_function = loss_function38        self.learning_rate = learning_rate39        self.lr_lambda = lr_lambda40 41    def training_step(self, batch_data_dict: Dict, batch_idx: int) -> torch.float:42        r"""Forward a mini-batch data to model, calculate loss function, and43        train for one step. A mini-batch data is evenly distributed to multiple44        devices (if there are) for parallel training.45 46        Args:47            batch_data_dict: e.g. {48                'vocals': (batch_size, channels_num, segment_samples),49                'accompaniment': (batch_size, channels_num, segment_samples),50                'mixture': (batch_size, channels_num, segment_samples)51            }52            batch_idx: int53 54        Returns:55            loss: float, loss function of this mini-batch56        """57        input_dict, target_dict = self.batch_data_preprocessor(batch_data_dict)58        # input_dict: {59        #     'waveform': (batch_size, channels_num, segment_samples),60        #     (if_exist) 'condition': (batch_size, channels_num),61        # }62        # target_dict: {63        #     'waveform': (batch_size, target_sources_num * channels_num, segment_samples),64        # }65 66        # Forward.67        self.model.train()68 69        output_dict = self.model(input_dict)70        # output_dict: {71        #     'waveform': (batch_size, target_sources_num * channels_num, segment_samples),72        # }73 74        outputs = output_dict['waveform']75        # outputs:, e.g, (batch_size, target_sources_num * channels_num, segment_samples)76 77        # Calculate loss.78        loss = self.loss_function(79            output=outputs,80            target=target_dict['waveform'],81            mixture=input_dict['waveform'],82        )83 84        return loss85 86    def configure_optimizers(self) -> Any:87        r"""Configure optimizer."""88 89        if self.optimizer_type == "Adam":90            optimizer = optim.Adam(91                self.model.parameters(),92                lr=self.learning_rate,93                betas=(0.9, 0.999),94                eps=1e-08,95                weight_decay=0.0,96                amsgrad=True,97            )98 99        elif self.optimizer_type == "AdamW":100            optimizer = optim.AdamW(101                self.model.parameters(),102                lr=self.learning_rate,103                betas=(0.9, 0.999),104                eps=1e-08,105                weight_decay=0.0,106                amsgrad=True,107            )108 109        else:110            raise NotImplementedError111 112        scheduler = {113            'scheduler': LambdaLR(optimizer, self.lr_lambda),114            'interval': 'step',115            'frequency': 1,116        }117 118        return [optimizer], [scheduler]119 120 121def get_model_class(model_type):122    r"""Get model.123 124    Args:125        model_type: str, e.g., 'ResUNet143_DecouplePlusInplaceABN'126 127    Returns:128        nn.Module129    """130    if model_type == 'ResUNet143_DecouplePlusInplaceABN_ISMIR2021':131        from bytesep.models.resunet_ismir2021 import (132            ResUNet143_DecouplePlusInplaceABN_ISMIR2021,133        )134 135        return ResUNet143_DecouplePlusInplaceABN_ISMIR2021136 137    elif model_type == 'UNet':138        from bytesep.models.unet import UNet139 140        return UNet141 142    elif model_type == 'UNetSubbandTime':143        from bytesep.models.unet_subbandtime import UNetSubbandTime144 145        return UNetSubbandTime146 147    elif model_type == 'ResUNet143_Subbandtime':148        from bytesep.models.resunet_subbandtime import ResUNet143_Subbandtime149 150        return ResUNet143_Subbandtime151 152    elif model_type == 'ResUNet143_DecouplePlus':153        from bytesep.models.resunet import ResUNet143_DecouplePlus154 155        return ResUNet143_DecouplePlus156 157    elif model_type == 'ConditionalUNet':158        from bytesep.models.conditional_unet import ConditionalUNet159 160        return ConditionalUNet161 162    elif model_type == 'LevelRNN':163        from bytesep.models.levelrnn import LevelRNN164 165        return LevelRNN166 167    elif model_type == 'WavUNet':168        from bytesep.models.wavunet import WavUNet169 170        return WavUNet171 172    elif model_type == 'WavUNetLevelRNN':173        from bytesep.models.wavunet_levelrnn import WavUNetLevelRNN174 175        return WavUNetLevelRNN176 177    elif model_type == 'TTnet':178        from bytesep.models.ttnet import TTnet179 180        return TTnet181 182    elif model_type == 'TTnetNoTransformer':183        from bytesep.models.ttnet_no_transformer import TTnetNoTransformer184 185        return TTnetNoTransformer186 187    else:188        raise NotImplementedError189