File size: 2,691 Bytes
ad16788 |
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 |
#!/usr/bin/env python3
# Copyright 2020 Tomoki Hayashi
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
"""Variance predictor related modules."""
import torch
from typeguard import check_argument_types
from espnet.nets.pytorch_backend.transformer.layer_norm import LayerNorm
class VariancePredictor(torch.nn.Module):
"""Variance predictor module.
This is a module of variacne predictor described in `FastSpeech 2:
Fast and High-Quality End-to-End Text to Speech`_.
.. _`FastSpeech 2: Fast and High-Quality End-to-End Text to Speech`:
https://arxiv.org/abs/2006.04558
"""
def __init__(
self,
idim: int,
n_layers: int = 2,
n_chans: int = 384,
kernel_size: int = 3,
bias: bool = True,
dropout_rate: float = 0.5,
):
"""Initilize duration predictor module.
Args:
idim (int): Input dimension.
n_layers (int, optional): Number of convolutional layers.
n_chans (int, optional): Number of channels of convolutional layers.
kernel_size (int, optional): Kernel size of convolutional layers.
dropout_rate (float, optional): Dropout rate.
"""
assert check_argument_types()
super().__init__()
self.conv = torch.nn.ModuleList()
for idx in range(n_layers):
in_chans = idim if idx == 0 else n_chans
self.conv += [
torch.nn.Sequential(
torch.nn.Conv1d(
in_chans,
n_chans,
kernel_size,
stride=1,
padding=(kernel_size - 1) // 2,
bias=bias,
),
torch.nn.ReLU(),
LayerNorm(n_chans, dim=1),
torch.nn.Dropout(dropout_rate),
)
]
self.linear = torch.nn.Linear(n_chans, 1)
def forward(self, xs: torch.Tensor, x_masks: torch.Tensor = None) -> torch.Tensor:
"""Calculate forward propagation.
Args:
xs (Tensor): Batch of input sequences (B, Tmax, idim).
x_masks (ByteTensor, optional):
Batch of masks indicating padded part (B, Tmax).
Returns:
Tensor: Batch of predicted sequences (B, Tmax, 1).
"""
xs = xs.transpose(1, -1) # (B, idim, Tmax)
for f in self.conv:
xs = f(xs) # (B, C, Tmax)
xs = self.linear(xs.transpose(1, 2)) # (B, Tmax, 1)
if x_masks is not None:
xs = xs.masked_fill(x_masks, 0.0)
return xs
|