附录｜WaveNet 代码架构、运行入口与完整源码

"""Educational conditional WaveNet; synthetic plumbing test, not a speech model."""
import argparse
import copy
import json
import math
import platform
import time
from pathlib import Path

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


def mu_encode(x, mu=255):
    x = x.clamp(-1, 1)
    y = x.sign() * torch.log1p(mu * x.abs()) / math.log1p(mu)
    return ((y + 1) * mu / 2).round().long()


def mu_decode(q, mu=255):
    y = 2 * q.float() / mu - 1
    return y.sign() * torch.expm1(y.abs() * math.log1p(mu)) / mu


class CausalConv(nn.Module):
    def __init__(self, channels, dilation):
        super().__init__()
        self.dilation = dilation
        self.conv = nn.Conv1d(channels, 2 * channels, 2, dilation=dilation)

    def forward(self, x):
        return self.conv(F.pad(x, (self.dilation, 0)))

    def step(self, x, history):
        # x: [B,C], history: [B,C,d]; weights ordered old -> current.
        w = self.conv.weight
        y = F.linear(history[:, :, 0], w[:, :, 0])
        y = y + F.linear(x, w[:, :, 1], self.conv.bias)
        history = torch.cat((history[:, :, 1:], x.unsqueeze(-1)), dim=-1)
        return y, history


class Block(nn.Module):
    def __init__(self, channels, cond_channels, dilation):
        super().__init__()
        self.causal = CausalConv(channels, dilation)
        self.condition = nn.Conv1d(cond_channels, 2 * channels, 1)
        self.residual = nn.Conv1d(channels, channels, 1)
        self.skip = nn.Conv1d(channels, channels, 1)

    def forward(self, x, cond):
        a, b = (self.causal(x) + self.condition(cond)).chunk(2, dim=1)
        z = torch.tanh(a) * torch.sigmoid(b)
        return (x + self.residual(z)) / math.sqrt(2), self.skip(z)

    def step(self, x, cond, history):
        a, history = self.causal.step(x, history)
        a = a + self.condition(cond.unsqueeze(-1)).squeeze(-1)
        f, g = a.chunk(2, dim=1)
        z = (f.tanh() * g.sigmoid()).unsqueeze(-1)
        residual = (x + self.residual(z).squeeze(-1)) / math.sqrt(2)
        return residual, self.skip(z).squeeze(-1), history


class WaveNet(nn.Module):
    def __init__(self, channels=24, cond_channels=2,
                 dilations=(1, 2, 4, 8, 16, 32, 1, 2, 4, 8, 16, 32)):
        super().__init__()
        self.channels = channels
        self.dilations = tuple(dilations)
        self.receptive_field = 1 + sum(dilations)
        self.embedding = nn.Embedding(256, channels)
        self.blocks = nn.ModuleList([
            Block(channels, cond_channels, d) for d in dilations
        ])
        self.head = nn.Sequential(nn.ReLU(), nn.Conv1d(channels, channels, 1),
                                  nn.ReLU(), nn.Conv1d(channels, 256, 1))

    def forward(self, q_input, cond_target):
        # q_input[:,n] predicts q[:,n+1]; cond_target is aligned to q[:,n+1].
        x = self.embedding(q_input).transpose(1, 2)
        skips = []
        for block in self.blocks:
            x, skip = block(x, cond_target)
            skips.append(skip)
        return self.head(torch.stack(skips).sum(0) / math.sqrt(len(skips)))

    def initial_state(self, batch_size):
        p = next(self.parameters())
        return [p.new_zeros(batch_size, self.channels, d) for d in self.dilations]

    def step(self, q_previous, cond_target, state):
        x = self.embedding(q_previous)
        skips, next_state = [], []
        for block, history in zip(self.blocks, state):
            x, skip, history = block.step(x, cond_target, history)
            skips.append(skip)
            next_state.append(history)
        s = torch.stack(skips).sum(0) / math.sqrt(len(skips))
        return self.head(s.unsqueeze(-1)).squeeze(-1), next_state


@torch.no_grad()
def generate(model, seed, cond_future, temperature=1.0):
    """One known seed sample; cond_future [B,H,N] aligns to N new samples."""
    if temperature <= 0:
        raise ValueError('temperature must be positive')
    model.eval()
    state = model.initial_state(seed.shape[0])
    previous, samples = seed, []
    for i in range(cond_future.shape[-1]):
        logits, state = model.step(previous, cond_future[:, :, i], state)
        probs = torch.softmax(logits / temperature, dim=1)
        previous = torch.multinomial(probs, 1).squeeze(1)
        samples.append(previous)
    return torch.stack(samples, dim=1)


