附录｜Tacotron 教学实现：代码导读与完整源码

"""CPU teaching model: CBHG + additive attention + reduction decoder.

This is NOT a trained speech synthesizer. Run this file to verify the tensor
contracts and numerical examples used in the accompanying articles.
"""
import json
import platform
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F


class SameConv(nn.Module):
    def __init__(self, cin, cout, kernel, activate=True):
        super().__init__()
        self.kernel, self.activate = kernel, activate
        self.conv = nn.Conv1d(cin, cout, kernel)
        self.bn = nn.BatchNorm1d(cout)

    def forward(self, x):
        # [B,C,N]; even kernels require asymmetric padding.
        left = (self.kernel - 1) // 2
        x = self.conv(F.pad(x, (left, self.kernel - 1 - left)))
        x = self.bn(x)
        return F.relu(x) if self.activate else x


class Highway(nn.Module):
    def __init__(self, width):
        super().__init__()
        self.h, self.g = nn.Linear(width, width), nn.Linear(width, width)
        nn.init.constant_(self.g.bias, -1.)

    def forward(self, x):
        gate = torch.sigmoid(self.g(x))
        return gate * F.relu(self.h(x)) + (1 - gate) * x


class CBHG(nn.Module):
    def __init__(self, width=16, bank_size=4):
        super().__init__()
        self.bank = nn.ModuleList(SameConv(width, width, k)
                                  for k in range(1, bank_size + 1))
        self.p1 = SameConv(width * bank_size, width, 3)
        self.p2 = SameConv(width, width, 3, activate=False)
        self.highway = nn.Sequential(*(Highway(width) for _ in range(4)))
        self.rnn = nn.GRU(width, width, batch_first=True, bidirectional=True)

    def forward(self, x):
        residual = x.transpose(1, 2)
        bank = torch.cat([conv(residual) for conv in self.bank], dim=1)
        # Pool width 2, stride 1; -inf prevents padded values winning max.
        pooled = F.max_pool1d(F.pad(bank, (0, 1), value=-float('inf')), 2, 1)
        x = (self.p2(self.p1(pooled)) + residual).transpose(1, 2)
        x = self.highway(x)
        return self.rnn(x)[0]  # [B,N,2*width]


class Attention(nn.Module):
    def __init__(self, memory_dim, query_dim, attention_dim):
        super().__init__()
        self.key = nn.Linear(memory_dim, attention_dim, bias=False)
        self.query = nn.Linear(query_dim, attention_dim, bias=False)
        self.score = nn.Linear(attention_dim, 1, bias=False)
        self.bias = nn.Parameter(torch.zeros(attention_dim))

    def forward(self, query, memory, valid):
        if not valid.any(dim=1).all():
            raise ValueError('Every text must contain at least one valid token')
        energy = self.score(torch.tanh(self.key(memory)
                            + self.query(query)[:, None, :] + self.bias)).squeeze(-1)
        energy = energy.masked_fill(~valid, -float('inf'))
        weights = energy.softmax(dim=-1)  # normalization over text positions
        context = torch.bmm(weights[:, None, :], memory).squeeze(1)
        return context, weights


