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)