152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
import torch
|
|
import torch.optim as optim
|
|
import matplotlib
|
|
matplotlib.use('Agg')
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import os
|
|
from torch.utils.data import DataLoader
|
|
|
|
# Imports
|
|
from hsmm_model import BatchedGaussianHSMM
|
|
from hsmm_inference import viterbi_decode
|
|
from real_data import get_real_dataloaders # <--- NEW IMPORT
|
|
|
|
# --- CONFIGURATION ---
|
|
CONFIG = {
|
|
# Path to your file
|
|
"DATA_PATH": "/u/schmitt/experiments/2025_10_02_unsup_asr_shared_enc/work/i6_experiments/users/schmitt/experiments/exp2025_10_02_shared_enc/librispeech/data/audio_preprocessing/Wav2VecUFeaturizeAudioJob.mkGrrp0YWy8y/output/audio_features/train.npy",
|
|
|
|
"N_STATES": 50, # Increased for real speech (approx 40-50 phonemes)
|
|
"PCA_DIM": 30, # Reduce 512 -> 30 dimensions
|
|
"MAX_DUR": 30, # Max duration in frames (30 * 20ms = 600ms)
|
|
"LR": 0.01, # Slightly lower LR for real data
|
|
"EPOCHS": 20,
|
|
"BATCH_SIZE": 64, # 64 chunks of 3 seconds each
|
|
"CROP_LEN": 150, # Training window (150 frames = 3 seconds)
|
|
|
|
"CHECKPOINT_PATH": "hsmm_librispeech.pth",
|
|
"RESUME": False
|
|
}
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"--- Using device: {device} ---")
|
|
|
|
def save_checkpoint(model, optimizer, epoch, loss, path):
|
|
torch.save({
|
|
'epoch': epoch,
|
|
'model_state_dict': model.state_dict(),
|
|
'optimizer_state_dict': optimizer.state_dict(),
|
|
'loss': loss,
|
|
}, path)
|
|
|
|
def load_checkpoint(model, optimizer, path):
|
|
if os.path.exists(path):
|
|
print(f"Loading checkpoint from {path}...")
|
|
checkpoint = torch.load(path)
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
|
return checkpoint['epoch'] + 1
|
|
return 0
|
|
|
|
def train():
|
|
# 1. Load Real Data
|
|
print("--- 1. Loading Real Data ---")
|
|
train_ds, val_ds = get_real_dataloaders(
|
|
CONFIG["DATA_PATH"],
|
|
batch_size=CONFIG["BATCH_SIZE"],
|
|
crop_len=CONFIG["CROP_LEN"],
|
|
pca_dim=CONFIG["PCA_DIM"]
|
|
)
|
|
|
|
# Loader for Training (Batched, Cropped)
|
|
train_loader = DataLoader(
|
|
train_ds,
|
|
batch_size=CONFIG["BATCH_SIZE"],
|
|
shuffle=True,
|
|
num_workers=4,
|
|
pin_memory=True,
|
|
drop_last=True # Avoid partial batch issues
|
|
)
|
|
|
|
# 2. Init Model
|
|
print("--- 2. Initializing Model ---")
|
|
model = BatchedGaussianHSMM(CONFIG["N_STATES"], CONFIG["PCA_DIM"], CONFIG["MAX_DUR"])
|
|
model.to(device)
|
|
|
|
# Smart Init (using PCA-reduced data from the dataset)
|
|
if not CONFIG["RESUME"] and not os.path.exists(CONFIG["CHECKPOINT_PATH"]):
|
|
print("--- 2b. Running Smart Initialization ---")
|
|
# Grab a batch to init means
|
|
init_batch = next(iter(train_loader)).to(device) # (B, T, D)
|
|
flat_data = init_batch.view(-1, CONFIG["PCA_DIM"])
|
|
|
|
# Pick random frames
|
|
indices = torch.randperm(flat_data.size(0))[:CONFIG["N_STATES"]]
|
|
model.means.data.copy_(flat_data[indices])
|
|
print("Means initialized.")
|
|
|
|
optimizer = optim.Adam(model.parameters(), lr=CONFIG["LR"])
|
|
start_epoch = 0
|
|
if CONFIG["RESUME"]:
|
|
start_epoch = load_checkpoint(model, optimizer, CONFIG["CHECKPOINT_PATH"])
|
|
|
|
# 3. Training Loop
|
|
print(f"--- 3. Training Loop ---")
|
|
|
|
for epoch in range(start_epoch, CONFIG["EPOCHS"]):
|
|
total_loss = 0
|
|
model.train()
|
|
|
|
for batch_idx, batch_data in enumerate(train_loader):
|
|
batch_data = batch_data.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
loss = model(batch_data)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
total_loss += loss.item()
|
|
|
|
if batch_idx % 50 == 0:
|
|
print(f"Epoch {epoch} | Batch {batch_idx} | Loss {loss.item():.4f}")
|
|
|
|
avg_loss = total_loss / len(train_loader)
|
|
print(f"Epoch {epoch:02d} DONE | Avg NLL: {avg_loss:.4f}")
|
|
save_checkpoint(model, optimizer, epoch, avg_loss, CONFIG["CHECKPOINT_PATH"])
|
|
|
|
# 4. Visualization (Using Validation Set - Full Length)
|
|
print("\n--- 4. Visualizing Inference on Real Audio ---")
|
|
|
|
# Grab the first file from validation set (Index 0)
|
|
# val_ds[0] returns (Time, Dim) -> add batch dim -> (1, T, D)
|
|
test_seq = val_ds[0].unsqueeze(0).to(device)
|
|
|
|
# Run Inference
|
|
path = viterbi_decode(model, test_seq)
|
|
|
|
# Move to CPU for plotting
|
|
raw_data = test_seq.squeeze(0).cpu().numpy()
|
|
|
|
# Plotting
|
|
fig, ax = plt.subplots(2, 1, figsize=(15, 6), sharex=True)
|
|
|
|
# Plot PCA Features
|
|
ax[0].imshow(raw_data.T, aspect='auto', cmap='viridis', interpolation='nearest')
|
|
ax[0].set_title(f"PCA Reduced Features ({CONFIG['PCA_DIM']} Dim)")
|
|
ax[0].set_ylabel("PCA Dim")
|
|
|
|
# Plot States
|
|
path_img = np.array(path)[np.newaxis, :]
|
|
ax[1].imshow(path_img, aspect='auto', cmap='tab20', interpolation='nearest')
|
|
ax[1].set_title("Inferred Phoneme States")
|
|
ax[1].set_ylabel("State ID")
|
|
ax[1].set_xlabel("Time (Frames)")
|
|
|
|
plt.tight_layout()
|
|
plt.savefig("librispeech_result.png")
|
|
print("Saved librispeech_result.png")
|
|
|
|
if __name__ == "__main__":
|
|
train()
|