Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,075 Bytes
8e8cd3e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
#!/usr/bin/env python -u
# -*- coding: utf-8 -*-
# Copyright 2018 Northwestern Polytechnical University (author: Ke Wang)
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import torch.nn as nn
class CLayerNorm(nn.LayerNorm):
"""Channel-wise layer normalization."""
def __init__(self, *args, **kwargs):
super(CLayerNorm, self).__init__(*args, **kwargs)
def forward(self, sample):
"""Forward function.
Args:
sample: [batch_size, channels, length]
"""
if sample.dim() != 3:
raise RuntimeError('{} only accept 3-D tensor as input'.format(
self.__name__))
# [N, C, T] -> [N, T, C]
sample = torch.transpose(sample, 1, 2)
# LayerNorm
sample = super().forward(sample)
# [N, T, C] -> [N, C, T]
sample = torch.transpose(sample, 1, 2)
return sample
class ILayerNorm(nn.InstanceNorm1d):
"""Channel-wise layer normalization."""
def __init__(self, *args, **kwargs):
super(ILayerNorm, self).__init__(*args, **kwargs)
def forward(self, sample):
"""Forward function.
Args:
sample: [batch_size, channels, length]
"""
if sample.dim() != 3:
raise RuntimeError('{} only accept 3-D tensor as input'.format(
self.__name__))
# [N, C, T] -> [N, T, C]
sample = torch.transpose(sample, 1, 2)
# LayerNorm
sample = super().forward(sample)
# [N, T, C] -> [N, C, T]
sample = torch.transpose(sample, 1, 2)
return sample
class GLayerNorm(nn.Module):
"""Global Layer Normalization for TasNet."""
def __init__(self, channels, eps=1e-5):
super(GLayerNorm, self).__init__()
self.eps = eps
self.norm_dim = channels
self.gamma = nn.Parameter(torch.Tensor(channels))
self.beta = nn.Parameter(torch.Tensor(channels))
self.reset_parameters()
def reset_parameters(self):
nn.init.ones_(self.gamma)
nn.init.zeros_(self.beta)
def forward(self, sample):
"""Forward function.
Args:
sample: [batch_size, channels, length]
"""
if sample.dim() != 3:
raise RuntimeError('{} only accept 3-D tensor as input'.format(
self.__name__))
# [N, C, T] -> [N, T, C]
sample = torch.transpose(sample, 1, 2)
# Mean and variance [N, 1, 1]
mean = torch.mean(sample, (1, 2), keepdim=True)
var = torch.mean((sample - mean)**2, (1, 2), keepdim=True)
sample = (sample - mean) / torch.sqrt(var + self.eps) * \
self.gamma + self.beta
# [N, T, C] -> [N, C, T]
sample = torch.transpose(sample, 1, 2)
return sample
class _LayerNorm(nn.Module):
"""Layer Normalization base class."""
def __init__(self, channel_size):
super(_LayerNorm, self).__init__()
self.channel_size = channel_size
self.gamma = nn.Parameter(torch.ones(channel_size),
requires_grad=True)
self.beta = nn.Parameter(torch.zeros(channel_size),
requires_grad=True)
def apply_gain_and_bias(self, normed_x):
""" Assumes input of size `[batch, chanel, *]`. """
return (self.gamma * normed_x.transpose(1, -1) +
self.beta).transpose(1, -1)
class GlobLayerNorm(_LayerNorm):
"""Global Layer Normalization (globLN)."""
def forward(self, x):
""" Applies forward pass.
Works for any input size > 2D.
Args:
x (:class:`torch.Tensor`): Shape `[batch, chan, *]`
Returns:
:class:`torch.Tensor`: gLN_x `[batch, chan, *]`
"""
dims = list(range(1, len(x.shape)))
mean = x.mean(dim=dims, keepdim=True)
var = torch.pow(x - mean, 2).mean(dim=dims, keepdim=True)
return self.apply_gain_and_bias((x - mean) / (var + 1e-8).sqrt())
|