File size: 6,959 Bytes
ea101a8 |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import numpy as np
import time
from tqdm import tqdm
from optimizers import build_optimizer
def train_aligner(config, accelerator, train_dataloader, val_dataloader, device, log_dir, epochs=100):
# Create model
aligner = AlignerModel().to(device)
# Define loss function
forward_sum_loss = ForwardSumLoss()
# Setup optimizer
scheduler_params = {
"max_lr": float(config['optimizer_params'].get('lr', 5e-4)),
"pct_start": float(config['optimizer_params'].get('pct_start', 0.0)),
"epochs": epochs,
"steps_per_epoch": len(train_dataloader),
}
optimizer, scheduler = build_optimizer(
{"params": aligner.parameters(), "optimizer_params":{}, "scheduler_params": scheduler_params})
# Setup TensorBoard writer
writer = SummaryWriter(log_dir=log_dir)
# Create directories for model checkpoints
os.makedirs(os.path.join(log_dir, 'checkpoints'), exist_ok=True)
# Track best validation loss
best_val_loss = float('inf')
# Loss weights
fwd_sum_loss_weight = config.get('fwd_sum_loss_weight', 1.0)
# Training loop
for epoch in range(1, epochs + 1):
aligner.train()
train_losses = []
train_fwd_losses = []
start_time = time.time()
# Training phase
pbar = tqdm(train_dataloader, desc=f"Epoch {epoch}/{epochs} [Train]")
for i, batch in enumerate(pbar):
batch = [b.to(device) for b in batch]
text_input, text_input_length, mel_input, mel_input_length, attn_prior = batch
# Forward pass
attn_soft, attn_logprob = aligner(spec=mel_input,
spec_len=mel_input_length,
text=text_input,
text_len=text_input_length,
attn_prior=attn_prior)
# Calculate loss
loss = forward_sum_loss(attn_logprob=attn_logprob,
in_lens=text_input_length,
out_lens=mel_input_length)
# Backward pass and optimization
optimizer.zero_grad()
loss.backward()
# Optional gradient clipping
grad_norm = nn.utils.clip_grad_norm_(aligner.parameters(), config.get('grad_clip', 5.0))
optimizer.step()
if scheduler is not None:
scheduler.step()
# Log to TensorBoard
global_step = (epoch - 1) * len(train_dataloader) + i
writer.add_scalar('train/total_loss', loss.item(), global_step)
writer.add_scalar('train/grad_norm', grad_norm, global_step)
# Update progress bar
train_losses.append(loss.item())
train_fwd_losses.append(loss.item())
# Update the progress bar description
pbar.set_description(f"Epoch {epoch}/{epochs} [Train] Loss: {loss.item():.4f}")
# Calculate average training loss for this epoch
avg_train_loss = sum(train_losses) / len(train_losses)
# Validation phase
aligner.eval()
val_losses = []
with torch.no_grad():
for batch in tqdm(val_dataloader, desc=f"Epoch {epoch}/{epochs} [Val]"):
batch = [b.to(device) for b in batch]
text_input, text_input_length, mel_input, mel_input_length, attn_prior = batch
# Forward pass
attn_soft, attn_logprob = aligner(spec=mel_input,
spec_len=mel_input_length,
text=text_input,
text_len=text_input_length,
attn_prior=attn_prior)
# Calculate loss
val_loss = forward_sum_loss(attn_logprob=attn_logprob,
in_lens=text_input_length,
out_lens=mel_input_length)
val_losses.append(val_loss.item())
# Calculate average validation loss
avg_val_loss = sum(val_losses) / len(val_losses)
# Log to TensorBoard
writer.add_scalar('epoch/train_loss', avg_train_loss, epoch)
writer.add_scalar('epoch/val_loss', avg_val_loss, epoch)
# Save model if it's the best so far
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
torch.save({
'epoch': epoch,
'model_state_dict': aligner.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'train_loss': avg_train_loss,
'val_loss': avg_val_loss,
}, os.path.join(log_dir, 'checkpoints', 'best_model.pt'))
# Save checkpoint every N epochs
if epoch % config.get('save_every', 10) == 0:
torch.save({
'epoch': epoch,
'model_state_dict': aligner.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'train_loss': avg_train_loss,
'val_loss': avg_val_loss,
}, os.path.join(log_dir, 'checkpoints', f'checkpoint_epoch_{epoch}.pt'))
# Print summary for this epoch
epoch_time = time.time() - start_time
print(f"Epoch {epoch}/{epochs} completed in {epoch_time:.2f}s | "
f"Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
# Plot and save attention matrices for visualization
if epoch % config.get('plot_every', 10) == 0:
plot_attention_matrices(aligner, val_dataloader, device,
os.path.join(log_dir, 'attention_plots', f'epoch_{epoch}'),
num_samples=4)
writer.close()
print(f"Training completed. Best validation loss: {best_val_loss:.4f}")
return aligner
# Main execution
if __name__ == "__main__":
def length_to_mask(lengths):
mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
mask = torch.gt(mask+1, lengths.unsqueeze(1))
return mask
# Assuming these variables are defined in your main script
train_aligner(
config=config,
train_dataloader=train_dataloader,
val_dataloader=val_dataloader,
device=device,
log_dir=log_dir,
epochs=epoch
) |