class TinyTacotron(nn.Module):
    def __init__(self, vocab=12, width=16, mel_dim=8, reduction=2):
        super().__init__()
        self.mel_dim, self.reduction, self.width = mel_dim, reduction, width
        self.embedding = nn.Embedding(vocab, width)
        self.encoder_prenet = nn.Sequential(nn.Linear(width, width), nn.ReLU(),
                                           nn.Dropout(.5))
        self.encoder = CBHG(width)
        self.decoder_prenet = nn.Sequential(nn.Linear(mel_dim, width), nn.ReLU(),
                                           nn.Dropout(.5))
        self.attention_rnn = nn.GRUCell(width, width * 2)
        self.attention = Attention(width * 2, width * 2, width)
        self.input_projection = nn.Linear(width * 4, width * 2)
        self.decoder1 = nn.GRUCell(width * 2, width * 2)
        self.decoder2 = nn.GRUCell(width * 2, width * 2)
        self.output = nn.Linear(width * 2, reduction * mel_dim)

    def forward(self, tokens, targets=None, steps=None):
        # Demo batch uses equal-length unpadded text. Attention itself accepts
        # a mask; variable-length encoder batching requires packing/masking too.
        if (tokens == 0).any():
            raise ValueError('Demo encoder requires unpadded, equal-length text')
        memory = self.encoder(self.encoder_prenet(self.embedding(tokens)))
        valid = torch.ones_like(tokens, dtype=torch.bool)
        batch = tokens.shape[0]
        previous = memory.new_zeros(batch, self.mel_dim)  # GO frame
        aq = memory.new_zeros(batch, self.width * 2)
        d1, d2 = torch.zeros_like(aq), torch.zeros_like(aq)
        if targets is not None:
            if targets.shape[1] % self.reduction:
                raise ValueError('Pad target length to a multiple of reduction')
            steps = targets.shape[1] // self.reduction
        if steps is None:
            raise ValueError('Demo inference requires an explicit step limit')
        blocks, alignments = [], []
        for s in range(steps):
            aq = self.attention_rnn(self.decoder_prenet(previous), aq)
            context, weights = self.attention(aq, memory, valid)
            decoder_input = self.input_projection(torch.cat([aq, context], -1))
            d1 = self.decoder1(decoder_input, d1)
            d2 = self.decoder2(d1 + decoder_input, d2)
            block = self.output(d2 + d1 + decoder_input)
            block = block.reshape(batch, self.reduction, self.mel_dim)
            blocks.append(block)
            alignments.append(weights)
            previous = (targets[:, (s + 1) * self.reduction - 1]
                        if targets is not None else block[:, -1])
        return torch.cat(blocks, 1), torch.stack(alignments, 1)


def verify():
    torch.manual_seed(7)
    torch.set_num_threads(1)
    # Exact hand example: exp energies proportional to [1,2,7].
    weights = torch.tensor([1., 2., 7.]).log().softmax(-1)
    context = weights @ torch.tensor([[1., 0.], [0., 2.], [3., 1.]])
    torch.testing.assert_close(context, torch.tensor([2.2, 1.1]))
    att = Attention(4, 4, 3)
    memory, query = torch.randn(2, 4, 4), torch.randn(2, 4)
    mask = torch.tensor([[1, 1, 1, 0], [1, 1, 0, 0]], dtype=torch.bool)
    _, a = att(query, memory, mask)
    altered = memory.clone()
    altered[~mask] = 1e4
    c1, _ = att(query, memory, mask)
    c2, _ = att(query, altered, mask)
    torch.testing.assert_close(c1, c2)
    assert a[~mask].count_nonzero() == 0
    torch.testing.assert_close(a.sum(-1), torch.ones(2))
    model = TinyTacotron().eval()
    text = torch.tensor([[1, 2, 3, 4], [4, 3, 2, 1]])
    target = torch.randn(2, 6, 8)
    pred, alignment = model(text, target)
    changed = target.clone()
    changed[:, 3] += 9  # 0-based frame 3 feeds decoder step 2 only.
    perturbed, _ = model(text, changed)
    prefix_error = (pred[:, :4] - perturbed[:, :4]).abs().max().item()
    suffix_delta = (pred[:, 4:] - perturbed[:, 4:]).abs().max().item()
    assert prefix_error == 0 and suffix_delta > 1e-5
    generated, _ = model(text, steps=3)
    assert pred.shape == generated.shape == (2, 6, 8)
    assert alignment.shape == (2, 3, 4)
    model.train()
    training_pred, _ = model(text, target)
    loss = F.l1_loss(training_pred, target)
    loss.backward()
    gradients = {name: p.grad.abs().sum().item() for name, p in model.named_parameters()
                 if p.grad is not None}
    assert all(torch.isfinite(p.grad).all() for p in model.parameters() if p.grad is not None)
    assert gradients['embedding.weight'] > 0 and gradients['attention.score.weight'] > 0
    result = dict(python=platform.python_version(), torch=torch.__version__, device='CPU',
                  seed=7, weights=weights.tolist(), context=context.tolist(),
                  predicted_shape=list(pred.shape), alignment_shape=list(alignment.shape),
                  padding_context_error=(c1-c2).abs().max().item(),
                  target_prefix_error=prefix_error, target_suffix_delta=suffix_delta,
                  l1_loss=loss.item(), gradient_tensors=len(gradients),
                  parameters=sum(p.numel() for p in model.parameters()),
                  limitation='Random initialization; no speech training or MOS evaluation')
    out = Path(__file__).parent / 'results'
    out.mkdir(exist_ok=True)
    (out / 'checks.json').write_text(json.dumps(result, indent=2), encoding='utf-8')
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    verify()