speech_baseline_cnn.py
Speech Commands CNN baseline - 919 lines.
View on GitHub (speech/speech_baseline_cnn.py).
Source
Section titled “Source”"""
Speech Commands Baseline CNN (non-SNN)
Pipeline: Waveform -> 8kHz -> Spectral Gate -> STFT log-magnitude -> CNN -> Classification
This serves as a conventional baseline for comparison with DyNED+SNN approaches.
Uses the same preprocessing (STFT) but feeds log-magnitude spectrograms into a
standard convolutional neural network instead of spiking neurons.
"""
import argparse
import os
import sys
import time
import zipfile
import matplotlib
try:
matplotlib.use("Agg")
except Exception:
pass
import matplotlib.pyplot as plt
import numpy as np
import torch
torch.backends.nnpack.enabled = False
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
from torch.utils.data import DataLoader
from torchaudio.datasets import SPEECHCOMMANDS
from torchaudio.transforms import Resample
try:
from noisereduce.torchgate import TorchGate
HAS_TORCHGATE = True
except ImportError:
HAS_TORCHGATE = False
print("Warning: noisereduce not installed - skipping spectral gate cleaning")
try:
from sklearn.manifold import TSNE
HAS_TSNE = True
except ImportError:
HAS_TSNE = False
try:
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
_SCRIPT_DIR = os.getcwd()
OUTPUT_DIR = os.path.join(_SCRIPT_DIR, "speech_baseline_cnn_output")
os.makedirs(OUTPUT_DIR, exist_ok=True)
SPEECH_LABELS = [
"backward",
"bed",
"bird",
"cat",
"dog",
"down",
"eight",
"five",
"follow",
"forward",
"four",
"go",
"happy",
"house",
"learn",
"left",
"marvin",
"nine",
"no",
"off",
"on",
"one",
"right",
"seven",
"sheila",
"six",
"stop",
"three",
"tree",
"two",
"up",
"visual",
"wow",
"yes",
"zero",
]
# =============================================================================
# Dataset - STFT log-magnitude (no DyNED)
# =============================================================================
class STFTSpeechDataset(torch.utils.data.Dataset):
"""Speech Commands dataset with pre-computed STFT log-magnitude cached to disk."""
def __init__(
self,
subset="training",
n_fft=1024,
hop_length=160,
add_noise=True,
noise_level=0.01,
target_sr=8000,
cache_dir="../assets",
):
self.n_fft = n_fft
self.hop_length = hop_length
self.is_training = subset == "training"
self.add_noise = add_noise and self.is_training
self.noise_level = noise_level
self.target_sr = target_sr
self.labels = SPEECH_LABELS
self.label_to_idx = {label: idx for idx, label in enumerate(self.labels)}
cache_path = (
Path(cache_dir)
/ "speech_baseline_cnn"
/ f"stft_cache_{subset}_sr{target_sr}_nfft{n_fft}_hop{hop_length}.pt"
)
if cache_path.exists():
print(f"Loading cached spectrograms from {cache_path}...")
try:
cache = torch.load(cache_path, weights_only=True)
self.spectrograms = cache["spectrograms"]
self.label_indices = cache["labels"]
print(f"Loaded {len(self.spectrograms)} samples ({self.spectrograms.shape})")
return
except (RuntimeError, EOFError, zipfile.BadZipFile) as e:
print(f"Corrupted cache file, rebuilding: {e}")
cache_path.unlink(missing_ok=True)
print(f"Cache not found - pre-computing STFT spectrograms for {subset}...")
self._build_cache(subset, cache_path, cache_dir)
def _build_cache(self, subset, cache_path, cache_dir):
cache_path.parent.mkdir(parents=True, exist_ok=True)
raw_dataset = SPEECHCOMMANDS(cache_dir, download=True, subset=subset)
resampler = Resample(orig_freq=16000, new_freq=self.target_sr)
torchgate = TorchGate(sr=self.target_sr, nonstationary=False) if HAS_TORCHGATE else None
all_specs = []
all_labels = []
total = len(raw_dataset)
for i in range(total):
waveform, sample_rate, label, _, _ = raw_dataset[i]
if waveform.shape[-1] < 16000:
waveform = F.pad(waveform, (0, 16000 - waveform.shape[-1]))
elif waveform.shape[-1] > 16000:
waveform = waveform[..., :16000]
waveform = resampler(waveform)
if torchgate is not None:
with torch.no_grad():
waveform = torchgate(waveform)
if waveform.shape[-1] < self.target_sr:
waveform = F.pad(waveform, (0, self.target_sr - waveform.shape[-1]))
elif waveform.shape[-1] > self.target_sr:
waveform = waveform[..., : self.target_sr]
# Normalize
mean = waveform.mean()
std = waveform.std()
waveform = (waveform - mean) / (std + 1e-8)
# Pre-emphasis
wf_pad = F.pad(waveform, (1, 0))
waveform = wf_pad[..., 1:] - 0.97 * wf_pad[..., :-1]
# STFT -> log-magnitude (same as DyNED pipeline, but stop here)
window = torch.hann_window(self.n_fft)
stft = torch.stft(
waveform.squeeze(0),
n_fft=self.n_fft,
hop_length=self.hop_length,
window=window,
return_complex=True,
center=True,
pad_mode="reflect",
)
magnitude = torch.abs(stft)
log_mag = torch.log1p(magnitude) # [n_freq, n_time]
all_specs.append(log_mag)
all_labels.append(self.label_to_idx[label])
if (i + 1) % 1000 == 0 or (i + 1) == total:
pct = (i + 1) / total * 100
print(f" [{subset}] Processed {i + 1}/{total} samples ({pct:.1f}%)")
self.spectrograms = torch.stack(all_specs)
self.label_indices = torch.tensor(all_labels, dtype=torch.long)
torch.save({"spectrograms": self.spectrograms, "labels": self.label_indices}, cache_path)
size_mb = cache_path.stat().st_size / (1024 * 1024)
print(f"Cached {len(self.spectrograms)} spectrograms to {cache_path} ({size_mb:.1f} MB)")
def __len__(self):
return len(self.spectrograms)
def __getitem__(self, n):
spec = self.spectrograms[n]
label_idx = self.label_indices[n]
if self.is_training:
# Time shift augmentation
max_shift = spec.shape[1] // 10
shift = torch.randint(-max_shift, max_shift + 1, (1,)).item()
if shift != 0:
spec = torch.roll(spec, shifts=shift, dims=1)
# Frequency masking
if torch.rand(1).item() < 0.3:
n_freq = spec.shape[0]
mask_width = torch.randint(1, n_freq // 10 + 1, (1,)).item()
mask_start = torch.randint(0, n_freq - mask_width, (1,)).item()
spec = spec.clone()
spec[mask_start : mask_start + mask_width, :] = 0
# Time masking
if torch.rand(1).item() < 0.3:
n_time = spec.shape[1]
mask_width = torch.randint(1, n_time // 10 + 1, (1,)).item()
mask_start = torch.randint(0, n_time - mask_width, (1,)).item()
spec = spec.clone()
spec[:, mask_start : mask_start + mask_width] = 0
# Additive Gaussian noise
if self.add_noise:
spec = spec + self.noise_level * torch.randn_like(spec)
return spec, label_idx
# =============================================================================
# CNN Model
# =============================================================================
class SpeechCNN(nn.Module):
"""CNN for speech command classification on STFT spectrograms.
Architecture:
Input [1, n_freq, n_time]
-> Conv2d(1, 64, 3) + BN + ReLU + MaxPool(2)
-> Conv2d(64, 128, 3) + BN + ReLU + MaxPool(2)
-> Conv2d(128, 256, 3) + BN + ReLU + MaxPool(2)
-> Conv2d(256, 256, 3) + BN + ReLU + AdaptiveAvgPool(1)
-> FC(256, hidden) + ReLU + Dropout
-> FC(hidden, 35)
"""
def __init__(self, hidden_size=256, num_outputs=35, dropout=0.3):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d(1),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(256, hidden_size),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(hidden_size, num_outputs),
)
def forward(self, x):
"""
Args:
x: [batch, n_freq, n_time] - log-magnitude spectrogram
Returns:
logits: [batch, num_outputs]
"""
x = x.unsqueeze(1) # [B, 1, F, T]
x = self.features(x)
x = self.classifier(x)
return x
# =============================================================================
# Data Setup
# =============================================================================
def setup_dataloaders(batch_size=512, num_workers=4, n_fft=1024, hop_length=160, target_sr=8000):
train_dataset = STFTSpeechDataset(
subset="training",
n_fft=n_fft,
hop_length=hop_length,
target_sr=target_sr,
)
test_dataset = STFTSpeechDataset(
subset="testing",
n_fft=n_fft,
hop_length=hop_length,
add_noise=False,
target_sr=target_sr,
)
loader_kwargs = dict(pin_memory=True)
if num_workers > 0:
loader_kwargs["persistent_workers"] = False
loader_kwargs["prefetch_factor"] = 4
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
**loader_kwargs,
)
test_loader = DataLoader(
test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
**loader_kwargs,
)
return train_loader, test_loader
# =============================================================================
# Training Loop
# =============================================================================
def train_network(
net, train_loader, test_loader, num_epochs=250, device="cuda", lr=1e-3, weight_decay=1e-4, trial=None
):
optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode="max",
factor=0.7,
patience=5,
min_lr=1e-6,
)
scaler = torch.amp.GradScaler("cuda") if device == "cuda" else None
best_acc = 0
metrics = {
"epoch": [],
"train_loss": [],
"test_accuracy": [],
"learning_rate": [],
"epoch_time": [],
"per_class_accuracy": [],
}
print(f"Training on device: {device}")
print(f"Batch size: {train_loader.batch_size}")
if device == "cuda":
print(f"GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB allocated")
for epoch in range(num_epochs):
epoch_start = time.time()
net.train()
epoch_loss = 0.0
num_batches = 0
running_loss = 0.0
for i, (data, targets) in enumerate(train_loader):
data = data.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)
optimizer.zero_grad()
if scaler is not None:
with torch.amp.autocast("cuda"):
output = net(data)
loss = F.cross_entropy(output, targets)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
else:
output = net(data)
loss = F.cross_entropy(output, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
optimizer.step()
batch_loss = loss.item()
running_loss += batch_loss
epoch_loss += batch_loss
num_batches += 1
if i % 50 == 49:
avg_loss = running_loss / 50
print(f"Epoch {epoch + 1}, Batch {i + 1}: Loss = {avg_loss:.4f}", end="")
if device == "cuda":
print(f" | GPU: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
else:
print()
running_loss = 0.0
epoch_time = time.time() - epoch_start
avg_epoch_loss = epoch_loss / num_batches
current_lr = optimizer.param_groups[0]["lr"]
eval_result = evaluate(net, test_loader, device)
test_acc = eval_result["accuracy"]
scheduler.step(test_acc)
metrics["epoch"].append(epoch + 1)
metrics["train_loss"].append(avg_epoch_loss)
metrics["test_accuracy"].append(test_acc)
metrics["learning_rate"].append(current_lr)
metrics["epoch_time"].append(epoch_time)
metrics["per_class_accuracy"].append(eval_result["per_class_accuracy"])
print(
f"Epoch {epoch + 1}: Acc = {test_acc:.4f} | Loss = {avg_epoch_loss:.4f} | "
f"LR = {current_lr:.6f} | {epoch_time:.1f}s"
)
if device == "cuda":
torch.cuda.empty_cache()
if test_acc > best_acc:
best_acc = test_acc
torch.save(
{
"epoch": epoch,
"model_state_dict": net.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"accuracy": test_acc,
},
os.path.join(OUTPUT_DIR, "best_baseline_cnn_model.pth"),
)
if trial is not None:
import optuna
trial.report(test_acc, epoch)
if trial.should_prune():
raise optuna.TrialPruned()
return metrics
# =============================================================================
# Evaluation
# =============================================================================
def evaluate(net, test_loader, device, collect_representations=False):
net.eval()
correct = 0
total = 0
all_predictions = []
all_targets = []
all_representations = []
num_classes = len(SPEECH_LABELS)
per_class_correct = np.zeros(num_classes)
per_class_total = np.zeros(num_classes)
with torch.no_grad():
for data, targets in test_loader:
data = data.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)
output = net(data)
predicted_classes = output.argmax(dim=1)
correct += (predicted_classes == targets).sum().item()
total += targets.size(0)
all_predictions.extend(predicted_classes.cpu().numpy())
all_targets.extend(targets.cpu().numpy())
for cls in range(num_classes):
mask = targets == cls
per_class_correct[cls] += (predicted_classes[mask] == targets[mask]).sum().item()
per_class_total[cls] += mask.sum().item()
if collect_representations:
all_representations.append(output.cpu())
accuracy = correct / total if total > 0 else 0.0
per_class_acc = per_class_correct / (per_class_total + 1e-8)
all_predictions = np.array(all_predictions)
all_targets = np.array(all_targets)
cm = np.zeros((num_classes, num_classes), dtype=np.int64)
for pred_cls, true_cls in zip(all_predictions, all_targets):
cm[true_cls][pred_cls] += 1
result = {
"accuracy": accuracy,
"per_class_accuracy": {SPEECH_LABELS[i]: float(per_class_acc[i]) for i in range(num_classes)},
"confusion_matrix": cm,
"predictions": all_predictions,
"targets": all_targets,
}
if collect_representations and all_representations:
result["representations"] = torch.cat(all_representations, dim=0).numpy()
return result
# =============================================================================
# Visualizations
# =============================================================================
def analyze_training_metrics(metrics):
losses = np.array(metrics["train_loss"])
accuracies = np.array(metrics["test_accuracy"])
epochs = np.array(metrics["epoch"])
# Save CSV
header_parts = ["epoch", "train_loss", "test_accuracy", "learning_rate", "epoch_time"]
header_parts.extend(f"acc_{cls}" for cls in SPEECH_LABELS)
header = ",".join(header_parts)
rows = []
for i in range(len(metrics["epoch"])):
row = [
metrics["epoch"][i],
metrics["train_loss"][i],
metrics["test_accuracy"][i],
metrics["learning_rate"][i],
metrics["epoch_time"][i],
]
pca = metrics["per_class_accuracy"][i]
for cls in SPEECH_LABELS:
row.append(pca.get(cls, 0))
rows.append(row)
filepath = os.path.join(OUTPUT_DIR, "training_metrics.csv")
np.savetxt(filepath, np.array(rows), delimiter=",", header=header, comments="")
print(f"Saved metrics to {filepath}")
best_epoch = np.argmax(accuracies)
print(f"\nBest epoch: {best_epoch + 1} - accuracy: {accuracies[best_epoch]:.4f}")
# Loss/accuracy plot
fig, ax1 = plt.subplots(figsize=(12, 6))
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Loss", color="tab:blue")
ax1.plot(epochs, losses, color="tab:blue", label="Training Loss", alpha=0.8)
ax1.tick_params(axis="y", labelcolor="tab:blue")
ax2 = ax1.twinx()
ax2.set_ylabel("Accuracy (%)", color="tab:red")
ax2.plot(epochs, accuracies * 100, color="tab:red", label="Test Accuracy", alpha=0.8)
ax2.tick_params(axis="y", labelcolor="tab:red")
plt.title("Baseline CNN - Speech Commands Training Progress")
lines1, labels1 = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines1 + lines2, labels1 + labels2, loc="center right")
plt.savefig(os.path.join(OUTPUT_DIR, "training_progress.png"), dpi=150, bbox_inches="tight")
plt.close()
# LR schedule
plt.figure(figsize=(10, 4))
plt.plot(epochs, metrics["learning_rate"])
plt.xlabel("Epoch")
plt.ylabel("Learning Rate")
plt.title("Learning Rate Schedule")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "lr_schedule.png"), dpi=150)
plt.close()
def visualize_confusion_matrix_plot(cm):
num_classes = len(SPEECH_LABELS)
fig, ax = plt.subplots(figsize=(14, 12))
cm_normalized = cm.astype(float) / (cm.sum(axis=1, keepdims=True) + 1e-8)
im = ax.imshow(cm_normalized, interpolation="nearest", cmap="Blues")
ax.figure.colorbar(im, ax=ax)
ax.set(
xticks=range(num_classes),
yticks=range(num_classes),
xticklabels=SPEECH_LABELS,
yticklabels=SPEECH_LABELS,
ylabel="True Label",
xlabel="Predicted Label",
title="Confusion Matrix (Normalized)",
)
plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="anchor", fontsize=7)
plt.setp(ax.get_yticklabels(), fontsize=7)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "confusion_matrix.png"), dpi=150)
plt.close()
np.savetxt(
os.path.join(OUTPUT_DIR, "confusion_matrix.csv"),
cm,
delimiter=",",
fmt="%d",
header=",".join(SPEECH_LABELS),
comments="",
)
print("Saved confusion matrix")
def visualize_tsne(representations, targets):
if not HAS_TSNE:
return
max_samples = 5000
if len(representations) > max_samples:
indices = np.random.choice(len(representations), max_samples, replace=False)
representations = representations[indices]
targets = targets[indices]
print("Computing t-SNE...")
tsne = TSNE(n_components=2, random_state=42, perplexity=30)
embedded = tsne.fit_transform(representations)
plt.figure(figsize=(14, 12))
scatter = plt.scatter(embedded[:, 0], embedded[:, 1], c=targets, cmap="nipy_spectral", alpha=0.6, s=5)
cbar = plt.colorbar(scatter, ticks=range(len(SPEECH_LABELS)))
cbar.set_ticklabels(SPEECH_LABELS)
cbar.ax.tick_params(labelsize=6)
plt.title("t-SNE of Output Representations")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "tsne.png"), dpi=150)
plt.close()
print("Saved t-SNE")
# =============================================================================
# Main
# =============================================================================
def main(json_path=None):
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
np.random.seed(42)
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
n_fft = 1024
hop_length = 160
target_sr = 8000
num_outputs = 35
if json_path:
import json
with open(json_path) as f:
params = json.load(f)
print(f"Loaded hyperparameters from: {json_path}")
lr = params["lr"]
weight_decay = params.get("weight_decay", 1e-4)
hidden_size = params["hidden_size"]
dropout = params.get("dropout", 0.3)
batch_size = params.get("batch_size", 512)
else:
lr = 1e-3
weight_decay = 1e-4
hidden_size = 256
dropout = 0.3
batch_size = 512
num_epochs = 250
num_workers = min(8, os.cpu_count() - 2) if os.cpu_count() > 2 else 0
print("=" * 80)
print("Baseline CNN for Speech Commands")
print("=" * 80)
print(f"Device: {device_str.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Pipeline: Waveform -> 8kHz -> Spectral Gate -> STFT log-mag -> CNN")
print(f"Batch: {batch_size} | Hidden: {hidden_size} | Dropout: {dropout}")
print(f"LR: {lr} | Weight decay: {weight_decay}")
print("=" * 80)
net = SpeechCNN(
hidden_size=hidden_size,
num_outputs=num_outputs,
dropout=dropout,
).to(device)
param_count = sum(p.numel() for p in net.parameters() if p.requires_grad)
print(f"Trainable parameters: {param_count:,}")
train_loader, test_loader = setup_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_fft=n_fft,
hop_length=hop_length,
target_sr=target_sr,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
weight_decay=weight_decay,
)
print("\n" + "=" * 80)
print("POST-TRAINING ANALYSIS")
print("=" * 80)
analyze_training_metrics(metrics)
final_eval = evaluate(net, test_loader, device, collect_representations=True)
visualize_confusion_matrix_plot(final_eval["confusion_matrix"])
if "representations" in final_eval:
visualize_tsne(final_eval["representations"], final_eval["targets"])
print("\nFinal Per-Class Accuracy:")
for cls, acc in sorted(final_eval["per_class_accuracy"].items(), key=lambda x: x[1], reverse=True):
print(f" {cls:>12s}: {acc:.4f}")
torch.save(
{
"model_state_dict": net.state_dict(),
"metrics": {
"train_losses": metrics["train_loss"],
"test_accuracies": metrics["test_accuracy"],
"final_accuracy": metrics["test_accuracy"][-1],
},
},
os.path.join(OUTPUT_DIR, "final_model.pth"),
)
print(f"\nDone! Accuracy: {metrics['test_accuracy'][-1]:.4f}")
print(f"Outputs: {OUTPUT_DIR}")
except KeyboardInterrupt:
print("Interrupted - saving...")
torch.save({"model_state_dict": net.state_dict()}, os.path.join(OUTPUT_DIR, "interrupted_model.pth"))
def optuna_objective(trial):
import optuna
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
n_fft = 1024
hop_length = 160
target_sr = 8000
num_outputs = 35
lr = trial.suggest_float("lr", 1e-4, 1e-2, log=True)
weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-2, log=True)
hidden_size = trial.suggest_categorical("hidden_size", [128, 256, 512])
dropout = trial.suggest_float("dropout", 0.1, 0.5)
batch_size = trial.suggest_categorical("batch_size", [256, 512, 1024])
num_epochs = 50
num_workers = min(8, os.cpu_count() - 2) if os.cpu_count() > 2 else 0
net = SpeechCNN(
hidden_size=hidden_size,
num_outputs=num_outputs,
dropout=dropout,
).to(device)
train_loader, test_loader = setup_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_fft=n_fft,
hop_length=hop_length,
target_sr=target_sr,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
weight_decay=weight_decay,
trial=trial,
)
except optuna.TrialPruned:
raise
finally:
del net, train_loader, test_loader
if device_str == "cuda":
torch.cuda.empty_cache()
return max(metrics["test_accuracy"])
def optuna_optimize(n_trials=50, study_name="speech_baseline_cnn", n_jobs=1):
import optuna
storage = f"sqlite:///{os.path.join(OUTPUT_DIR, study_name + '.db')}"
study = optuna.create_study(
study_name=study_name,
storage=storage,
direction="maximize",
load_if_exists=True,
pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=10),
)
print(f"Optuna: {n_trials} trials, {n_jobs} jobs")
print(f"DB: {storage}")
study.optimize(optuna_objective, n_trials=n_trials, n_jobs=n_jobs, gc_after_trial=True)
print(f"\nBest accuracy: {study.best_trial.value:.4f}")
for key, value in study.best_trial.params.items():
print(f" {key}: {value}")
import json
with open(os.path.join(OUTPUT_DIR, f"{study_name}_best_params.json"), "w") as f:
json.dump({"accuracy": study.best_trial.value, "optuna_epochs": 50, **study.best_trial.params}, f, indent=2)
return study
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Speech Commands Baseline CNN")
parser.add_argument("--optuna", action="store_true")
parser.add_argument("--n-trials", type=int, default=50)
parser.add_argument("--study-name", type=str, default="speech_baseline_cnn")
parser.add_argument("--n-jobs", type=int, default=1)
parser.add_argument("--json", type=str, default=None)
parser.add_argument("--gen-data", action="store_true", help="Pre-generate all cache files, then exit")
parser.add_argument("--cpu", action="store_true", help="Force CPU (disable CUDA)")
args = parser.parse_args()
if args.cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
if args.gen_data:
import multiprocessing
def _build_with_prefix(prefix, cls, kwargs):
"""Build cache with prefixed output lines."""
import sys
_orig_write = sys.stdout.write
_orig_flush = sys.stdout.flush
def _prefixed_write(s):
if s.strip():
_orig_write(f" [{prefix}] {s}")
else:
_orig_write(s)
_orig_flush()
sys.stdout.write = _prefixed_write
try:
cls(**kwargs)
finally:
sys.stdout.write = _orig_write
combos = []
for subset in ["training", "testing"]:
stag = "train" if subset == "training" else "test"
prefix = stag
combos.append((prefix, dict(subset=subset, n_fft=1024, hop_length=160, target_sr=8000)))
print(f"Pre-generating {len(combos)} cache files...")
# Ensure raw data is downloaded before spawning parallel workers
# (one archive for all subsets - subset only filters which wavs to list)
SPEECHCOMMANDS("../assets", download=True, subset="training")
for prefix, kwargs in combos:
_build_with_prefix(prefix, STFTSpeechDataset, kwargs)
print("Cache generation complete.")
sys.exit(0)
if args.optuna:
optuna_optimize(n_trials=args.n_trials, study_name=args.study_name, n_jobs=args.n_jobs)
else:
main(json_path=args.json)