From 779589a2de84e4e70d6722e870558d651bdd3e0f Mon Sep 17 00:00:00 2001 From: Erwan Date: Tue, 14 Jun 2022 08:49:50 +0200 Subject: [PATCH] Add RNN params to parser --- .../ASR/rnn_lm/compute_perplexity.py | 1 - egs/librispeech/ASR/rnn_lm/dataset.py | 6 +- egs/librispeech/ASR/rnn_lm/test_dataset.py | 5 +- egs/librispeech/ASR/rnn_lm/test_model.py | 2 +- egs/librispeech/ASR/rnn_lm/train-small.py | 607 ------------------ egs/librispeech/ASR/rnn_lm/train.py | 94 ++- 6 files changed, 71 insertions(+), 644 deletions(-) delete mode 100755 egs/librispeech/ASR/rnn_lm/train-small.py diff --git a/egs/librispeech/ASR/rnn_lm/compute_perplexity.py b/egs/librispeech/ASR/rnn_lm/compute_perplexity.py index ee64ca0d5..e754b9534 100755 --- a/egs/librispeech/ASR/rnn_lm/compute_perplexity.py +++ b/egs/librispeech/ASR/rnn_lm/compute_perplexity.py @@ -144,7 +144,6 @@ def main(): args.lm_data = Path(args.lm_data) params = AttributeDict(vars(args)) - print(params) setup_logger(f"{params.exp_dir}/log-ppl/") logging.info("Computing perplexity started") diff --git a/egs/librispeech/ASR/rnn_lm/dataset.py b/egs/librispeech/ASR/rnn_lm/dataset.py index 04adf0ae1..0fc5b3d84 100644 --- a/egs/librispeech/ASR/rnn_lm/dataset.py +++ b/egs/librispeech/ASR/rnn_lm/dataset.py @@ -70,9 +70,9 @@ class LmDataset(torch.utils.data.Dataset): # in the worst case, the subsequent sentences also have # this number of tokens, we should reduce the batch size # so that this batch will not contain too many tokens - actucal_batch_size = batch_size // sz + 1 - actucal_batch_size = min(actucal_batch_size, batch_size) - end = cur + actucal_batch_size + actual_batch_size = batch_size // sz + 1 + actual_batch_size = min(actual_batch_size, batch_size) + end = cur + actual_batch_size end = min(end, num_sentences) this_batch_indexes = torch.arange(cur, end).tolist() batch_indexes.append(this_batch_indexes) diff --git a/egs/librispeech/ASR/rnn_lm/test_dataset.py b/egs/librispeech/ASR/rnn_lm/test_dataset.py index 7d08515e0..bf961f54b 100755 --- a/egs/librispeech/ASR/rnn_lm/test_dataset.py +++ b/egs/librispeech/ASR/rnn_lm/test_dataset.py @@ -56,10 +56,7 @@ def main(): max_sent_len=3, batch_size=4, ) - print(dataset.sentences) - print(dataset.words) - print(dataset.batch_indexes) - print(len(dataset)) + collate_fn = LmDatasetCollate(sos_id=1, eos_id=-1, blank_id=0) dataloader = torch.utils.data.DataLoader( dataset, batch_size=1, collate_fn=collate_fn diff --git a/egs/librispeech/ASR/rnn_lm/test_model.py b/egs/librispeech/ASR/rnn_lm/test_model.py index 5a216a3fb..e0876727d 100755 --- a/egs/librispeech/ASR/rnn_lm/test_model.py +++ b/egs/librispeech/ASR/rnn_lm/test_model.py @@ -40,7 +40,7 @@ def test_rnn_lm_model(): ) lengths = torch.tensor([4, 3, 2]) nll_loss = model(x, y, lengths) - print(nll_loss) + """ tensor([[1.1180, 1.3059, 1.2426, 1.7773], [1.4231, 1.2783, 1.7321, 0.0000], diff --git a/egs/librispeech/ASR/rnn_lm/train-small.py b/egs/librispeech/ASR/rnn_lm/train-small.py deleted file mode 100755 index 1225f9b8f..000000000 --- a/egs/librispeech/ASR/rnn_lm/train-small.py +++ /dev/null @@ -1,607 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang) -# -# See ../../../../LICENSE for clarification regarding multiple authors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Usage: - ./rnn_lm/train.py \ - --start-epoch 0 \ - --num-epochs 20 \ - --batch-size 200 \ - -If you want to use DDP training, e.g., a single node with 4 GPUs, -use: - - python -m torch.distributed.launch \ - --use_env \ - --nproc_per_node 4 \ - ./rnn_lm/train.py \ - --use-ddp-launch true \ - --start-epoch 0 \ - --num-epochs 10 \ - --batch-size 200 -""" - -import argparse -import logging -import math -from pathlib import Path -from shutil import copyfile -from typing import Optional, Tuple - -import torch -import torch.multiprocessing as mp -import torch.nn as nn -import torch.optim as optim -from lhotse.utils import fix_random_seed -from rnn_lm.dataset import get_dataloader -from rnn_lm.model import RnnLmModel -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.nn.utils import clip_grad_norm_ -from torch.utils.tensorboard import SummaryWriter - -from icefall.checkpoint import load_checkpoint -from icefall.checkpoint import save_checkpoint as save_checkpoint_impl -from icefall.dist import ( - cleanup_dist, - get_local_rank, - get_rank, - get_world_size, - setup_dist, -) -from icefall.utils import ( - AttributeDict, - MetricsTracker, - get_env_info, - setup_logger, - str2bool, -) - - -def get_parser(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - - parser.add_argument( - "--world-size", - type=int, - default=1, - help="Number of GPUs for DDP training.", - ) - - parser.add_argument( - "--master-port", - type=int, - default=12354, - help="Master port to use for DDP training.", - ) - - parser.add_argument( - "--tensorboard", - type=str2bool, - default=True, - help="Should various information be logged in tensorboard.", - ) - - parser.add_argument( - "--num-epochs", - type=int, - default=10, - help="Number of epochs to train.", - ) - - parser.add_argument( - "--start-epoch", - type=int, - default=0, - help="""Resume training from from this epoch. - If it is positive, it will load checkpoint from - exp_dir/epoch-{start_epoch-1}.pt - """, - ) - - parser.add_argument( - "--exp-dir", - type=str, - default="rnn_lm/exp_small", - help="""The experiment dir. - It specifies the directory where all training related - files, e.g., checkpoints, logs, etc, are saved - """, - ) - - parser.add_argument( - "--batch-size", - type=int, - default=50, - ) - - parser.add_argument( - "--use-ddp-launch", - type=str2bool, - default=False, - help="True if using torch.distributed.launch", - ) - return parser - - -def get_params() -> AttributeDict: - """Return a dict containing training parameters.""" - - params = AttributeDict( - { - # LM training/validation data - "lm_data": "data/lm_training_bpe_500/sorted_lm_data.pt", - "lm_data_valid": "data/lm_training_bpe_500/sorted_lm_data-valid.pt", - "max_sent_len": 200, - "sos_id": 1, - "eos_id": 1, - "blank_id": 0, - # model related - # - # vocab size of the BPE model - "vocab_size": 500, - "embedding_dim": 1024, - "hidden_dim": 1024, - "num_layers": 2, - # - "lr": 1e-3, - "weight_decay": 1e-6, - # - "best_train_loss": float("inf"), - "best_valid_loss": float("inf"), - "best_train_epoch": -1, - "best_valid_epoch": -1, - "batch_idx_train": 0, - "log_interval": 200, - "reset_interval": 2000, - "valid_interval": 30000, - "env_info": get_env_info(), - } - ) - return params - - -def load_checkpoint_if_available( - params: AttributeDict, - model: nn.Module, - optimizer: Optional[torch.optim.Optimizer] = None, - scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None, -) -> None: - """Load checkpoint from file. - - If params.start_epoch is positive, it will load the checkpoint from - `params.start_epoch - 1`. Otherwise, this function does nothing. - - Apart from loading state dict for `model`, `optimizer` and `scheduler`, - it also updates `best_train_epoch`, `best_train_loss`, `best_valid_epoch`, - and `best_valid_loss` in `params`. - - Args: - params: - The return value of :func:`get_params`. - model: - The training model. - optimizer: - The optimizer that we are using. - scheduler: - The learning rate scheduler we are using. - Returns: - Return None. - """ - if params.start_epoch <= 0: - return - - filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt" - logging.info(f"Loading checkpoint: {filename}") - saved_params = load_checkpoint( - filename, - model=model, - optimizer=optimizer, - scheduler=scheduler, - ) - - keys = [ - "best_train_epoch", - "best_valid_epoch", - "batch_idx_train", - "best_train_loss", - "best_valid_loss", - ] - for k in keys: - params[k] = saved_params[k] - - return saved_params - - -def save_checkpoint( - params: AttributeDict, - model: nn.Module, - optimizer: Optional[torch.optim.Optimizer] = None, - scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None, - rank: int = 0, -) -> None: - """Save model, optimizer, scheduler and training stats to file. - - Args: - params: - It is returned by :func:`get_params`. - model: - The training model. - """ - if rank != 0: - return - filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" - save_checkpoint_impl( - filename=filename, - model=model, - params=params, - optimizer=optimizer, - scheduler=scheduler, - rank=rank, - ) - - if params.best_train_epoch == params.cur_epoch: - best_train_filename = params.exp_dir / "best-train-loss.pt" - copyfile(src=filename, dst=best_train_filename) - - if params.best_valid_epoch == params.cur_epoch: - best_valid_filename = params.exp_dir / "best-valid-loss.pt" - copyfile(src=filename, dst=best_valid_filename) - - -def compute_loss( - model: nn.Module, - x: torch.Tensor, - y: torch.Tensor, - sentence_lengths: torch.Tensor, - is_training: bool, -) -> Tuple[torch.Tensor, MetricsTracker]: - """Compute the negative log-likelihood loss given a model and its input. - Args: - model: - The NN model, e.g., RnnLmModel. - x: - A 2-D tensor. Each row contains BPE token IDs for a sentence. Also, - each row starts with SOS ID. - y: - A 2-D tensor. Each row is a shifted version of the corresponding row - in `x` but ends with an EOS ID (before padding). - sentence_lengths: - A 1-D tensor containing number of tokens of each sentence - before padding. - is_training: - True for training. False for validation. - """ - with torch.set_grad_enabled(is_training): - device = model.device - x = x.to(device) - y = y.to(device) - sentence_lengths = sentence_lengths.to(device) - - nll = model(x, y, sentence_lengths) - loss = nll.sum() - - num_tokens = sentence_lengths.sum().item() - - loss_info = MetricsTracker() - # Note: Due to how MetricsTracker() is designed, - # we use "frames" instead of "num_tokens" as a key here - loss_info["frames"] = num_tokens - loss_info["loss"] = loss.detach().item() - return loss, loss_info - - -def compute_validation_loss( - params: AttributeDict, - model: nn.Module, - valid_dl: torch.utils.data.DataLoader, - world_size: int = 1, -) -> MetricsTracker: - """Run the validation process. The validation loss - is saved in `params.valid_loss`. - """ - model.eval() - - tot_loss = MetricsTracker() - - for batch_idx, batch in enumerate(valid_dl): - x, y, sentence_lengths = batch - - loss, loss_info = compute_loss( - model=model, - x=x, - y=y, - sentence_lengths=sentence_lengths, - is_training=False, - ) - - assert loss.requires_grad is False - tot_loss = tot_loss + loss_info - - if world_size > 1: - tot_loss.reduce(loss.device) - - loss_value = tot_loss["loss"] / tot_loss["frames"] - if loss_value < params.best_valid_loss: - params.best_valid_epoch = params.cur_epoch - params.best_valid_loss = loss_value - - return tot_loss - - -def train_one_epoch( - params: AttributeDict, - model: nn.Module, - optimizer: torch.optim.Optimizer, - train_dl: torch.utils.data.DataLoader, - valid_dl: torch.utils.data.DataLoader, - tb_writer: Optional[SummaryWriter] = None, - world_size: int = 1, -) -> None: - """Train the model for one epoch. - - The training loss from the mean of all sentences is saved in - `params.train_loss`. It runs the validation process every - `params.valid_interval` batches. - - Args: - params: - It is returned by :func:`get_params`. - model: - The model for training. - optimizer: - The optimizer we are using. - train_dl: - Dataloader for the training dataset. - valid_dl: - Dataloader for the validation dataset. - tb_writer: - Writer to write log messages to tensorboard. - world_size: - Number of nodes in DDP training. If it is 1, DDP is disabled. - """ - model.train() - - tot_loss = MetricsTracker() - - for batch_idx, batch in enumerate(train_dl): - params.batch_idx_train += 1 - x, y, sentence_lengths = batch - batch_size = x.size(0) - - loss, loss_info = compute_loss( - model=model, - x=x, - y=y, - sentence_lengths=sentence_lengths, - is_training=True, - ) - - # summary stats - tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info - - optimizer.zero_grad() - loss.backward() - clip_grad_norm_(model.parameters(), 5.0, 2.0) - optimizer.step() - - if batch_idx % params.log_interval == 0: - # Note: "frames" here means "num_tokens" - this_batch_ppl = math.exp(loss_info["loss"] / loss_info["frames"]) - tot_ppl = math.exp(tot_loss["loss"] / tot_loss["frames"]) - - logging.info( - f"Epoch {params.cur_epoch}, " - f"batch {batch_idx}, loss[{loss_info}, ppl: {this_batch_ppl}] " - f"tot_loss[{tot_loss}, ppl: {tot_ppl}], " - f"batch size: {batch_size}" - ) - - if tb_writer is not None: - loss_info.write_summary( - tb_writer, "train/current_", params.batch_idx_train - ) - tot_loss.write_summary( - tb_writer, "train/tot_", params.batch_idx_train - ) - - tb_writer.add_scalar( - "train/current_ppl", this_batch_ppl, params.batch_idx_train - ) - - tb_writer.add_scalar( - "train/tot_ppl", tot_ppl, params.batch_idx_train - ) - - if batch_idx > 0 and batch_idx % params.valid_interval == 0: - logging.info("Computing validation loss") - - valid_info = compute_validation_loss( - params=params, - model=model, - valid_dl=valid_dl, - world_size=world_size, - ) - model.train() - - valid_ppl = math.exp(valid_info["loss"] / valid_info["frames"]) - logging.info( - f"Epoch {params.cur_epoch}, validation: {valid_info}, " - f"ppl: {valid_ppl}" - ) - - if tb_writer is not None: - valid_info.write_summary( - tb_writer, "train/valid_", params.batch_idx_train - ) - - tb_writer.add_scalar( - "train/valid_ppl", valid_ppl, params.batch_idx_train - ) - - loss_value = tot_loss["loss"] / tot_loss["frames"] - params.train_loss = loss_value - if params.train_loss < params.best_train_loss: - params.best_train_epoch = params.cur_epoch - params.best_train_loss = params.train_loss - - -def run(rank, world_size, args): - """ - Args: - rank: - It is a value between 0 and `world_size-1`, which is - passed automatically by `mp.spawn()` in :func:`main`. - The node with rank 0 is responsible for saving checkpoint. - world_size: - Number of GPUs for DDP training. - args: - The return value of get_parser().parse_args() - """ - params = get_params() - params.update(vars(args)) - - if params.use_ddp_launch: - local_rank = get_local_rank() - else: - local_rank = rank - - logging.warning( - f"rank: {rank}, world_size: {world_size}, local_rank: {local_rank}" - ) - - fix_random_seed(42) - if world_size > 1: - setup_dist(rank, world_size, params.master_port, params.use_ddp_launch) - - setup_logger( - f"{params.exp_dir}/log/log-train", rank=rank, world_size=world_size - ) - logging.info("Training started") - logging.info(params) - - if args.tensorboard and rank == 0: - tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") - else: - tb_writer = None - - device = torch.device("cpu") - if torch.cuda.is_available(): - device = torch.device("cuda", local_rank) - - logging.info(f"Device: {device}, rank: {rank}, local_rank: {local_rank}") - - logging.info("About to create model") - model = RnnLmModel( - vocab_size=params.vocab_size, - embedding_dim=params.embedding_dim, - hidden_dim=params.hidden_dim, - num_layers=params.num_layers, - ) - - checkpoints = load_checkpoint_if_available(params=params, model=model) - - model.to(device) - if world_size > 1: - model = DDP(model, device_ids=[local_rank]) - - model.device = device - - optimizer = optim.Adam( - model.parameters(), - lr=params.lr, - weight_decay=params.weight_decay, - ) - if checkpoints: - logging.info("Load optimizer state_dict from checkpoint") - optimizer.load_state_dict(checkpoints["optimizer"]) - - logging.info(f"Loading LM training data from {params.lm_data}") - train_dl = get_dataloader( - filename=params.lm_data, - is_distributed=world_size > 1, - params=params, - ) - - logging.info(f"Loading LM validation data from {params.lm_data_valid}") - valid_dl = get_dataloader( - filename=params.lm_data_valid, - is_distributed=world_size > 1, - params=params, - ) - - # Note: No learning rate scheduler is used here - for epoch in range(params.start_epoch, params.num_epochs): - if world_size > 1: - train_dl.sampler.set_epoch(epoch) - - params.cur_epoch = epoch - - train_one_epoch( - params=params, - model=model, - optimizer=optimizer, - train_dl=train_dl, - valid_dl=valid_dl, - tb_writer=tb_writer, - world_size=world_size, - ) - - save_checkpoint( - params=params, - model=model, - optimizer=optimizer, - rank=rank, - ) - - logging.info("Done!") - - if world_size > 1: - torch.distributed.barrier() - cleanup_dist() - - -def main(): - parser = get_parser() - args = parser.parse_args() - args.exp_dir = Path(args.exp_dir) - - if args.use_ddp_launch: - # for torch.distributed.lanunch - rank = get_rank() - world_size = get_world_size() - print(f"rank: {rank}, world_size: {world_size}") - # This following is a hack as the default log level - # is warning - logging.info = logging.warning - run(rank=rank, world_size=world_size, args=args) - return - - world_size = args.world_size - assert world_size >= 1 - if world_size > 1: - mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) - else: - run(rank=0, world_size=1, args=args) - - -torch.set_num_threads(1) -torch.set_num_interop_threads(1) - -if __name__ == "__main__": - main() diff --git a/egs/librispeech/ASR/rnn_lm/train.py b/egs/librispeech/ASR/rnn_lm/train.py index 69c811734..a6640948a 100755 --- a/egs/librispeech/ASR/rnn_lm/train.py +++ b/egs/librispeech/ASR/rnn_lm/train.py @@ -124,6 +124,13 @@ def get_parser(): """, ) + parser.add_argument( + "--use-fp16", + type=str2bool, + default=False, + help="Whether to use half precision training.", + ) + parser.add_argument( "--batch-size", type=int, @@ -136,6 +143,49 @@ def get_parser(): default=False, help="True if using torch.distributed.launch", ) + + parser.add_argument( + "--lm-data", + type=str, + default="data/lm_training_bpe_500/sorted_lm_data.pt", + help="LM training data", + ) + + parser.add_argument( + "--lm-data-valid", + type=str, + default="data/lm_training_bpe_500/sorted_lm_data-valid.pt", + help="LM validation data", + ) + + parser.add_argument( + "--vocab-size", + type=int, + default=500, + help="Vocabulary size of the model", + ) + + parser.add_argument( + "--embedding-dim", + type=int, + default=2048, + help="Embedding dim of the model", + ) + + parser.add_argument( + "--hidden-dim", + type=int, + default=2048, + help="Hidden dim of the model", + ) + + parser.add_argument( + "--num-layers", + type=int, + default=4, + help="Number of RNN layers the model", + ) + return parser @@ -144,24 +194,12 @@ def get_params() -> AttributeDict: params = AttributeDict( { - # LM training/validation data - "lm_data": "data/lm_training_bpe_500/sorted_lm_data.pt", - "lm_data_valid": "data/lm_training_bpe_500/sorted_lm_data-valid.pt", "max_sent_len": 200, "sos_id": 1, "eos_id": 1, "blank_id": 0, - # model related - # - # vocab size of the BPE model - "vocab_size": 500, - "embedding_dim": 2048, - "hidden_dim": 2048, - "num_layers": 4, - # "lr": 1e-3, "weight_decay": 1e-6, - # "best_train_loss": float("inf"), "best_valid_loss": float("inf"), "best_train_epoch": -1, @@ -321,14 +359,14 @@ def compute_validation_loss( for batch_idx, batch in enumerate(valid_dl): x, y, sentence_lengths = batch - - loss, loss_info = compute_loss( - model=model, - x=x, - y=y, - sentence_lengths=sentence_lengths, - is_training=False, - ) + with torch.cuda.amp.autocast(enabled=params.use_fp16): + loss, loss_info = compute_loss( + model=model, + x=x, + y=y, + sentence_lengths=sentence_lengths, + is_training=False, + ) assert loss.requires_grad is False tot_loss = tot_loss + loss_info @@ -383,14 +421,14 @@ def train_one_epoch( params.batch_idx_train += 1 x, y, sentence_lengths = batch batch_size = x.size(0) - - loss, loss_info = compute_loss( - model=model, - x=x, - y=y, - sentence_lengths=sentence_lengths, - is_training=True, - ) + with torch.cuda.amp.autocast(enabled=params.use_fp16): + loss, loss_info = compute_loss( + model=model, + x=x, + y=y, + sentence_lengths=sentence_lengths, + is_training=True, + ) # summary stats tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info