wonkitty/apple_oh

1# Copyright (c) Facebook, Inc. and its affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import random8import torch as th9from torch import nn10 11 12class Shift(nn.Module):13    """14    Randomly shift audio in time by up to `shift` samples.15    """16    def __init__(self, shift=8192):17        super().__init__()18        self.shift = shift19 20    def forward(self, wav):21        batch, sources, channels, time = wav.size()22        length = time - self.shift23        if self.shift > 0:24            if not self.training:25                wav = wav[..., :length]26            else:27                offsets = th.randint(self.shift, [batch, sources, 1, 1], device=wav.device)28                offsets = offsets.expand(-1, -1, channels, -1)29                indexes = th.arange(length, device=wav.device)30                wav = wav.gather(3, indexes + offsets)31        return wav32 33 34class FlipChannels(nn.Module):35    """36    Flip left-right channels.37    """38    def forward(self, wav):39        batch, sources, channels, time = wav.size()40        if self.training and wav.size(2) == 2:41            left = th.randint(2, (batch, sources, 1, 1), device=wav.device)42            left = left.expand(-1, -1, -1, time)43            right = 1 - left44            wav = th.cat([wav.gather(2, left), wav.gather(2, right)], dim=2)45        return wav46 47 48class FlipSign(nn.Module):49    """50    Random sign flip.51    """52    def forward(self, wav):53        batch, sources, channels, time = wav.size()54        if self.training:55            signs = th.randint(2, (batch, sources, 1, 1), device=wav.device, dtype=th.float32)56            wav = wav * (2 * signs - 1)57        return wav58 59 60class Remix(nn.Module):61    """62    Shuffle sources to make new mixes.63    """64    def __init__(self, group_size=4):65        """66        Shuffle sources within one batch.67        Each batch is divided into groups of size `group_size` and shuffling is done within68        each group separatly. This allow to keep the same probability distribution no matter69        the number of GPUs. Without this grouping, using more GPUs would lead to a higher70        probability of keeping two sources from the same track together which can impact71        performance.72        """73        super().__init__()74        self.group_size = group_size75 76    def forward(self, wav):77        batch, streams, channels, time = wav.size()78        device = wav.device79 80        if self.training:81            group_size = self.group_size or batch82            if batch % group_size != 0:83                raise ValueError(f"Batch size {batch} must be divisible by group size {group_size}")84            groups = batch // group_size85            wav = wav.view(groups, group_size, streams, channels, time)86            permutations = th.argsort(th.rand(groups, group_size, streams, 1, 1, device=device),87                                      dim=1)88            wav = wav.gather(1, permutations.expand(-1, -1, -1, channels, time))89            wav = wav.view(batch, streams, channels, time)90        return wav91 92 93class Scale(nn.Module):94    def __init__(self, proba=1., min=0.25, max=1.25):95        super().__init__()96        self.proba = proba97        self.min = min98        self.max = max99 100    def forward(self, wav):101        batch, streams, channels, time = wav.size()102        device = wav.device103        if self.training and random.random() < self.proba:104            scales = th.empty(batch, streams, 1, 1, device=device).uniform_(self.min, self.max)105            wav *= scales106        return wav107