2.TasNet

from paderbox.database import keys
from padertorch.contrib.jensheit.data import SequenceProvider, MaskTransformer, STFT
import sacred
import torch
import torch.nn as nn
import padertorch as pt
import padertorch.train.optimizer as pt_opt
import torch.nn.functional as F
from padertorch.train.trainer import Trainer
from padertorch.ops.losses.loss import pit_loss
import numpy as np
import torch
from einops import rearrange
from padertorch.contrib.je.modules.conv import Pad, Trim
from tasnet_stft.desecting_tasnet.tasnet import si_loss
from tasnet_stft.desecting_tasnet.tasnet import Encoder
from tasnet_stft.desecting_tasnet.tasnet import Decoder

ex = sacred.Experiment()

def transform(example):
    example[keys.OBSERVATION] = example[keys.OBSERVATION].astype(np.float32)
    example[keys.SPEECH_SOURCE] = example[keys.SPEECH_SOURCE].astype(np.float32)
    return example

#def get_iterators():
provider_config = dict(
                        database= dict(factory='paderbox.database.merl_mixtures.MerlMixtures'),
                        audio_keys=[keys.OBSERVATION, keys.SPEECH_SOURCE],
                        batch_size=None,
                        batch_size_eval=None,
                        )
provider_config = SequenceProvider.get_config(provider_config)
provider = SequenceProvider.from_config(provider_config)
provider.transform = transform
train_iterator = provider.get_train_iterator()
test_iterator = provider.get_eval_iterator(num_examples=100)
#return train_iterator, test_iterator







#hyperparameter
# B = Batch
# T = Sample time
# L = Lenth of filters
# T/L = Padding_L
# N = number of output channels
# C = number of sources

class Tasnet(pt.Model):
    """
    >>> num_samples = 1000
    >>> signal = torch.ones(1000)
    >>> target = torch.ones(2, 1000)
    >>> net = Tasnet(40, 500,20,False,1000,4,True,False,2,False,False,False)
    >>> out = net(dict(observation=signal,num_samples=num_samples))
    >>> print(out.shape)

    """
    def __init__(self,L, N,stride,get_ipd_features,
                 hidden_size,num_layers,batch_first,bidirectional,C,
                    remove_ipd_features,beamforming, griffon_lim):

        super(Tasnet, self).__init__()

        self.L = L
        self.N = N
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.batch_first = batch_first
        self.bidirectional = bidirectional
        self.C = C

        self.encoder = Encoder(L, N, stride, get_ipd_features)
        self.separator = Separator(N,hidden_size,num_layers,batch_first,bidirectional,C)
        self.decoder = Decoder(L, N, stride, remove_ipd_features,
                 beamforming, griffon_lim)

    def forward(self, mixture):
        weight = self.encoder(mixture)                          # [B,Padding_T/L,N]
        est_mask = self.separator(weight)                       # [B,N,C,Padding_T/L]
        out = self.decoder(weight, est_mask,mixture)                    # [Padding_T]

        return out                                                          #[B,C,T]


    def review(self,mixture,output):

        clean = mixture[keys.SPEECH_SOURCE][None].permute(2,1,0)            # C T -> T C 1
        output = output.permute(2,1,0)
        #print("output.shape=",output.shape)

        def si_snr_loss(output,clean):
            estimate = output.permute(2,1,0)
            target = clean.permute(2,1,0)

            loss = si_loss(estimate, target, eps=1e-10)              #(B C T)
            return loss


        mixture = mixture['observation']                                   # T C 1-> T
        #print('mixture=',mixture.shape)
        speaker1 = 2*torch.squeeze(output[:,0])                                # T C 1-> T
        #print('speaker',speaker1.shape)
        speaker2 = 2*torch.squeeze(output[:,1])                                # T C 1-> T



        return {'loss': pit_loss(output, clean, si_snr_loss),
                'audios':{'mixture': (mixture,8000),
                          'speaker1': (speaker1,8000),
                          'speaker2': (speaker2, 8000),
                         }
                }




# signal = torch.ones(4,1000)
# enc = Encoder(20, 500)
# out = enc(signal)
# assert out.shape == (4, 500, 1000/20)

class Encoder(pt.Module):
    """
    >>> signal = torch.ones(1000,dtype=torch.float32)
    >>> enc = Encoder(40, 500,20,False)
    >>> out = enc(dict(observation=signal))
    >>> print(out.shape)
     """
    """Estimation of the nonnegative mixture weight by a 1-D conv layer.
    """
    def __init__(self, L, N, stride, get_ipd_features):
        super().__init__()
        # Hyper-parameter
        self.L, self.N, self.stride = L, N, stride

        # Components
        # 50% overlap
        self.conv1d_U = torch.nn.Conv1d(1, N, kernel_size=L,
                                        stride=stride, bias=False)
        self.pad = Pad()
        self.get_ipd_features = get_ipd_features

    def forward(self, mixture):
        """
        Args:
            mixture: [C, T], C is channel size, T is #samples
        Returns:
            mixture_w: [C, Fr, N], where Fr = (T-L)/(L/2)+1 = 2T/L-1

        """
        x = mixture['observation']  # mixture[T]
        if x.dim() == 1:
            mixture = torch.unsqueeze(x, 0)  # [1,T]
        pad_size = self.L - 1 - ((mixture.shape[-1] - 1) % self.stride)
        mixture = torch.unsqueeze(mixture, 1) # [C, 1, T]
        mixture = self.pad(mixture, pad_size)

        encoder_out = F.relu(self.conv1d_U(mixture)).permute(0,2,1)  # [C, Fr, N]
        return encoder_out


