I’m trying to build a simple ASR model using the LibriSpeech dataset built into torchvision, but I seem to have a problem with the CTCLoss function.
When I try to use nn.CTCLoss I get None gradient for my loss.
I’m using pytorch 1.4.0 and torchaudio 0.4.0.
I get the following output from running ./train.py:
Using device: cuda
0%| | 0/676 [00:00<?, ?it/s]Training epoch: 1
Loss: 3228.21142578125: 0%| | 0/676 [00:01<?, ?it/s]
None
Loss: 7.94204209823249e+22: 0%|2 | 1/676 [00:02<18:37, 1.65s/it]
None
Loss: nan: 0%|4 | 2/676 [00:03<14:56, 1.33s/it]
By using the code:
train.py
#!/usr/bin/env python3
import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_sequence
from torch.nn import CTCLoss
from torch.utils.data import DataLoader
from torch.optim import Adam, SGD
from torchaudio.datasets import LIBRISPEECH
from models.Basic import Basic
from tqdm import tqdm
import multiprocessing
dictionary = "ABCDEFGHIJKLMNOPQRSTUVWXYZ' "
def text_to_tensor(text, dictionary = dictionary):
"""
This function will convert a string of text
to a tensor of character indicies in the given dictionary.
The indicies will start from 1, as 0 means the blank
character.
"""
return torch.tensor([
dictionary.index(c) + 1 if c in dictionary else 0
for c in list(text.upper())
])
def pad_collate(datapoints):
waveforms, sample_rates, utterances, speaker_ids, chapter_ids, utterance_ids = zip(*datapoints)
batch_size = len(datapoints)
waveform_lengths = torch.tensor([waveform.shape[1] for waveform in waveforms])
waveforms = pad_sequence([wave.T for wave in waveforms], batch_first=True)
utterance_lengths = torch.tensor([len(utterance) for utterance in utterances])
# We convert our label text to tensor of dictionary indicies
# and reshape the data to (N, S) where N is batch size
# and S is max target length. Is needed for CTCLoss:
# https://pytorch.org/docs/stable/nn.html#torch.nn.CTCLoss
utterances = torch.cat([text_to_tensor(utterance) for utterance in utterances])
return batch_size, waveforms, waveform_lengths, utterances, utterance_lengths
def train(num_epochs=10, batch_size=4, num_workers=multiprocessing.cpu_count()):
dataset = LIBRISPEECH("../data", "dev-clean", download=True)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=pad_collate, num_workers=num_workers, pin_memory=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Basic(n_classes = len(dictionary) + 1).to(device)
optimizer = SGD(model.parameters(), lr=0.0001)
loss_fn = CTCLoss()
print(f"Using device: {device}")
tqdm_dataloader = tqdm(dataloader)
for epoch in range(num_epochs):
print(f"Training epoch: {epoch+1}")
for batch_size, X, X_lengths, y, y_lengths in tqdm_dataloader:
# First we zero our gradients, to make everything work nicely.
optimizer.zero_grad()
X = X.permute(0, 2, 1).to(device)
X_lengths = X_lengths.to(device)
y = y.to(device)
# We predict the outputs using our model
# and reshape the data to size (T, N, C) where
# T is target length, N is batch size and C is number of classes.
# In our case that is the length of the dictionary + 1
# as we also need one more class for the blank character.
pred_y = model(X)
pred_y = pred_y.permute(2, 0, 1)
pred_y_lengths = model.forward_shape(X_lengths).to(device)
loss = loss_fn(pred_y, y, pred_y_lengths, y_lengths)
tqdm_dataloader.set_description(f"Loss: {loss.item()}")
loss.backward()
print(loss.grad)
optimizer.step()
if __name__ == "__main__":
train()
With model
models/Basic.py
import torch.nn as nn
import torch.nn.functional as F
class Basic(nn.Module):
def __init__(self, n_classes):
super(Basic, self).__init__()
self.c1 = nn.Conv1d(1, 128, 32)
self.c2 = nn.Conv1d(128, n_classes, 64)
def forward(self, X):
a = self.c1(X)
b = F.relu(a)
c = self.c2(b)
return F.log_softmax(c, dim=1)
def forward_shape(self, lengths):
return lengths - 32 - 64 + 1 + 1