To extract speaker embeddings, provide a 1-channel 16 kHz mono audio signal to the model. The model expects an input tensor of shape [N, T], where N is the batch size and T is the number of samples.
Inference Workflow
- Audio Requirements: 16,000 Hz sampling rate, mono channel.
- Device Setup: Move the model to the target device (
cuda or cpu). - Precision: Use
torch.float16 (AMP) for GPU inference to improve performance, or torch.float32 for CPU. - Mode: Set the model to
.eval() mode and use torch.no_grad(). - Output: The model returns an embedding of shape
[N, 192].
import torch
import torchaudio
# 1. Load audio samples (must be 16kHz mono)
samples, fs = torchaudio.load("assets/audio.wav") # shape [1, T]
assert fs == 16000, f"Audio sampling rate {fs} != 16000"
assert samples.shape[0] == 1, f"Expected mono audio, but got {samples.shape[0]} channels"
# 2. Load model (e.g., mixed dataset version)
model = torch.hub.load(
"IDRnD/ReDimNet",
"ReDimNet",
model_name="M",
train_type="ft_mix",
dataset="vb2+vox2+cnc"
)
# 3. Select device and setup inference precision
device_type = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_type)
precision = torch.float16 if device_type == "cuda" else torch.float32
# 4. Setup model evaluation mode
model = model.to(device)
model.eval()
# 5. Run inference
with torch.no_grad(), torch.autocast(device_type=device_type, dtype=precision):
# Input shape [N, T]
embedding = model(samples.to(device))
print(embedding.shape, embedding.dtype) # Output shape: [N, 192]
