Delete models/av_mossformer2_tse/av_mossformer_tmp.py

models/av_mossformer2_tse/av_mossformer_tmp.py

@@ -1,252 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchaudio
-
-import math
-
-from .mossformer.utils.one_path_flash_fsmn import Dual_Path_Model, SBFLASHBlock_DualA
-from models.av_mossformer2_tse.visual_frontend import VisualFrontend
-
-EPS = 1e-8
-
-class avMossformer(nn.Module):
-    def __init__(self, args):
-        super(avMossformer, self).__init__()
-        
-        N, L, = args.network_audio.encoder_out_nchannels, args.network_audio.encoder_kernel_size
-
-        self.encoder = Encoder(L, N)
-        self.separator = Separator(args)
-        self.decoder = Decoder(args, N, L)
-
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_normal_(p)
-
-    def forward(self, mixture, visual):
-        """
-        Args:
-            mixture: [M, T], M is batch size, T is #samples
-        Returns:
-            est_source: [M, C, T]
-        """
-        mixture_w = self.encoder(mixture)
-        est_mask = self.separator(mixture_w, visual)
-        est_source = self.decoder(mixture_w, est_mask)
-
-        # T changed after conv1d in encoder, fix it here
-        T_origin = mixture.size(-1)
-        T_conv = est_source.size(-1)
-        est_source = F.pad(est_source, (0, T_origin - T_conv))
-        return est_source
-
-class Encoder(nn.Module):
-    def __init__(self, L, N):
-        super(Encoder, self).__init__()
-        self.L, self.N = L, N
-        self.conv1d_U = nn.Conv1d(1, N, kernel_size=L, stride=L // 2, bias=False)
-
-    def forward(self, mixture):
-        """
-        Args:
-            mixture: [M, T], M is batch size, T is #samples
-        Returns:
-            mixture_w: [M, N, K], where K = (T-L)/(L/2)+1 = 2T/L-1
-        """
-        mixture = torch.unsqueeze(mixture, 1)  # [M, 1, T]
-        mixture_w = F.relu(self.conv1d_U(mixture))  # [M, N, K]
-        return mixture_w
-
-
-class Decoder(nn.Module):
-    def __init__(self, args, N, L):
-        super(Decoder, self).__init__()
-        self.N, self.L, self.args = N, L, args
-        self.basis_signals = nn.Linear(N, L, bias=False)
-
-    def forward(self, mixture_w, est_mask):
-        """
-        Args:
-            mixture_w: [M, N, K]
-            est_mask: [M, C, N, K]
-        Returns:
-            est_source: [M, C, T]
-        """
-        est_source = mixture_w * est_mask
-        est_source = torch.transpose(est_source, 2, 1) # [M,  K, N]
-        est_source = self.basis_signals(est_source)  # [M,  K, L]
-        est_source = overlap_and_add(est_source, self.L//2) # M x C x T
-        return est_source
-
-
-
-
-class Separator(nn.Module):
-    def __init__(self, args):
-        super(Separator, self).__init__()
-
-        self.layer_norm = nn.GroupNorm(1, args.network_audio.encoder_out_nchannels, eps=1e-8)
-        self.bottleneck_conv1x1 = nn.Conv1d(args.network_audio.encoder_out_nchannels, args.network_audio.encoder_out_nchannels, 1, bias=False)
-
-        # mossformer 2
-        intra_model = SBFLASHBlock_DualA(
-            num_layers=args.network_audio.intra_numlayers,
-            d_model=args.network_audio.encoder_out_nchannels,
-            nhead=args.network_audio.intra_nhead,
-            d_ffn=args.network_audio.intra_dffn,
-            dropout=args.network_audio.intra_dropout,
-            use_positional_encoding=args.network_audio.intra_use_positional,
-            norm_before=args.network_audio.intra_norm_before
-        )
-
-        self.masknet = Dual_Path_Model(
-            in_channels=args.network_audio.encoder_out_nchannels,
-            out_channels=args.network_audio.encoder_out_nchannels,
-            intra_model=intra_model,
-            num_layers=args.network_audio.masknet_numlayers,
-            norm=args.network_audio.masknet_norm,
-            K=args.network_audio.masknet_chunksize,
-            num_spks=args.network_audio.masknet_numspks,
-            skip_around_intra=args.network_audio.masknet_extraskipconnection,
-            linear_layer_after_inter_intra=args.network_audio.masknet_useextralinearlayer
-        )
-
-        # reference
-        # visual
-        stacks = []
-        for x in range(5):
-            stacks +=[VisualConv1D(V=256, H=512)]
-        self.visual_conv = nn.Sequential(*stacks)
-        self.v_ds = nn.Conv1d(512, 256, 1, bias=False)
-        self.av_conv = nn.Conv1d(args.network_audio.encoder_out_nchannels+args.network_reference.emb_size, args.network_audio.encoder_out_nchannels, 1, bias=True)
-
-
-    def forward(self, x, visual):
-        """
-        Keep this API same with TasNet
-        Args:
-            mixture_w: [M, N, K], M is batch size
-        returns:
-            est_mask: [M, C, N, K]
-        """
-        M, N, D = x.size()
-
-        x = self.layer_norm(x)
-        x = self.bottleneck_conv1x1(x)
-
-
-        visual = visual.transpose(1,2)
-        visual = self.v_ds(visual)
-        visual = self.visual_conv(visual)
-        visual = F.interpolate(visual, (D), mode='linear')
-        
-        x = torch.cat((x, visual),1)
-        x  = self.av_conv(x)
-
-        x = self.masknet(x)
-
-        x = x.squeeze(0)
-
-        return x
-
-
-
-def overlap_and_add(signal, frame_step):
-    """Reconstructs a signal from a framed representation.
-
-    Adds potentially overlapping frames of a signal with shape
-    `[..., frames, frame_length]`, offsetting subsequent frames by `frame_step`.
-    The resulting tensor has shape `[..., output_size]` where
-
-        output_size = (frames - 1) * frame_step + frame_length
-
-    Args:
-        signal: A [..., frames, frame_length] Tensor. All dimensions may be unknown, and rank must be at least 2.
-        frame_step: An integer denoting overlap offsets. Must be less than or equal to frame_length.
-
-    Returns:
-        A Tensor with shape [..., output_size] containing the overlap-added frames of signal's inner-most two dimensions.
-        output_size = (frames - 1) * frame_step + frame_length
-
-    Based on https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/contrib/signal/python/ops/reconstruction_ops.py
-    """
-    outer_dimensions = signal.size()[:-2]
-    frames, frame_length = signal.size()[-2:]
-
-    subframe_length = math.gcd(frame_length, frame_step)  # gcd=Greatest Common Divisor
-    subframe_step = frame_step // subframe_length
-    subframes_per_frame = frame_length // subframe_length
-    output_size = frame_step * (frames - 1) + frame_length
-    output_subframes = output_size // subframe_length
-
-    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length)
-
-    frame = torch.arange(0, output_subframes).unfold(0, subframes_per_frame, subframe_step)
-    frame = signal.new_tensor(frame).long().cuda()  # signal may in GPU or CPU
-    frame = frame.contiguous().view(-1)
-
-    result = signal.new_zeros(*outer_dimensions, output_subframes, subframe_length)
-    result.index_add_(-2, frame, subframe_signal)
-    result = result.view(*outer_dimensions, -1)
-    return result
-
-
-class av_mossformer_tmp(nn.Module):
-    def __init__(self, args):
-        super(av_mossformer_tmp, self).__init__()
-        args.causal=0
-        self.sep_network = avMossformer(args)
-        self.v_front_end = VisualFrontend(args)
-
-    def forward(self, mixture, ref):
-        ref = self.v_front_end(ref.unsqueeze(1)).transpose(1,2)
-        return self.sep_network(mixture, ref)
-
-
-
-class VisualConv1D(nn.Module):
-    def __init__(self, V=256, H=512):
-        super(VisualConv1D, self).__init__()
-        relu_0 = nn.ReLU()
-        norm_0 = GlobalLayerNorm(V)
-        conv1x1 = nn.Conv1d(V, H, 1, bias=False)
-        relu = nn.ReLU()
-        norm_1 = GlobalLayerNorm(H)
-        dsconv = nn.Conv1d(H, H, 3, stride=1, padding=1,dilation=1, groups=H, bias=False)
-        prelu = nn.PReLU()
-        norm_2 = GlobalLayerNorm(H)
-        pw_conv = nn.Conv1d(H, V, 1, bias=False)
-        self.net = nn.Sequential(relu_0, norm_0, conv1x1, relu, norm_1 ,dsconv, prelu, norm_2, pw_conv)
-
-    def forward(self, x):
-        out = self.net(x)
-        return out + x
-
-class GlobalLayerNorm(nn.Module):
-    """Global Layer Normalization (gLN)"""
-    def __init__(self, channel_size):
-        super(GlobalLayerNorm, self).__init__()
-        self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
-        self.beta = nn.Parameter(torch.Tensor(1, channel_size,1 ))  # [1, N, 1]
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        self.gamma.data.fill_(1)
-        self.beta.data.zero_()
-
-    def forward(self, y):
-        """
-        Args:
-            y: [M, N, K], M is batch size, N is channel size, K is length
-        Returns:
-            gLN_y: [M, N, K]
-        """
-        # TODO: in torch 1.0, torch.mean() support dim list
-        mean = y.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True) #[M, 1, 1]
-        var = (torch.pow(y-mean, 2)).mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)
-        gLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta
-        return gLN_y
-
-