def train_one_epoch(dataloader: torch.utils.data.DataLoader, valid_dataloader: torch.utils.data.DataLoader, model: AcousticModel, device: torch.device, graph_compiler: CtcTrainingGraphCompiler, optimizer: torch.optim.Optimizer, accum_grad: int, att_rate: float, current_epoch: int, tb_writer: SummaryWriter, num_epochs: int, global_batch_idx_train: int): """One epoch training and validation. Args: dataloader: Training dataloader valid_dataloader: Validation dataloader model: Acoustic model to be trained device: Training device, torch.device("cpu") or torch.device("cuda", device_id) graph_compiler: MMI training graph compiler optimizer: Training optimizer accum_grad: Number of gradient accumulation att_rate: Attention loss rate, final loss is att_rate * att_loss + (1-att_rate) * other_loss current_epoch: current training epoch, for logging only tb_writer: tensorboard SummaryWriter num_epochs: total number of training epochs, for logging only global_batch_idx_train: global training batch index before this epoch, for logging only Returns: A tuple of 3 scalar: (total_objf / total_frames, valid_average_objf, global_batch_idx_train) - `total_objf / total_frames` is the average training loss - `valid_average_objf` is the average validation loss - `global_batch_idx_train` is the global training batch index after this epoch """ total_objf, total_frames, total_all_frames = 0., 0., 0. valid_average_objf = float('inf') time_waiting_for_batch = 0 forward_count = 0 prev_timestamp = datetime.now() model.train() for batch_idx, batch in enumerate(dataloader): forward_count += 1 if forward_count == accum_grad: is_update = True forward_count = 0 else: is_update = False global_batch_idx_train += 1 timestamp = datetime.now() time_waiting_for_batch += (timestamp - prev_timestamp).total_seconds() curr_batch_objf, curr_batch_frames, curr_batch_all_frames = get_objf( batch=batch, model=model, device=device, graph_compiler=graph_compiler, is_training=True, is_update=is_update, accum_grad=accum_grad, att_rate=att_rate, tb_writer=tb_writer, global_batch_idx_train=global_batch_idx_train, optimizer=optimizer) total_objf += curr_batch_objf total_frames += curr_batch_frames total_all_frames += curr_batch_all_frames if batch_idx % 10 == 0: logging.info( 'batch {}, epoch {}/{} ' 'global average objf: {:.6f} over {} ' 'frames ({:.1f}% kept), current batch average objf: {:.6f} over {} frames ({:.1f}% kept) ' 'avg time waiting for batch {:.3f}s'.format( batch_idx, current_epoch, num_epochs, total_objf / total_frames, total_frames, 100.0 * total_frames / total_all_frames, curr_batch_objf / (curr_batch_frames + 0.001), curr_batch_frames, 100.0 * curr_batch_frames / curr_batch_all_frames, time_waiting_for_batch / max(1, batch_idx))) if tb_writer is not None: tb_writer.add_scalar('train/global_average_objf', total_objf / total_frames, global_batch_idx_train) tb_writer.add_scalar( 'train/current_batch_average_objf', curr_batch_objf / (curr_batch_frames + 0.001), global_batch_idx_train) # if batch_idx >= 10: # print("Exiting early to get profile info") # sys.exit(0) if batch_idx > 0 and batch_idx % 200 == 0: total_valid_objf, total_valid_frames, total_valid_all_frames = get_validation_objf( dataloader=valid_dataloader, model=model, device=device, graph_compiler=graph_compiler) valid_average_objf = total_valid_objf / total_valid_frames model.train() logging.info( 'Validation average objf: {:.6f} over {} frames ({:.1f}% kept)' .format(valid_average_objf, total_valid_frames, 100.0 * total_valid_frames / total_valid_all_frames)) if tb_writer is not None: tb_writer.add_scalar('train/global_valid_average_objf', valid_average_objf, global_batch_idx_train) model.write_tensorboard_diagnostics( tb_writer, global_step=global_batch_idx_train) prev_timestamp = datetime.now() return total_objf / total_frames, valid_average_objf, global_batch_idx_train
def train_one_epoch(dataloader: torch.utils.data.DataLoader, valid_dataloader: torch.utils.data.DataLoader, model: AcousticModel, P: k2.Fsa, device: torch.device, graph_compiler: MmiTrainingGraphCompiler, optimizer: torch.optim.Optimizer, current_epoch: int, tb_writer: SummaryWriter, num_epochs: int, global_batch_idx_train: int): total_objf, total_frames, total_all_frames = 0., 0., 0. valid_average_objf = float('inf') time_waiting_for_batch = 0 prev_timestamp = datetime.now() model.train() for batch_idx, batch in enumerate(dataloader): global_batch_idx_train += 1 timestamp = datetime.now() time_waiting_for_batch += (timestamp - prev_timestamp).total_seconds() P.set_scores_stochastic_(model.P_scores) assert P.is_cpu assert P.requires_grad is True curr_batch_objf, curr_batch_frames, curr_batch_all_frames = get_objf( batch=batch, model=model, P=P, device=device, graph_compiler=graph_compiler, is_training=True, tb_writer=tb_writer, global_batch_idx_train=global_batch_idx_train, optimizer=optimizer) total_objf += curr_batch_objf total_frames += curr_batch_frames total_all_frames += curr_batch_all_frames if batch_idx % 10 == 0: logging.info( 'batch {}, epoch {}/{} ' 'global average objf: {:.6f} over {} ' 'frames ({:.1f}% kept), current batch average objf: {:.6f} over {} frames ({:.1f}% kept) ' 'avg time waiting for batch {:.3f}s'.format( batch_idx, current_epoch, num_epochs, total_objf / total_frames, total_frames, 100.0 * total_frames / total_all_frames, curr_batch_objf / (curr_batch_frames + 0.001), curr_batch_frames, 100.0 * curr_batch_frames / curr_batch_all_frames, time_waiting_for_batch / max(1, batch_idx))) tb_writer.add_scalar('train/global_average_objf', total_objf / total_frames, global_batch_idx_train) tb_writer.add_scalar('train/current_batch_average_objf', curr_batch_objf / (curr_batch_frames + 0.001), global_batch_idx_train) # if batch_idx >= 10: # print("Exiting early to get profile info") # sys.exit(0) if batch_idx > 0 and batch_idx % 200 == 0: total_valid_objf, total_valid_frames, total_valid_all_frames = get_validation_objf( dataloader=valid_dataloader, model=model, P=P, device=device, graph_compiler=graph_compiler) valid_average_objf = total_valid_objf / total_valid_frames model.train() logging.info( 'Validation average objf: {:.6f} over {} frames ({:.1f}% kept)' .format(valid_average_objf, total_valid_frames, 100.0 * total_valid_frames / total_valid_all_frames)) tb_writer.add_scalar('train/global_valid_average_objf', valid_average_objf, global_batch_idx_train) model.write_tensorboard_diagnostics( tb_writer, global_step=global_batch_idx_train) prev_timestamp = datetime.now() return total_objf / total_frames, valid_average_objf, global_batch_idx_train