Update modules.py

This commit is contained in:
Stardust·减
2023-08-20 15:25:13 +08:00
committed by GitHub
parent 5746b90b7b
commit 28e450fd5b

View File

@@ -16,6 +16,19 @@ from attentions import Encoder
LRELU_SLOPE = 0.1 LRELU_SLOPE = 0.1
class LayerNorm(nn.Module):
def __init__(self, channels, eps=1e-5):
super().__init__()
self.channels = channels
self.eps = eps
self.gamma = nn.Parameter(torch.ones(channels))
self.beta = nn.Parameter(torch.zeros(channels))
def forward(self, x):
x = x.transpose(1, -1)
x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
return x.transpose(1, -1)
class ConvReluNorm(nn.Module): class ConvReluNorm(nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout): def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):