input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# measure data loading time
data_time.update(time.time() - end)
if args.cuda:
inputs = inputs.cuda()
out, output_sizes = model(inputs, input_sizes)
out = out.transpose(0, 1) # TxNxH
loss = criterion(out, targets, output_sizes, target_sizes)
loss = loss / inputs.size(0) # average the loss by minibatch
inf = float("inf")
if args.distributed:
loss_value = reduce_tensor(loss, args.world_size)[0]
else:
loss_value = loss.item()
if loss_value == inf or loss_value == -inf:
print("WARNING: received an inf loss, setting loss value to 0")
loss_value = 0
avg_loss += loss_value
losses.update(loss_value, inputs.size(0))
# compute gradient
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
# SGD step
def train_main(args):
args.distributed = args.world_size > 1
main_proc = True
if args.distributed:
if args.gpu_rank:
torch.cuda.set_device(int(args.gpu_rank))
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size, rank=args.rank)
main_proc = args.rank == 0 # Only the first proc should save models
save_folder = args.save_folder
loss_results, cer_results, wer_results = torch.Tensor(args.epochs), torch.Tensor(args.epochs), torch.Tensor(
args.epochs)
best_wer = None
if args.visdom and main_proc:
from visdom import Visdom
viz = Visdom()
opts = dict(title=args.id, ylabel='', xlabel='Epoch', legend=['Loss', 'WER', 'CER'])
viz_window = None
epochs = torch.arange(1, args.epochs + 1)
if args.tensorboard and main_proc:
os.makedirs(args.log_dir, exist_ok=True)
from tensorboardX import SummaryWriter
tensorboard_writer = SummaryWriter(args.log_dir)
os.makedirs(save_folder, exist_ok=True)
avg_loss, start_epoch, start_iter = 0, 0, 0
if args.continue_from: # Starting from previous model
print("Loading checkpoint model %s" % args.continue_from)
package = torch.load(args.continue_from, map_location=lambda storage, loc: storage)
model = DeepSpeech.load_model_package(package)
labels = DeepSpeech.get_labels(model)
audio_conf = DeepSpeech.get_audio_conf(model)
parameters = model.parameters()
optimizer = torch.optim.SGD(parameters, lr=args.lr,
momentum=args.momentum, nesterov=True)
if not args.finetune: # Don't want to restart training
if args.cuda:
model.cuda()
optimizer.load_state_dict(package['optim_dict'])
start_epoch = int(package.get('epoch', 1)) - 1 # Index start at 0 for training
start_iter = package.get('iteration', None)
if start_iter is None:
start_epoch += 1 # We saved model after epoch finished, start at the next epoch.
start_iter = 0
else:
start_iter += 1
avg_loss = int(package.get('avg_loss', 0))
loss_results, cer_results, wer_results = package['loss_results'], package[
'cer_results'], package['wer_results']
if main_proc and args.visdom and \
package[
'loss_results'] is not None and start_epoch > 0: # Add previous scores to visdom graph
x_axis = epochs[0:start_epoch]
y_axis = torch.stack(
(loss_results[0:start_epoch], wer_results[0:start_epoch], cer_results[0:start_epoch]),
dim=1)
viz_window = viz.line(
X=x_axis,
Y=y_axis,
opts=opts,
)
if main_proc and args.tensorboard and \
package[
'loss_results'] is not None and start_epoch > 0: # Previous scores to tensorboard logs
for i in range(start_epoch):
values = {
'Avg Train Loss': loss_results[i],
'Avg WER': wer_results[i],
'Avg CER': cer_results[i]
}
tensorboard_writer.add_scalars(args.id, values, i + 1)
else:
with open(args.labels_path) as label_file:
labels = str(''.join(json.load(label_file)))
audio_conf = dict(sample_rate=args.sample_rate,
window_size=args.window_size,
window_stride=args.window_stride,
window=args.window,
noise_dir=args.noise_dir,
noise_prob=args.noise_prob,
noise_levels=(args.noise_min, args.noise_max))
rnn_type = args.rnn_type.lower()
assert rnn_type in supported_rnns, "rnn_type should be either lstm, rnn or gru"
model = DeepSpeech(rnn_hidden_size=args.hidden_size,
nb_layers=args.hidden_layers,
labels=labels,
rnn_type=supported_rnns[rnn_type],
audio_conf=audio_conf,
bidirectional=args.bidirectional)
parameters = model.parameters()
optimizer = torch.optim.SGD(parameters, lr=args.lr,
momentum=args.momentum, nesterov=True)
criterion = CTCLoss()
decoder = GreedyDecoder(labels)
train_dataset = SpectrogramDataset(audio_conf=audio_conf, manifest_filepath=args.train_manifest, labels=labels,
normalize=True, augment=args.augment)
test_dataset = SpectrogramDataset(audio_conf=audio_conf, manifest_filepath=args.val_manifest, labels=labels,
normalize=True, augment=False)
if not args.distributed:
train_sampler = BucketingSampler(train_dataset, batch_size=args.batch_size)
else:
train_sampler = DistributedBucketingSampler(train_dataset, batch_size=args.batch_size,
num_replicas=args.world_size, rank=args.rank)
train_loader = AudioDataLoader(train_dataset,
num_workers=args.num_workers, batch_sampler=train_sampler)
test_loader = AudioDataLoader(test_dataset, batch_size=args.batch_size,
num_workers=args.num_workers)
if (not args.no_shuffle and start_epoch != 0) or args.no_sorta_grad:
print("Shuffling batches for the following epochs")
train_sampler.shuffle(start_epoch)
if args.cuda:
model.cuda()
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=(int(args.gpu_rank),) if args.rank else None)
print(model)
print("Number of parameters: %d" % DeepSpeech.get_param_size(model))
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
for epoch in range(start_epoch, args.epochs):
model.train()
end = time.time()
start_epoch_time = time.time()
for i, (data) in enumerate(train_loader, start=start_iter):
if i == len(train_sampler):
break
inputs, targets, input_percentages, target_sizes = data
input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# measure data loading time
data_time.update(time.time() - end)
if args.cuda:
inputs = inputs.cuda()
out, output_sizes = model(inputs, input_sizes)
out = out.transpose(0, 1) # TxNxH
loss = criterion(out, targets, output_sizes, target_sizes)
loss = loss / inputs.size(0) # average the loss by minibatch
inf = float("inf")
if args.distributed:
loss_value = reduce_tensor(loss, args.world_size)[0]
else:
loss_value = loss.item()
if loss_value == inf or loss_value == -inf:
print("WARNING: received an inf loss, setting loss value to 0")
loss_value = 0
avg_loss += loss_value
losses.update(loss_value, inputs.size(0))
# compute gradient
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
# SGD step
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if not args.silent:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
(epoch + 1), (i + 1), len(train_sampler), batch_time=batch_time, data_time=data_time, loss=losses))
if args.checkpoint_per_batch > 0 and i > 0 and (i + 1) % args.checkpoint_per_batch == 0 and main_proc:
file_path = '%s/deepspeech_checkpoint_epoch_%d_iter_%d.pth' % (save_folder, epoch + 1, i + 1)
print("Saving checkpoint model to %s" % file_path)
torch.save(DeepSpeech.serialize(model, optimizer=optimizer, epoch=epoch, iteration=i,
loss_results=loss_results,
wer_results=wer_results, cer_results=cer_results, avg_loss=avg_loss),
file_path)
del loss
del out
avg_loss /= len(train_sampler)
epoch_time = time.time() - start_epoch_time
print('Training Summary Epoch: [{0}]\t'
'Time taken (s): {epoch_time:.0f}\t'
'Average Loss {loss:.3f}\t'.format(epoch + 1, epoch_time=epoch_time, loss=avg_loss))
start_iter = 0 # Reset start iteration for next epoch
total_cer, total_wer = 0, 0
model.eval()
with torch.no_grad():
for i, (data) in tqdm(enumerate(test_loader), total=len(test_loader)):
inputs, targets, input_percentages, target_sizes = data
input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# unflatten targets
split_targets = []
offset = 0
for size in target_sizes:
split_targets.append(targets[offset:offset + size])
offset += size
if args.cuda:
inputs = inputs.cuda()
out, output_sizes = model(inputs, input_sizes)
decoded_output, _ = decoder.decode(out.data, output_sizes)
target_strings = decoder.convert_to_strings(split_targets)
wer, cer = 0, 0
for x in range(len(target_strings)):
transcript, reference = decoded_output[x][0], target_strings[x][0]
wer += decoder.wer(transcript, reference) / float(len(reference.split()))
cer += decoder.cer(transcript, reference) / float(len(reference))
total_cer += cer
total_wer += wer
del out
wer = total_wer / len(test_loader.dataset)
cer = total_cer / len(test_loader.dataset)
wer *= 100
cer *= 100
loss_results[epoch] = avg_loss
wer_results[epoch] = wer
cer_results[epoch] = cer
print('Validation Summary Epoch: [{0}]\t'
'Average WER {wer:.3f}\t'
'Average CER {cer:.3f}\t'.format(epoch + 1, wer=wer, cer=cer))
if args.visdom and main_proc:
x_axis = epochs[0:epoch + 1]
y_axis = torch.stack(
(loss_results[0:epoch + 1], wer_results[0:epoch + 1], cer_results[0:epoch + 1]), dim=1)
if viz_window is None:
viz_window = viz.line(
X=x_axis,
Y=y_axis,
opts=opts,
)
else:
viz.line(
X=x_axis.unsqueeze(0).expand(y_axis.size(1), x_axis.size(0)).transpose(0, 1), # Visdom fix
Y=y_axis,
win=viz_window,
update='replace',
)
if args.tensorboard and main_proc:
values = {
'Avg Train Loss': avg_loss,
'Avg WER': wer,
'Avg CER': cer
}
tensorboard_writer.add_scalars(args.id, values, epoch + 1)
if args.log_params:
for tag, value in model.named_parameters():
tag = tag.replace('.', '/')
tensorboard_writer.add_histogram(tag, to_np(value), epoch + 1)
tensorboard_writer.add_histogram(tag + '/grad', to_np(value.grad), epoch + 1)
if args.checkpoint and main_proc:
file_path = '%s/deepspeech_%d.pth' % (save_folder, epoch + 1)
torch.save(DeepSpeech.serialize(model, optimizer=optimizer, epoch=epoch, loss_results=loss_results,
wer_results=wer_results, cer_results=cer_results),
file_path)
# anneal lr
optim_state = optimizer.state_dict()
optim_state['param_groups'][0]['lr'] = optim_state['param_groups'][0]['lr'] / args.learning_anneal
optimizer.load_state_dict(optim_state)
print('Learning rate annealed to: {lr:.6f}'.format(lr=optim_state['param_groups'][0]['lr']))
if (best_wer is None or best_wer > wer) and main_proc:
print("Found better validated model, saving to %s" % args.model_path)
torch.save(DeepSpeech.serialize(model, optimizer=optimizer, epoch=epoch, loss_results=loss_results,
wer_results=wer_results, cer_results=cer_results), args.model_path)
best_wer = wer
avg_loss = 0
if not args.no_shuffle:
print("Shuffling batches...")
train_sampler.shuffle(epoch)
def train_batch(self, epoch, batch_id, data):
inputs, targets, filenames, input_percentages, target_sizes = data
input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# measure data loading time
data_time.update(time.time() - self.end)
inputs = inputs.to(device)
input_sizes = input_sizes.to(device)
logits, probs, output_sizes = model(inputs, input_sizes)
assert logits.is_cuda
assert probs.is_cuda
assert output_sizes.is_cuda
split_targets = []
offset = 0
for size in target_sizes:
split_targets.append(targets[offset:offset + size])
offset += size
decoded_output, _ = decoder.decode(probs, output_sizes)
target_strings = decoder.convert_to_strings(split_targets)
for x in range(len(target_strings)):
transcript, reference = decoded_output[x][0], target_strings[x][0]
wer, cer, wer_ref, cer_ref = get_cer_wer(decoder, transcript,
reference)
train_dataset.update_curriculum(filenames[x], reference,
transcript, None, cer / cer_ref,
wer / wer_ref)
self.train_wer += wer
self.train_cer += cer
self.num_words += wer_ref
self.num_chars += cer_ref
logits = logits.transpose(0, 1) # TxNxH
if torch.isnan(logits).any(): # and args.nan == 'zero':
# work around bad data
print("WARNING: Working around NaNs in data")
logits[torch.isnan(logits)] = 0
loss = criterion(logits, targets, output_sizes.cpu(), target_sizes)
loss = loss / inputs.size(0) # average the loss by minibatch
loss = loss.to(device)
inf = float("inf")
if args.distributed:
loss_value = reduce_tensor(loss, args.world_size).item()
else:
loss_value = loss.item()
if loss_value == inf or loss_value == -inf:
print("WARNING: received an inf loss, setting loss value to 1000")
loss_value = 1000
loss_value = float(loss_value)
losses.update(loss_value, inputs.size(0))
# update_curriculum
# compute gradient
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
if torch.isnan(logits).any():
# work around bad data
print("WARNING: Skipping NaNs in backward step")
else:
# SGD step
optimizer.step()
if args.enorm:
enorm.step()
# measure elapsed time
batch_time.update(time.time() - self.end)
if not args.silent:
print('GPU-{0} Epoch {1} [{2}/{3}]\t'
'Time {batch_time.val:.2f} ({batch_time.avg:.2f})\t'
'Data {data_time.val:.2f} ({data_time.avg:.2f})\t'
'Loss {loss.val:.2f} ({loss.avg:.2f})\t'.format(
args.gpu_rank or VISIBLE_DEVICES[0],
epoch + 1,
batch_id + 1,
len(train_sampler),
batch_time=batch_time,
data_time=data_time,
loss=losses))
del inputs, targets, input_percentages, input_sizes
del logits, probs, output_sizes, target_sizes, loss
return loss_value
def check_model_quality(epoch, checkpoint, train_loss, train_cer, train_wer):
gc.collect()
torch.cuda.empty_cache()
val_cer_sum, val_wer_sum, val_loss_sum = 0, 0, 0
num_chars, num_words, num_losses = 0, 0, 0
model.eval()
with torch.no_grad():
for i, data in tq(enumerate(test_loader), total=len(test_loader)):
inputs, targets, filenames, input_percentages, target_sizes = data
input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# unflatten targets
split_targets = []
offset = 0
for size in target_sizes:
split_targets.append(targets[offset:offset + size])
offset += size
inputs = inputs.to(device)
logits, probs, output_sizes = model(inputs, input_sizes)
loss = criterion(logits.transpose(0, 1), targets,
output_sizes.cpu(), target_sizes)
loss = loss / inputs.size(0) # average the loss by minibatch
inf = float("inf")
if args.distributed:
loss_value = reduce_tensor(loss, args.world_size).item()
else:
loss_value = loss.item()
if loss_value == inf or loss_value == -inf:
print(
"WARNING: received an inf loss, setting loss value to 1000"
)
loss_value = 1000
loss_value = float(loss_value)
val_loss_sum = (val_loss_sum * 0.998 + loss_value * 0.002
) # discount earlier losses
val_loss_sum += loss_value
num_losses += 1
decoded_output, _ = decoder.decode(probs, output_sizes)
target_strings = decoder.convert_to_strings(split_targets)
for x in range(len(target_strings)):
transcript, reference = decoded_output[x][0], target_strings[
x][0]
wer, cer, wer_ref, cer_ref = get_cer_wer(
decoder, transcript, reference)
if x < 1:
print("CER: {:6.2f}% WER: {:6.2f}% Filename: {}".format(
cer / cer_ref * 100, wer / wer_ref * 100,
filenames[x]))
print('Reference:', reference, '\nTranscript:', transcript)
val_wer_sum += wer
val_cer_sum += cer
num_words += wer_ref
num_chars += cer_ref
del inputs, targets, input_percentages, target_sizes
del logits, probs, output_sizes, input_sizes
del split_targets, loss
if args.cuda:
torch.cuda.synchronize()
val_wer = 100 * val_wer_sum / num_words
val_cer = 100 * val_cer_sum / num_chars
print('Validation Summary Epoch: [{0}]\t'
'Average WER {wer:.3f}\t'
'Average CER {cer:.3f}\t'.format(epoch + 1,
wer=val_wer,
cer=val_cer))
val_loss = val_loss_sum / num_losses
plots.loss_results[epoch] = train_loss
plots.wer_results[epoch] = train_wer
plots.cer_results[epoch] = train_cer
plots.epochs[epoch] = epoch + 1
checkpoint_plots.loss_results[checkpoint] = val_loss
checkpoint_plots.wer_results[checkpoint] = val_wer
checkpoint_plots.cer_results[checkpoint] = val_cer
checkpoint_plots.epochs[checkpoint] = checkpoint + 1
plots.plot_progress(epoch, train_loss, train_cer, train_wer)
checkpoint_plots.plot_progress(checkpoint, val_loss, val_cer, val_wer)
if args.checkpoint_anneal != 1.0:
global lr_plots
lr_plots.loss_results[checkpoint] = val_loss
lr_plots.epochs[checkpoint] = get_lr()
zero_loss = lr_plots.loss_results == 0
lr_plots.loss_results[zero_loss] = val_loss
lr_plots.epochs[zero_loss] = get_lr()
lr_plots.plot_progress(checkpoint, val_loss, val_cer, val_wer)
return val_wer, val_cer
input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# measure data loading time
data_time.update(time.time() - end)
if args.cuda:
inputs = inputs.cuda()
out, output_sizes = model(inputs, input_sizes)
out = out.transpose(0, 1) # TxNxH
loss = criterion(out, targets, output_sizes, target_sizes)
loss = loss / inputs.size(0) # average the loss by minibatch
inf = float("inf")
if args.distributed:
loss_value = reduce_tensor(loss, args.world_size)[0]
else:
loss_value = loss.item()
if loss_value == inf or loss_value == -inf:
print("WARNING: received an inf loss, setting loss value to 0")
loss_value = 0
avg_loss += loss_value
losses.update(loss_value, inputs.size(0))
# compute gradient
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_norm)
# SGD step
def calculate_trainval_quality_metrics(checkpoint,
epoch,
loader,
plots_handle):
val_cer_sum, val_wer_sum, val_loss_sum = 0, 0, 0
num_chars, num_words, num_losses = 0, 0, 0
model.eval()
with torch.no_grad():
for i, data in enumerate(loader):#tq(enumerate(loader), total=len(loader)):
inputs, targets, filenames, input_percentages, target_sizes = data
input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
# unflatten targets
split_targets = []
offset = 0
for size in target_sizes:
split_targets.append(targets[offset:offset + size])
offset += size
inputs = inputs.to(device)
logits, probs, output_sizes = model(inputs, input_sizes)
loss = criterion(logits.transpose(0, 1), targets, output_sizes.cpu(), target_sizes)
loss = loss / inputs.size(0) # average the loss by minibatch
inf = float("inf")
if args.distributed:
loss_value = reduce_tensor(loss, args.world_size).item()
else:
loss_value = loss.item()
if loss_value == inf or loss_value == -inf:
print("WARNING: received an inf loss, setting loss value to 1000")
loss_value = 1000
loss_value = float(loss_value)
val_loss_sum = (val_loss_sum * 0.998 + loss_value * 0.002) # discount earlier losses
val_loss_sum += loss_value
num_losses += 1
decoded_output, _ = decoder.decode(probs, output_sizes)
target_strings = decoder.convert_to_strings(split_targets)
for x in range(len(target_strings)):
transcript, reference = decoded_output[x][0], target_strings[x][0]
wer, cer, wer_ref, cer_ref = get_cer_wer(decoder, transcript, reference)
if x < 1:
print("CER: {:6.2f}% WER: {:6.2f}% Filename: {}".format(cer/cer_ref*100, wer/wer_ref*100, filenames[x]))
print('Reference:', reference, '\nTranscript:', transcript)
times_used = trainval_dataset.curriculum[filenames[x]]['times_used']+1
trainval_dataset.update_curriculum(filenames[x],
reference, transcript,
None,
cer / cer_ref, wer / wer_ref,
times_used=times_used)
val_wer_sum += wer
val_cer_sum += cer
num_words += wer_ref
num_chars += cer_ref
del inputs, targets, input_percentages, target_sizes
del logits, probs, output_sizes, input_sizes
del split_targets, loss
if args.cuda:
torch.cuda.synchronize()
val_wer = 100 * val_wer_sum / num_words
val_cer = 100 * val_cer_sum / num_chars
print('TrainVal Summary Epoch: [{0}]\t'
'Average WER {wer:.3f}\t'
'Average CER {cer:.3f}\t'.format(epoch + 1, wer=val_wer, cer=val_cer))
val_loss = val_loss_sum / num_losses
plots_handle.loss_results[checkpoint] = val_loss
plots_handle.wer_results[checkpoint] = val_wer
plots_handle.cer_results[checkpoint] = val_cer
plots_handle.epochs[checkpoint] = checkpoint + 1
plots_handle.plot_progress(checkpoint, val_loss, val_cer, val_wer)
