C++ Inference Code by Using ONNX model

[shubhammore1310]

Hi Team, I am trying to get the timestamps with BFA in C++. I have converted the checkpoint model to ONNX model and taking output with ONNX runtime (ort). I am currently facing issue to convert the ViterbiDecoder class and its method in C++. I am attaching the Python file i am referring for inference, C++ code made so far and the exported ONNX model.

bfaonnx.py

main.cpp

[shubhammore1310]

Following Python code I used to export the checkpoint to ONNX model

`import torch
import torch.nn as nn
from bfaSingle import ContextFreePhonemeRecognizer # Import your classes

class CUPEPredictONNX(nn.Module):
def init(self, cupe_ckpt_path):
super().init()
self.model = ContextFreePhonemeRecognizer()

    # Load config and build model exactly as in CUPEEmbeddingsExtractor
    checkpoint = torch.load(cupe_ckpt_path, map_location='cpu', weights_only=True)
    if 'model_config' not in checkpoint:
        raise ValueError("Model config not found in checkpoint")
    
    self.model.load_config_state_dict(checkpoint['model_config'])
    self.model.make_model()
    
    # Load weights exactly as in CUPEEmbeddingsExtractor
    state_dict = checkpoint['state_dict']
    state_dict = {k.replace('model.', ''): v for k, v in state_dict.items()}
    # Filter valid keys (same as original)
    valid_prefixes = [
        'feature_extractor.', 'freq_attention.', 'temporal_stream.', 
        'spectral_stream.', 'fusion.', 'window_processor.', 
        'classifier.', 'group_classifier.'
    ]
    state_dict = {k: v for k, v in state_dict.items() 
                 if any(prefix in k for prefix in valid_prefixes)
                 and not any(skip in k for skip in ['quantizer.', 'prediction_head.', 'final_proj.'])}
    
    self.model.load_state_dict(state_dict, strict=False)
    self.model.eval()

def forward(self, audio_batch):
    """Exact replica of CUPEEmbeddingsExtractor.predict() without training noise/groups_only"""
    # Input: [batch_size, wav_length]
    x = audio_batch.unsqueeze(1)  # [B, 1, T]
    
    # Feature extraction
    features = self.model.feature_extractor(x)
    att = self.model.freq_attention(features)
    features = features * att
    
    # Dual streams
    temporal = self.model.temporal_stream(features)
    spectral = self.model.spectral_stream(features)
    
    # Fusion
    fused = torch.cat([temporal, spectral], dim=1)
    fused = self.model.fusion(fused)
    fused = fused.transpose(1, 2)  # [B, frames, dim]
    
    # Transformer + classifiers
    embeddings = self.model.window_processor(fused)
    logits_class = self.model.classifier(embeddings)
    logits_group = self.model.group_classifier(embeddings)
    
    return logits_class, logits_group, embeddings

Export

def export_to_onnx(cupe_ckpt_path, output_path="cupe_predict_english.onnx", dummy_wav_len=1891):
model = CUPEPredictONNX(cupe_ckpt_path)
device = next(model.parameters()).device
model.eval()

# Dummy input matching your example
dummy_input = torch.randn(2, dummy_wav_len, dtype=torch.float32).to(device)

torch.onnx.export(
    model,
    dummy_input,
    output_path,
    opset_version=18,
    do_constant_folding=True,
    input_names=['audio'],
    output_names=['logits_class', 'logits_group', 'embeddings'],
    dynamic_axes={
        'audio': {0: 'batch_size', 1: 'wav_length'},  # Variable audio length
        'logits_class': {0: 'batch_size', 1: 'frames'},
        'logits_group': {0: 'batch_size', 1: 'frames'},
        'embeddings': {0: 'batch_size', 1: 'frames'}
    },
    verbose=False
)
print(f"Exported {output_path}")
return output_path

Usage

ckpt_path = "en_libri1000_ua01c_e4_val_GER3D0.2186.ckpt"
onnx_path = export_to_onnx(ckpt_path)
`