class Separator(nn.Module):
    """
        >>> weight = torch.ones(1,49,500)
        >>> sep = Separator(500,1000,4,True,False,2)
        >>> out = sep(weight)
        >>> print(out.shape)
        torch.Size([1, 500, 2, 49])

    """


    def __init__(self,N,hidden_size,num_layers,batch_first,bidirectional,C):
        super(Separator, self).__init__()
        self.N = N
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.C = C

        self.layer_norm = nn.LayerNorm(N)   #layer nomalization

        self.lstm1 = nn.LSTM(N,
                            hidden_size,
                            num_layers,
                            batch_first,
                            bidirectional
                            )
        #skip connection #not sure add net or add output
        self.lstm2 = nn.LSTM(N,
                             hidden_size,
                             2,
                             batch_first,
                             bidirectional
                             )

        self.fc = nn.Linear(hidden_size, self.C * self.N)

    def forward(self,weight):
        normalization = self.layer_norm(weight)

        out1, _ = self.lstm1(normalization)                   #[B,Padding_T/L,hidden_size]
        out2, _ = self.lstm2(normalization)
        out = out1 + out2
        out = self.fc(out)                             #[B,Padding_T/L,C*N]
        out = out.view(1,-1,self.C,self.N)             #[B,Padding_T/L,C,N]
        out = out.permute(0,2,1,3)                     #[B,N,C,Padding_T/L]
        #softmax
        est_mask = F.softmax(out,dim=2)

        return est_mask                                 #[B,K,F,N]

class Decoder(pt.Module):
    """
        >>> num_samples=1000
        >>> weight = torch.ones(1,50,500)
        >>> est_mask = torch.ones(1, 2, 50,500)
        >>> dec = Decoder(40,500,20,False,False,False)
        >>> out= dec(weight,est_mask,num_samples)
        >>> print(out.shape)




    """
    def __init__(self, L, N, stride, remove_ipd_features,
                 beamforming, griffon_lim):
        super().__init__()
        # Hyper-parameter
        self.N, self.L, self.stride = N, L, stride
        self.griffon_lim = griffon_lim
        # Components
        self.basis_signals = torch.nn.ConvTranspose1d(
            N, 1, kernel_size=L, stride=stride, bias=False)
        self.remove_ipd_features = remove_ipd_features
        self.cut = Trim()
        self.beamforming = beamforming

    def forward(self, mixture_w, est_mask, mixture):
        """
        Args:
            mixture_w: [C, F, N]
            est_mask: [C, K , F, N]
        Returns:
            est_source: [C, K, T]
        """
        num_samples = mixture['num_samples']
        # D = W * M
        if self.remove_ipd_features:
            mixture_w = mixture_w[..., :self.N]
        source_w = torch.unsqueeze(mixture_w, dim=1) * est_mask  # [C, K, F, N]
        # S = DV
        C, K, F, N = source_w.shape
        source_w = rearrange(source_w, 'c k f n -> (c k) n f')
        est_source = self.basis_signals(source_w)[:, 0]  # [C, K, F*L]
        size = est_source.shape[-1] - num_samples
        est_source = rearrange(est_source, '(c k) t -> c k t', c=C, k=K)
        return self.cut(est_source, size=size)

@ex.config
def config():
        tasnet_config = {
                        'L': 40,
                        'N': 500,
                        'stride':20,
                        'get_ipd_features':False,
                        'hidden_size': 1000,
                        'num_layers' : 4,
                        'batch_first': True,
                        'bidirectional': False,
                        'C': 2,
                        'remove_ipd_features':False,
                        'beamforming' : False,
                        'griffon_lim' : False,
                        }
        use_pt = True
        epochs = 2
        storage_dir = '/net/vol/zhenyuw/2.TasNet'

@ex.automain
def main(tasnet_config,use_pt,epochs,storage_dir):

    if use_pt:
        model = Tasnet(**tasnet_config)
        optimizer = pt_opt.Adam()

        trainer = Trainer(model,
                          storage_dir=storage_dir,
                          optimizer=optimizer,
                          loss_weights=None,
                          summary_trigger=(1000, 'iteration'),
                          checkpoint_trigger=(500, 'iteration'),
                          )

        trainer.test_run(train_iterator,
                         test_iterator)


        trainer.train(train_iterator,
                      resume=False,
                      device=0
                      )

        trainer.validate(test_iterator)