358 lines
14 KiB
Python
358 lines
14 KiB
Python
import torch
|
|
from torch import Tensor, device, no_grad, cuda
|
|
from torch.optim import AdamW, Optimizer
|
|
from torch.optim.lr_scheduler import OneCycleLR
|
|
from torch.nn import CTCLoss, utils
|
|
from torch.utils.data import DataLoader
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
from pathlib import Path
|
|
from tqdm import tqdm
|
|
from datetime import datetime
|
|
from torch.amp.grad_scaler import GradScaler
|
|
from tokenizer import ASRTokenizer
|
|
from dataset import Batch, create_dataloader
|
|
from model import ASRModel
|
|
|
|
|
|
# ============ 全局配置 ============
|
|
CONFIG = {
|
|
# 数据配置
|
|
'batch_size': 96,
|
|
|
|
# 训练配置
|
|
'num_epochs': 50,
|
|
'learning_rate': 1e-4,
|
|
'weight_decay': 1e-4,
|
|
'grad_clip_norm': 1.0,
|
|
|
|
# 模型配置
|
|
'input_dim': 640,
|
|
'num_heads': 8,
|
|
'ffn_dim': 2048,
|
|
'num_layers': 8,
|
|
'dropout': 0.1,
|
|
|
|
# 保存和评估
|
|
'early_stopping_patience': 12,
|
|
}
|
|
|
|
|
|
def calculate_cer(pred: str, target: str) -> float:
|
|
""" 计算字符错误率 (Character Error Rate) 使用编辑距离算法"""
|
|
if len(target) == 0:
|
|
return 0.0 if len(pred) == 0 else 1.0
|
|
|
|
d = [[0] * (len(target) + 1) for _ in range(len(pred) + 1)]
|
|
|
|
for i in range(len(pred) + 1):
|
|
d[i][0] = i
|
|
for j in range(len(target) + 1):
|
|
d[0][j] = j
|
|
|
|
for i in range(1, len(pred) + 1):
|
|
for j in range(1, len(target) + 1):
|
|
if pred[i-1] == target[j-1]:
|
|
d[i][j] = d[i-1][j-1]
|
|
else:
|
|
d[i][j] = min(d[i-1][j], d[i][j-1], d[i-1][j-1]) + 1
|
|
|
|
return d[len(pred)][len(target)] / len(target)
|
|
|
|
|
|
def calculate_wer(pred: str, target: str) -> float:
|
|
""" 计算词错误率 (Word Error Rate) """
|
|
pred_words = pred.split()
|
|
target_words = target.split()
|
|
|
|
if len(target_words) == 0:
|
|
return 0.0 if len(pred_words) == 0 else 1.0
|
|
|
|
d = [[0] * (len(target_words) + 1) for _ in range(len(pred_words) + 1)]
|
|
|
|
for i in range(len(pred_words) + 1):
|
|
d[i][0] = i
|
|
for j in range(len(target_words) + 1):
|
|
d[0][j] = j
|
|
|
|
for i in range(1, len(pred_words) + 1):
|
|
for j in range(1, len(target_words) + 1):
|
|
if pred_words[i-1] == target_words[j-1]:
|
|
d[i][j] = d[i-1][j-1]
|
|
else:
|
|
d[i][j] = min(d[i-1][j], d[i][j-1], d[i-1][j-1]) + 1
|
|
|
|
return d[len(pred_words)][len(target_words)] / len(target_words)
|
|
|
|
def train_one_epoch(model: ASRModel, dataloader: DataLoader, criterion: CTCLoss, optimizer: Optimizer, scheduler: OneCycleLR, device: device, epoch: int, writer: SummaryWriter, global_step: int, scaler: GradScaler) -> tuple[float, int]:
|
|
model.train()
|
|
num_batches = 0
|
|
total_train_loss = 0
|
|
|
|
progress_bar = tqdm(dataloader, desc=f"Epoch {epoch}")
|
|
for batch_index, batch in enumerate(progress_bar):
|
|
batch: Batch
|
|
mel_specs = batch['mel_specs'].to(device)
|
|
targets = batch['targets'].to(device)
|
|
mel_lengths = batch['mel_lengths'].to(device)
|
|
target_lengths = batch['target_lengths'].to(device)
|
|
|
|
optimizer.zero_grad()
|
|
|
|
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
|
|
log_probs, lengths = model(mel_specs=mel_specs, mel_lengths=mel_lengths)
|
|
log_probs: Tensor
|
|
log_probs_ctc = log_probs.permute(1, 0, 2)
|
|
loss: Tensor = criterion(log_probs=log_probs_ctc, targets=targets, input_lengths=lengths, target_lengths=target_lengths)
|
|
|
|
scaler.scale(loss).backward()
|
|
scaler.unscale_(optimizer)
|
|
utils.clip_grad_norm_(model.parameters(), max_norm=CONFIG['grad_clip_norm'])
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
scheduler.step()
|
|
|
|
total_train_loss += loss.item()
|
|
num_batches += 1
|
|
global_step += 1
|
|
|
|
writer.add_scalar('Train/Loss', loss.item(), global_step)
|
|
writer.add_scalar('Train/LearningRate', optimizer.param_groups[0]['lr'], global_step)
|
|
progress_bar.set_postfix({'epoch': f"{epoch}/{CONFIG['num_epochs']}",'loss': f'{loss.item():.4f}', 'step': global_step})
|
|
|
|
train_avg_loss = total_train_loss / num_batches
|
|
return train_avg_loss, global_step
|
|
|
|
def validate(model: ASRModel, dataloader: DataLoader, criterion: CTCLoss, device: device, tokenizer: ASRTokenizer, writer: SummaryWriter, global_step: int) -> tuple[float, float, float]:
|
|
model.eval()
|
|
total_loss = 0
|
|
total_cer = 0
|
|
total_wer = 0
|
|
num_samples = 0
|
|
num_batches = 0
|
|
examples = []
|
|
|
|
with no_grad():
|
|
progress_bar = tqdm(dataloader, desc="Validate", leave=False)
|
|
for batch in progress_bar:
|
|
batch: Batch
|
|
mel_specs = batch['mel_specs'].to(device)
|
|
targets = batch['targets'].to(device)
|
|
mel_lengths = batch['mel_lengths'].to(device)
|
|
target_lengths = batch['target_lengths'].to(device)
|
|
|
|
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
|
|
log_probs, lengths = model(mel_specs=mel_specs, mel_lengths=mel_lengths)
|
|
log_probs: Tensor
|
|
log_probs_ctc = log_probs.permute(1, 0, 2)
|
|
loss: Tensor = criterion(log_probs=log_probs_ctc, targets=targets, input_lengths=lengths, target_lengths=target_lengths)
|
|
|
|
total_loss += loss.item()
|
|
|
|
for i in range(log_probs.shape[0]):
|
|
pred_text = tokenizer.ctc_greedy_decode(log_probs=log_probs[i])
|
|
true_text = batch['target_texts'][i]
|
|
cer = calculate_cer(pred=pred_text, target=true_text)
|
|
wer = calculate_wer(pred=pred_text, target=true_text)
|
|
total_cer += cer
|
|
total_wer += wer
|
|
num_samples += 1
|
|
|
|
if len(examples) < 3:
|
|
examples.append((true_text, pred_text, cer, wer))
|
|
|
|
num_batches += 1
|
|
|
|
avg_loss = total_loss / num_batches
|
|
avg_cer = total_cer / num_samples
|
|
avg_wer = total_wer / num_samples
|
|
|
|
writer.add_scalar('Val/Loss', avg_loss, global_step)
|
|
writer.add_scalar('Val/CER', avg_cer, global_step)
|
|
writer.add_scalar('Val/WER', avg_wer, global_step)
|
|
|
|
for index, (true_text, pred_text, cer, wer) in enumerate(examples):
|
|
writer.add_text(f'Val/Example_{index}', f'True: {true_text}\nPred: {pred_text}\nCER: {cer:.4f} | WER: {wer:.4f}', global_step)
|
|
|
|
model.train()
|
|
return avg_loss, avg_cer, avg_wer
|
|
|
|
def save_checkpoint(model: ASRModel, optimizer: Optimizer, scheduler: OneCycleLR, global_step: int, epoch: int, train_loss: float, val_loss: float, cer: float, wer: float, save_path: Path):
|
|
checkpoint = {
|
|
'model_state_dict': model.state_dict(),
|
|
'optimizer_state_dict': optimizer.state_dict(),
|
|
'scheduler_state_dict': scheduler.state_dict(),
|
|
'global_step': global_step,
|
|
'train_loss': train_loss,
|
|
'val_loss': val_loss,
|
|
'epoch': epoch,
|
|
'cer': cer,
|
|
'wer': wer,
|
|
}
|
|
torch.save(checkpoint, save_path)
|
|
|
|
def find_latest_checkpoint(checkpoint_dir: Path) -> Path | None:
|
|
checkpoints = sorted(checkpoint_dir.glob('checkpoint_epoch_*.pt'), key=lambda p: int(p.stem.split('_')[-1]))
|
|
return checkpoints[-1] if checkpoints else None
|
|
|
|
def load_checkpoint(file_path: Path, model: ASRModel, optimizer: AdamW, scheduler: OneCycleLR):
|
|
checkpoint = torch.load(file_path, weights_only=False)
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
|
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
|
|
|
|
epoch = checkpoint['epoch']
|
|
global_step = checkpoint['global_step']
|
|
best_cer = checkpoint['cer']
|
|
best_wer = checkpoint['wer']
|
|
return epoch, global_step, best_cer, best_wer
|
|
|
|
def main():
|
|
workspace_dir = Path(__file__).parent.parent
|
|
device = torch.device('cuda:0')
|
|
tokenizer = ASRTokenizer(workspace_dir / 'config/asr_vocab.json')
|
|
|
|
# ============ 创建数据加载器 ============
|
|
train_loader = create_dataloader(
|
|
tsv_path=workspace_dir / '.data/ug/train.tsv',
|
|
audio_dir=workspace_dir / '.data/ug/clips',
|
|
tokenizer=tokenizer,
|
|
batch_size=CONFIG['batch_size'],
|
|
shuffle=True,
|
|
augment=True
|
|
)
|
|
|
|
val_loader = create_dataloader(
|
|
tsv_path=workspace_dir / '.data/ug/dev.tsv',
|
|
audio_dir=workspace_dir / '.data/ug/clips',
|
|
tokenizer=tokenizer,
|
|
batch_size=CONFIG['batch_size'],
|
|
shuffle=False,
|
|
augment=False
|
|
)
|
|
|
|
# ============ 初始化模型 ============
|
|
model = ASRModel(
|
|
vocab_size=tokenizer.vocab_size(),
|
|
input_dim=CONFIG['input_dim'],
|
|
num_heads=CONFIG['num_heads'],
|
|
ffn_dim=CONFIG['ffn_dim'],
|
|
num_layers=CONFIG['num_layers'],
|
|
dropout=CONFIG['dropout'],
|
|
).to(device)
|
|
print(f"🤖 模型参数量: {model.get_num_params() / 1e6:.2f}M")
|
|
|
|
|
|
criterion = CTCLoss(blank=tokenizer.get_special_token_id('<BLANK>'), zero_infinity=True)
|
|
optimizer = AdamW(model.parameters(), lr=CONFIG['learning_rate'], weight_decay=CONFIG['weight_decay'], foreach=True)
|
|
|
|
# 使用 OneCycleLR：自动处理 warmup 和衰减
|
|
scheduler = OneCycleLR(
|
|
optimizer,
|
|
max_lr=CONFIG['learning_rate'],
|
|
epochs=CONFIG['num_epochs'],
|
|
steps_per_epoch=len(train_loader),
|
|
pct_start=0.1, # 前 10% 步数用于 warmup
|
|
anneal_strategy='cos',
|
|
div_factor=25.0, # 初始 lr = max_lr / 25
|
|
final_div_factor=1e4, # 最终 lr = max_lr / 10000
|
|
)
|
|
scaler = GradScaler()
|
|
|
|
# ============ TensorBoard ============
|
|
log_dir = workspace_dir / 'runs' / datetime.now().strftime('%Y%m%d_%H%M%S')
|
|
writer = SummaryWriter(log_dir)
|
|
|
|
config_text = '\n'.join([f'{k}: {v}' for k, v in CONFIG.items()])
|
|
writer.add_text('Config', config_text, 0)
|
|
print(f"📊 TensorBoard 日志: {log_dir}")
|
|
|
|
checkpoint_dir = workspace_dir / '.checkpoints'
|
|
checkpoint_dir.mkdir(exist_ok=True)
|
|
|
|
# ============ 训练循环 ============
|
|
best_cer = float('inf')
|
|
best_wer = float('inf')
|
|
patience_counter = 0
|
|
global_step = 0
|
|
start_epoch = 0
|
|
|
|
latest = find_latest_checkpoint(checkpoint_dir)
|
|
if latest:
|
|
start_epoch, global_step, best_cer, best_wer = load_checkpoint(latest, model, optimizer, scheduler)
|
|
start_epoch += 1 # 从下一个 epoch 开始
|
|
print(f"✅ 恢复训练: 从 epoch={start_epoch} 开始, global_step={global_step}")
|
|
else:
|
|
start_epoch = 0
|
|
print("🆕 从头开始训练\n")
|
|
|
|
for epoch in range(start_epoch, CONFIG['num_epochs']):
|
|
train_loss, global_step = train_one_epoch(
|
|
model=model,
|
|
dataloader=train_loader,
|
|
criterion=criterion,
|
|
optimizer=optimizer,
|
|
scheduler=scheduler,
|
|
device=device,
|
|
epoch=epoch,
|
|
writer=writer,
|
|
global_step=global_step,
|
|
scaler=scaler,
|
|
)
|
|
|
|
val_loss, val_cer, val_wer = validate(model=model, dataloader=val_loader, criterion=criterion, device=device, tokenizer=tokenizer, writer=writer, global_step=global_step)
|
|
print(f"\n📊 Step {global_step} | Val Loss: {val_loss:.4f} | Val CER: {val_cer:.4f} | Val WER: {val_wer:.4f}")
|
|
|
|
# 保存常规 checkpoint
|
|
checkpoint_path = checkpoint_dir / f'checkpoint_epoch_{epoch}.pt'
|
|
save_checkpoint(model=model, optimizer=optimizer, scheduler=scheduler, epoch=epoch, global_step=global_step, train_loss=train_loss, val_loss=val_loss, cer=val_cer, wer=val_wer, save_path=checkpoint_path)
|
|
|
|
# 分别保存最佳 CER 和 WER 模型
|
|
improved = False
|
|
if val_cer < best_cer:
|
|
best_cer = val_cer
|
|
best_cer_path = checkpoint_dir / 'best_cer_model.pt'
|
|
save_checkpoint(model=model, optimizer=optimizer, scheduler=scheduler, epoch=epoch, global_step=global_step, train_loss=train_loss, val_loss=val_loss, cer=val_cer, wer=val_wer, save_path=best_cer_path)
|
|
improved = True
|
|
|
|
if val_wer < best_wer:
|
|
best_wer = val_wer
|
|
best_wer_path = checkpoint_dir / 'best_wer_model.pt'
|
|
save_checkpoint(model=model, optimizer=optimizer, scheduler=scheduler, epoch=epoch, global_step=global_step, train_loss=train_loss, val_loss=val_loss, cer=val_cer, wer=val_wer, save_path=best_wer_path)
|
|
improved = True
|
|
|
|
# Early Stopping 逻辑
|
|
if improved:
|
|
patience_counter = 0
|
|
else:
|
|
patience_counter += 1
|
|
print(f"⚠️ 验证指标未改善，patience: {patience_counter}/{CONFIG['early_stopping_patience']}")
|
|
|
|
# 删除旧的 checkpoint（保留最近3个）
|
|
old_checkpoints = sorted(checkpoint_dir.glob('checkpoint_epoch_*.pt'), key=lambda p: int(p.stem.split('_')[-1]))
|
|
for old in old_checkpoints[:-3]:
|
|
old.unlink()
|
|
|
|
cuda.empty_cache()
|
|
|
|
writer.add_scalar('Val/EpochLoss', val_loss, epoch)
|
|
writer.add_scalar('Train/EpochLoss', train_loss, epoch)
|
|
print(f"✅ Epoch {epoch} 完成 | Train Avg Loss: {train_loss:.4f} | Val Avg Loss: {val_loss:.4f} | Best CER: {best_cer:.4f} | Best WER: {best_wer:.4f}\n")
|
|
|
|
# Early Stopping 检查
|
|
if patience_counter >= CONFIG['early_stopping_patience']:
|
|
print(f"\n🛑 Early Stopping: 验证指标连续 {CONFIG['early_stopping_patience']} 次没有改善")
|
|
print(f"🏆 最佳 CER: {best_cer:.4f}")
|
|
print(f"🏆 最佳 WER: {best_wer:.4f}")
|
|
break
|
|
|
|
writer.close()
|
|
print(f"\n{'='*60}")
|
|
print(f"🎉 训练完成！")
|
|
print(f"🏆 最佳 CER: {best_cer:.4f}")
|
|
print(f"🏆 最佳 WER: {best_wer:.4f}")
|
|
print(f"📊 TensorBoard: tensorboard --logdir={workspace_dir / 'runs'}")
|
|
print(f"{'='*60}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |