File size: 3,625 Bytes
916419b | 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 | from typing import Tuple
import pytest
import torch
import mel_spectrogram
def torch_mel_spectrogram(
samples: torch.Tensor,
filters: torch.Tensor,
fft_size: int,
fft_step: int,
n_frames: int,
) -> torch.Tensor:
"""Native PyTorch implementation of mel spectrogram generation."""
# Calculate frames based on CUDA implementation's frame count
frames = []
# Frame the audio to match CUDA implementation
for i in range(n_frames):
start = i * fft_step
end = min(start + fft_size, samples.size(0))
if end - start < fft_size:
# Zero-pad if needed
frame = torch.zeros(fft_size, device=samples.device)
frame[: end - start] = samples[start:end]
frames.append(frame)
else:
frames.append(samples[start:end])
# Stack frames
frames = torch.stack(frames) # [n_frames, fft_size]
# Apply Hanning window
window = torch.hann_window(fft_size, device=samples.device)
windowed = frames * window
# Compute FFT
fft_complex = torch.fft.rfft(windowed, dim=1) # [n_frames, n_fft]
fft_magnitudes = torch.abs(fft_complex) # [n_frames, n_fft]
# Apply mel filterbank
mel_spec = torch.matmul(fft_magnitudes, filters.T) # [n_frames, n_mel]
# Log-scale and normalization
eps = 1e-10
mel_spec = torch.log10(torch.clamp(mel_spec, min=eps))
# Normalize like the CUDA implementation
max_val = mel_spec.max()
min_val = max_val - 8.0
mel_spec = torch.clamp(mel_spec, min=min_val) / 4.0 + 1.0
# Transpose to match CUDA implementation output shape [n_mel, n_frames]
return mel_spec.T
@pytest.mark.parametrize(
"n_samples, n_mel, fft_size, fft_step, seed",
[
(16000, 80, 1024, 512, 42),
(32000, 40, 1024, 256, 123),
(8000, 60, 512, 256, 987),
],
)
def test_mel_spectrogram(
n_samples: int, n_mel: int, fft_size: int, fft_step: int, seed: int
) -> None:
"""Test the CUDA mel_spectrogram function against a native PyTorch implementation."""
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
# Set seed for reproducibility
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
device = "cuda"
# Calculate the number of frames and FFT bins
n_frames = n_samples // fft_step
n_fft = 1 + fft_size // 2 # Number of unique FFT bins
# Create tensors with the correct shapes
samples = torch.randn(n_samples, dtype=torch.float32, device=device)
filters = torch.abs(torch.randn(n_mel, n_fft, dtype=torch.float32, device=device))
filters = filters / filters.sum(dim=1, keepdim=True) # Normalize filters
# Run the CUDA implementation
cuda_output = torch.zeros(n_mel, n_frames, dtype=torch.float32, device=device)
mel_spectrogram.mel_spectrogram(
cuda_output,
samples,
filters,
fft_size,
fft_step,
)
# Get actual frame count from CUDA output
actual_frames = cuda_output.shape[1]
# Run the PyTorch implementation with matched frame count
torch_output = torch_mel_spectrogram(
samples, filters, fft_size, fft_step, actual_frames
)
# Check shapes match
assert cuda_output.shape == torch_output.shape
# Compare results (allowing for numerical differences)
mae = torch.abs(cuda_output - torch_output).mean().item()
print(f"Mean Absolute Error: {mae}")
# Simple validation
assert not torch.isnan(cuda_output).any(), "CUDA output contains NaN values"
assert not torch.allclose(cuda_output, torch.zeros_like(cuda_output))
|