@torch.no_grad()
def verify(model):
    model.eval()
    q = torch.randint(0, 256, (2, 80))
    cond = torch.randn(2, 2, 80)
    y = model(q, cond)
    q2, c2 = q.clone(), cond.clone()
    q2[:, 41:] = torch.randint(0, 256, q2[:, 41:].shape)
    c2[:, :, 41:] = torch.randn_like(c2[:, :, 41:])
    causal_error = (y[:, :, :41] - model(q2, c2)[:, :, :41]).abs().max().item()
    state = model.initial_state(2)
    steps = []
    for i in range(q.shape[1]):
        z, state = model.step(q[:, i], cond[:, :, i], state)
        steps.append(z)
    cache_error = (y - torch.stack(steps, dim=2)).abs().max().item()
    bins = torch.arange(256)
    assert torch.equal(mu_encode(mu_decode(bins)), bins)
    assert y.shape == (2, 256, 80)
    assert causal_error < 1e-6
    torch.testing.assert_close(y, torch.stack(steps, dim=2), atol=1e-4, rtol=1e-5)
    # Conv1d and Linear may use different float32 accumulation orders.
    # Double precision independently checks the cache indexing after training.
    precise = copy.deepcopy(model).double()
    full64 = precise(q, cond.double())
    state64, steps64 = precise.initial_state(2), []
    for i in range(q.shape[1]):
        z, state64 = precise.step(q[:, i], cond[:, :, i].double(), state64)
        steps64.append(z)
    cache64_error = (full64 - torch.stack(steps64, dim=2)).abs().max().item()
    assert cache64_error < 1e-10
    # Finite receptive field: earliest input cannot affect final position if T>R.
    long_q = torch.randint(0, 256, (1, model.receptive_field + 8))
    long_c = torch.randn(1, 2, long_q.shape[1])
    before = model(long_q, long_c)[:, :, -1]
    long_q[:, 0] = (long_q[:, 0] + 100) % 256
    rf_error = (before - model(long_q, long_c)[:, :, -1]).abs().max().item()
    assert rf_error < 1e-6
    return dict(causal_max_error=causal_error, cache_max_error=cache_error,
                cache_float64_max_error=cache64_error,
                outside_receptive_field_error=rf_error, quantizer_roundtrip=True,
                output_shape=list(y.shape))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--steps', type=int, default=300)
    parser.add_argument('--out', type=Path, default=Path('results'))
    args = parser.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(2)
    torch.manual_seed(7)
    model = WaveNet()
    checks = verify(model)
    # Fixed 4-example batch: deliberately an overfit test, no validation set.
    fs = 8000
    t = torch.arange(513) / fs
    f0 = torch.tensor([125., 125., 250., 250.])[:, None]
    phase = torch.tensor([0., 0.7, 0., 0.7])[:, None]
    amp = torch.tensor([0.55, 0.75, 0.55, 0.75])[:, None]
    wave = amp * (torch.sin(2 * math.pi * f0 * t + phase)
                  + .2 * torch.sin(4 * math.pi * f0 * t + 2 * phase)) / 1.2
    q = mu_encode(wave)
    cond = torch.stack(((f0 / 250).expand(-1, 513), amp.expand(-1, 513)), dim=1)
    inp, target, aligned_cond = q[:, :-1], q[:, 1:], cond[:, :, 1:]
    warmup = model.receptive_field - 1
    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    losses = []
    start = time.perf_counter()
    for step in range(args.steps + 1):
        model.train()
        logits = model(inp, aligned_cond)
        loss = F.cross_entropy(logits[:, :, warmup:], target[:, warmup:])
        if step % 10 == 0 or step == args.steps:
            losses.append({'step': step, 'loss': loss.item()})
        if step == args.steps:
            break
        opt.zero_grad(set_to_none=True)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    elapsed = time.perf_counter() - start
    assert losses[-1]['loss'] < losses[0]['loss'] * .7, 'overfit smoke test failed'
    post_checks = verify(model)
    samples = generate(model, q[:1, 0], cond[:1, :, 1:129])
    result = dict(python=platform.python_version(), torch=torch.__version__,
                  device='cpu', seed=7, threads=2, steps=args.steps,
                  parameters=sum(p.numel() for p in model.parameters()),
                  receptive_field=model.receptive_field, sample_rate=fs,
                  elapsed_seconds=elapsed, losses=losses,
                  checks_before=checks, checks_after=post_checks,
                  generated_shape=list(samples.shape),
                  warning='Fixed synthetic training batch; no speech quality/generalization claim.')
    (args.out / 'metrics.json').write_text(json.dumps(result, indent=2), encoding='utf-8')
    torch.save(model.state_dict(), args.out / 'toy_weights.pt')
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()