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
