I used the following github repo: Speech Recognition. But since it didn't have code to train and save the model, I looked online and added code to speech_recognition to train and save the model and create infer.py to make predictions. But now even using a .flac file from the training dataset, it gives garbled output like: Transcript: ' f ' i ' ' ' i ' a ' ' a ' s ' o ' f ' o ' f ' i ' m ' i ' a ' i ' f ' o ' r ' ' i ' ' i ' This is the code: data_preprocessing.py import torch import torch.nn as nn import torch.utils.data import torchaudio class TextTransform: '''Maps characters to integers and vice versa''' def __init__(self): char_map_str = ''' ' 0 <SPACE> 1 a 2 b 3 c 4 d 5 e 6 f 7 g 8 h 9 i 10 j 11 k 12 l 13 m 14 n 15 o 16 p 17 q 18 r 19 s 20 t 21 u 22 v 23 w 24 x 25 y 26 z 27 ''' self.char_map = {} self.index_map = {} for line in char_map_str.strip().split('\n'): ch, index = line.split() self.char_map[ch] = int(index) self.index_map[int(index)] = ch self.index_map[1] = ' ' def text_to_int(self, text): ''' Use a character map and convert text to an integer sequence ''' int_sequence = [] for c in text: if c == ' ': ch = self.char_map['<SPACE>'] else: ch = self.char_map[c] int_sequence.append(ch) return int_sequence def int_to_text(self, labels): ''' Use a character map and convert integer labels to an text sequence ''' string = [] for i in labels: string.append(self.index_map) return ''.join(string).replace('', ' ') # TODO: Questions: Log Mel Spectrogram vs Mel Spectrogram class LogMelSpectogram(nn.Module): # TODO: Understand the parameters def __init__(self, sample_rate=8000, n_mels=81, win_length=160, hop_length=80): super(LogMelSpectogram, self).__init__() self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate, n_mels=n_mels, win_length=win_length, hop_length=hop_length) def forward(self, waveform): mel_spectrogram = self.mel_spectrogram(waveform) # Add 1e-6 to avoid taking log of zero log_mel_spectrogram = torch.log(mel_spectrogram + 1e-9) return log_mel_spectrogram class PreprocessData(torch.utils.data.Dataset): def __init__(self, dataset, validation_set, sample_rate=8000, n_mels=81, win_length=160, hop_length=80): super(PreprocessData).__init__() self.dataset = dataset self.text_transform = TextTransform() if validation_set: self.preprocess_audio = LogMelSpectogram(sample_rate=sample_rate, n_mels=n_mels, win_length=win_length, hop_length=hop_length) else: # TODO: Why no frequency and time masking for validation set? self.preprocess_audio = nn.Sequential( LogMelSpectogram(sample_rate=sample_rate, n_mels=n_mels, win_length=win_length, hop_length=hop_length), torchaudio.transforms.FrequencyMasking(freq_mask_param=15), torchaudio.transforms.TimeMasking(time_mask_param=35) ) def __getitem__(self, index): # waveform, sample_rate, label, speaker_id, chapter_id, utterance_id waveform, _, label, _, _, _ = self.dataset[index] # Convert waveform to log mel spectrogram log_mel_spectrogram = self.preprocess_audio(waveform) # Get the length of the log mel spectrogram log_mel_spectrogram_len = log_mel_spectrogram.shape[0] // 2 # TODO: Why divide by 2? # Convert label text to integer sequence label_in_int = torch.tensor(self.text_transform.text_to_int(label.lower())) # Get the length of the label label_len = torch.tensor(len(label_in_int)) return log_mel_spectrogram, label_in_int, log_mel_spectrogram_len, label_len def __len__(self): return len(self.dataset) gather.py: import os import torchaudio def gather_data(): if not os.path.isdir('data'): os.makedirs('data') train_dataset = torchaudio.datasets.LIBRISPEECH('data/', url='train-clean-100', download=True) test_dataset = torchaudio.datasets.LIBRISPEECH('data/', url='test-clean', download=True) return train_dataset, test_dataset model.py: import torch import torch.nn as nn import torch.nn.functional as F import math class CNNLayerNorm(nn.Module): """Layer normalization built for cnns input""" def __init__(self, n_feats): super(CNNLayerNorm, self).__init__() self.layer_norm = nn.LayerNorm(n_feats) def forward(self, x): # x (batch, channel, feature, time) x = x.transpose(2, 3).contiguous() # (batch, channel, time, feature) x = self.layer_norm(x) return x.transpose(2, 3).contiguous() # (batch, channel, feature, time) class ResidualCNN(nn.Module): """Residual CNN inspired by https://arxiv.org/pdf/1603.05027.pdf except with layer norm instead of batch norm """ def __init__(self, in_channels, out_channels, kernel, stride, dropout, n_feats): super(ResidualCNN, self).__init__() self.cnn1 = nn.Conv2d(in_channels, out_channels, kernel, stride, padding=kernel//2) self.cnn2 = nn.Conv2d(out_channels, out_channels, kernel, stride, padding=kernel//2) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.layer_norm1 = CNNLayerNorm(n_feats) self.layer_norm2 = CNNLayerNorm(n_feats) def forward(self, x): residual = x # (batch, channel, feature, time) x = self.layer_norm1(x) x = F.gelu(x) x = self.dropout1(x) x = self.cnn1(x) x = self.layer_norm2(x) x = F.gelu(x) x = self.dropout2(x) x = self.cnn2(x) # residual connection - Skip connection x += residual return x # (batch, channel, feature, time) class BidirectionalGRU(nn.Module): def __init__(self, rnn_dim, hidden_size, dropout, batch_first): super(BidirectionalGRU, self).__init__() self.BiGRU = nn.GRU( input_size=rnn_dim, hidden_size=hidden_size, num_layers=1, batch_first=batch_first, bidirectional=True) self.layer_norm = nn.LayerNorm(rnn_dim) self.dropout = nn.Dropout(dropout) def forward(self, x): x = self.layer_norm(x) x = F.gelu(x) x, _ = self.BiGRU(x) x = self.dropout(x) return x class SpeechRecognitionModel(nn.Module): """Speech Recognition Model Inspired by DeepSpeech 2""" def __init__(self, n_cnn_layers, n_rnn_layers, rnn_dim, n_class, n_feats, stride=2, dropout=0.1): super(SpeechRecognitionModel, self).__init__() # n_feats = n_feats//2 n_feats = math.ceil(n_feats / stride) self.cnn = nn.Conv2d(1, 32, 3, stride=stride, padding=1) # cnn for extracting heirachal features # n residual cnn layers with filter size of 32 self.rescnn_layers = nn.Sequential(*[ ResidualCNN(32, 32, kernel=3, stride=1, dropout=dropout, n_feats=n_feats) for _ in range(n_cnn_layers) ]) self.fully_connected = nn.Linear(n_feats*32, rnn_dim) self.birnn_layers = nn.Sequential(*[ BidirectionalGRU(rnn_dim=rnn_dim if i==0 else rnn_dim*2, hidden_size=rnn_dim, dropout=dropout, batch_first=i==0) for i in range(n_rnn_layers) ]) self.classifier = nn.Sequential( nn.Linear(rnn_dim*2, rnn_dim), # birnn returns rnn_dim*2 nn.GELU(), nn.Dropout(dropout), nn.Linear(rnn_dim, n_class) ) def forward(self, x): x = self.cnn(x) x = self.rescnn_layers(x) sizes = x.size() x = x.view(sizes[0], sizes[1] * sizes[2], sizes[3]) # (batch, feature, time) x = x.transpose(1, 2) # (batch, time, feature) x = self.fully_connected(x) x = self.birnn_layers(x) x = self.classifier(x) return x speech_recognition.py import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from gather import gather_data from data_preprocessing import PreprocessData, TextTransform from model import SpeechRecognitionModel def collate_fn(batch): """ Pads variable-length audio features and labels to max lengths in batch. Returns: features_tensor: (batch, 1, n_mels, max_time) labels_tensor: (batch, max_label_len) feature_lengths: (batch,) label_lengths: (batch,) """ feats, labels, _, _ = zip(*batch) feat_lens = [f.shape[-1] for f in feats] label_lens = [l.shape[0] for l in labels] max_feat_len = max(feat_lens) padded_feats = [F.pad(f, (0, max_feat_len - f.shape[-1])) for f in feats] feats_tensor = torch.stack(padded_feats) max_label_len = max(label_lens) padded_labels = [F.pad(l, (0, max_label_len - l.shape[0]), value=0) for l in labels] labels_tensor = torch.stack(padded_labels) return feats_tensor, labels_tensor, torch.tensor(feat_lens), torch.tensor(label_lens) def train_one_epoch(model, device, dataloader, criterion, optimizer): model.train() total_loss = 0.0 for features, labels, _, label_lens in dataloader: features, labels = features.to(device), labels.to(device) batch_size = features.size(0) # Forward pass outputs = model(features) # (batch, time, classes) time_steps = outputs.size(1) log_probs = outputs.log_softmax(2).permute(1, 0, 2) # (time, batch, classes) # Create input_lengths = full length for each (after downsampling) input_lengths = torch.full((batch_size,), time_steps, dtype=torch.long).to(device) label_lens = label_lens.to(device) # Compute loss loss = criterion(log_probs, labels, input_lengths, label_lens) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() avg_loss = total_loss / len(dataloader) print(f"Train Loss: {avg_loss:.4f}") return avg_loss def main(batch_size=8, epochs=20, lr=1e-4, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')): # Prepare data train_ds, test_ds = gather_data() train_ds = PreprocessData(train_ds, validation_set=False) train_loader = DataLoader( train_ds, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, collate_fn=collate_fn ) # Determine n_mels from one sample sample_feat, _, _, _ = train_ds[0] n_mels = sample_feat.shape[1] text_transform = TextTransform() n_classes = len(text_transform.char_map) + 1 # +1 for CTC blank # Build model model = SpeechRecognitionModel( n_cnn_layers=3, n_rnn_layers=5, rnn_dim=512, n_class=n_classes, n_feats=n_mels ).to(device) # Loss & optimizer criterion = nn.CTCLoss(blank=0, zero_infinity=True) optimizer = torch.optim.Adam(model.parameters(), lr=lr) # Training loop best_loss = float('inf') for epoch in range(1, epochs + 1): print(f"=== Epoch {epoch}/{epochs} ===") try: loss = train_one_epoch(model, device, train_loader, criterion, optimizer) except RuntimeError as e: if 'out of memory' in str(e): print("⚠️ CUDA OOM: emptying cache and exiting training loop") torch.cuda.empty_cache() break else: raise # Save best model if loss < best_loss: best_loss = loss torch.save(model.state_dict(), 'speech_model.pth') print("Model saved: speech_model.pth") if device.type == 'cuda': torch.cuda.empty_cache() print("Training complete!") if __name__ == '__main__': main() and then I checked online and added this script: infer.py: import torch import torchaudio from model import SpeechRecognitionModel from data_preprocessing import TextTransform import torchaudio.transforms as T def load_model(model_path, device, input_dim, output_dim): text_transform = TextTransform() n_classes = len(text_transform.char_map) + 1 # +1 for CTC blank model = SpeechRecognitionModel( n_cnn_layers=3, n_rnn_layers=5, rnn_dim=512, n_class=n_classes, n_feats=input_dim ).to(device) checkpoint = torch.load(model_path, map_location=device) checkpoint_state = checkpoint if 'state_dict' in checkpoint: checkpoint_state = checkpoint['state_dict'] model_state = model.state_dict() for key in ['classifier.3.weight', 'classifier.3.bias']: if key in checkpoint_state and key in model_state: if checkpoint_state[key].shape != model_state[key].shape: print(f"Skipping incompatible layer: {key}") del checkpoint_state[key] model.load_state_dict(checkpoint_state, strict=False) model.eval() return model def greedy_decode(output, blank_index): arg_maxes = torch.argmax(output, dim=2) decodes = [] for i, args in enumerate(arg_maxes): decode = [] previous = -1 for j in args: if j != previous: if j != blank_index: decode.append(j.item()) previous = j decodes.append(decode) return decodes def transcribe(audio_path, model_path): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') tt = TextTransform() n_feats = 81 # Must match model's training config waveform, sample_rate = torchaudio.load(audio_path) waveform = waveform.to(device) resampler = T.Resample(orig_freq=sample_rate, new_freq=8000).to(device) waveform = resampler(waveform) mel_spec = T.MelSpectrogram( sample_rate=8000, n_mels=n_feats, win_length=160, hop_length=80 ).to(device) spec = mel_spec(waveform) # if spec.ndim == 3: # spec = spec.unsqueeze(0) # (1, channel, n_mels, time) spec = spec.unsqueeze(0) # CNN expects 4D: (batch, channel, features, time) # Model's layer_norm1 is using x.transpose(2, 3), so shape must be (B, C, F, T) model_input = spec # shape: (1, 1, n_feats, time) model = load_model(model_path, device, n_feats, len(tt.char_map)) with torch.no_grad(): output = model(model_input) # (batch, time, classes) print('output shape : ' ,output.shape) pred_indices = greedy_decode(output, blank_index=len(tt.char_map)) pred_str = tt.int_to_text(pred_indices[0]) return pred_str if __name__ == "__main__": import sys audio_file = sys.argv[1] if len(sys.argv) > 1 else "sample.flac" weights_file = sys.argv[2] if len(sys.argv) > 2 else "speech_model.pth" print("Transcript:", transcribe(audio_file, weights_file)) after training was done, it had a loss of 0.28. output of infer.py: output shape : torch.Size([1, 589, 29]) Transcript: ' f ' i ' ' ' i ' a ' ' a ' s ' o ' f ' o ' f ' i ' m ' i ' a ' i ' f ' o ' r ' ' i ' Continue reading...