File size: 4,371 Bytes
c40e77f | 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 | """Encodeur audio from scratch (Phase 7).
Pipeline entièrement maison :
1. **Mel-spectrogramme** calculé à la main (`src/data/audio_features.py`) : FFT de
PyTorch, banc de filtres mel triangulaires, log. Aucun modèle pré-entraîné.
2. **Patch embedding** : découpe le mel-spectrogramme en patchs (fréquence × temps)
et projette chacun en `n_embd` — comme un ViT mais sur une image 1 canal.
3. **Positional embedding appris** + blocs transformer **bidirectionnels** (l'audio
n'est pas causal : on voit tout le clip).
4. **Pooling** (CLS ou moyenne) → vecteur projeté dans l'espace texte.
Justification : docs/architecture.md §6.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn as nn
from src.model.block import TransformerBlock
@dataclass
class AudioConfig:
n_mels: int = 64
n_frames: int = 96
patch_mels: int = 16
patch_frames: int = 16
n_layer: int = 2
n_head: int = 4
n_embd: int = 128
dropout: float = 0.0
bias: bool = False
ffn_mult: int = 4
pool: str = "mean" # mean | cls
class MelPatchEmbedding(nn.Module):
"""Découpe un mel-spectrogramme (B, 1, n_mels, n_frames) en patchs."""
def __init__(self, n_mels: int, n_frames: int, patch_mels: int, patch_frames: int,
n_embd: int):
super().__init__()
if n_mels % patch_mels or n_frames % patch_frames:
raise ValueError("n_mels/n_frames doivent être divisibles par la taille de patch")
self.grid_mels = n_mels // patch_mels
self.grid_frames = n_frames // patch_frames
self.n_patches = self.grid_mels * self.grid_frames
self.proj = nn.Conv2d(1, n_embd, kernel_size=(patch_mels, patch_frames),
stride=(patch_mels, patch_frames))
def forward(self, mel: torch.Tensor) -> torch.Tensor:
x = self.proj(mel) # (B, D, gm, gf)
return x.flatten(2).transpose(1, 2) # (B, n_patches, D)
class AudioEncoder(nn.Module):
"""Encodeur de mel-spectrogrammes : patch embedding + position + blocs."""
def __init__(self, config: AudioConfig):
super().__init__()
self.config = config
self.patch_embed = MelPatchEmbedding(
config.n_mels, config.n_frames, config.patch_mels, config.patch_frames,
config.n_embd,
)
n_tokens = self.patch_embed.n_patches + (1 if config.pool == "cls" else 0)
self.cls_token = (
nn.Parameter(torch.zeros(1, 1, config.n_embd)) if config.pool == "cls" else None
)
self.pos_embedding = nn.Parameter(torch.zeros(1, n_tokens, config.n_embd))
self.dropout = nn.Dropout(config.dropout)
self.blocks = nn.ModuleList([
TransformerBlock(config.n_embd, config.n_head, dropout=config.dropout,
bias=config.bias, ffn_mult=config.ffn_mult,
use_rope=False, causal=False)
for _ in range(config.n_layer)
])
self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
self._init_weights()
def _init_weights(self) -> None:
nn.init.normal_(self.pos_embedding, std=0.02)
if self.cls_token is not None:
nn.init.normal_(self.cls_token, std=0.02)
for module in self.modules():
if isinstance(module, nn.Linear):
nn.init.normal_(module.weight, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.normal_(module.weight, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
@property
def n_audio_tokens(self) -> int:
return self.patch_embed.n_patches
def forward(self, mel: torch.Tensor) -> torch.Tensor:
x = self.patch_embed(mel)
if self.cls_token is not None:
x = torch.cat([self.cls_token.expand(x.size(0), -1, -1), x], dim=1)
x = self.dropout(x + self.pos_embedding)
for block in self.blocks:
x = block(x)
return self.ln_f(x)
def pool(self, tokens: torch.Tensor) -> torch.Tensor:
if self.config.pool == "cls":
return tokens[:, 0]
return tokens.mean(dim=1)
|