train_transformer.py

import os
import sys
import time
from collections import OrderedDict
from time import strftime, gmtime
from tensorboardX import SummaryWriter
from dataset import AsrDataset, DataLoader, AsrCollator
from models.transformer import Model
import hparams as hp
import argparse
import megengine as mge
import megengine.module as M
import megengine.functional as F
from megengine.functional import clip, concat, minimum, norm
from megengine.core._imperative_rt.core2 import pop_scope, push_scope
from typing import Iterable, Union
from megengine.tensor import Tensor
import megengine.distributed as dist
from megengine.data import SequentialSampler, RandomSampler, DataLoader
from criterions.label_smoothing_loss import LabelSmoothingLoss
from megengine.utils.network import Network as Net
import megengine.autodiff as autodiff
import megengine.data as data
import megengine
import multiprocessing

logging = megengine.logger.get_logger()


def clip_grad_norm(
    tensors: Union[Tensor, Iterable[Tensor]],
    max_norm: float,
    ord: float = 2.0,
):
    push_scope("clip_grad_norm")
    if isinstance(tensors, Tensor):
        tensors = [tensors]
    tensors = [t for t in tensors if t.grad is not None]
    norm_ = [norm(t.grad.flatten(), ord=ord) for t in tensors]
    if len(norm_) > 1:
        norm_ = norm(concat(norm_), ord=ord)
    else:
        norm_ = norm_[0]
    scale = max_norm / (norm_ + 1e-6)
    scale = minimum(scale, 1)
    for tensor in tensors:
        tensor.grad._reset(tensor.grad * scale)
    pop_scope("clip_grad_norm")
    return norm_


class exponential_ma:
    def __init__(self, ratio):
        self.value = 0
        self.weight = 0
        self.ratio = ratio

    def update(self, x):
        self.value = self.value * self.ratio + (1 - self.ratio) * x
        self.weight = self.weight * self.ratio + (1 - self.ratio)

    def get_value(self):
        if self.weight < 1e-8:
            return 0
        return self.value / self.weight


def update_train_log(monitor_vars_name, ma_dict, losses, ttrain, tdata):
    for n in monitor_vars_name:
        for ma in ma_dict["losses"]:
            ma[n].update(losses[n])
    for ma in ma_dict["ttrain"]:
        ma.update(ttrain)
    for ma in ma_dict["tdata"]:
        ma.update(tdata)


def print_train_log(sess, epoch, minibatch, ma_dict, minibatch_per_epoch):
    ma_output = "[{}] e:{}, {}/{} ".format(
        strftime("%Y-%m-%d %H:%M:%S", gmtime()), epoch, minibatch, minibatch_per_epoch
    )
    print(ma_output, file=sys.stderr)
    line = "    {:31}:".format("speed")
    for ma in ma_dict["ttrain"]:
        line += "{:10.2g}".format(1 / ma.get_value())
    print(line, file=sys.stderr)
    line = "    {:31}".format("dp/tot")
    for ma1, ma2 in zip(ma_dict["ttrain"], ma_dict["tdata"]):
        line += "{:10.2g}".format(ma2.get_value() / ma1.get_value())
    print(line, file=sys.stderr)
    for k in sess.loss_names:
        line = "    {:31}".format(k)
        for ma in ma_dict["losses"]:
            line += "{:10.2E}".format(ma[k].get_value())
        print(line, file=sys.stderr)
    line = "    {:31}:    {}".format("lr", sess.get_learning_rate())
    print(line, file=sys.stderr)
    sys.stderr.flush()


def adjust_learning_rate(optimizer, step_num, warmup_step=4000):
    lr = (
        hp.lr
        * warmup_step ** 0.5
        * min(step_num * warmup_step ** -1.5, step_num ** -0.5)
    )
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr


def set_grad(net, min, max):
    for param in net.parameters():
        param.grad = mge.random.uniform(min, max, param.shape)
        param.grad_backup = F.copy(param.grad)


class Session:
    def __init__(self, args):
        with open(os.path.join(hp.dataset_root, "vocab.txt")) as f:
            self.vocab = [w.strip() for w in f.readlines()]
            self.vocab = ["<pad>"] + self.vocab
            print(f"Vocab Size: {len(self.vocab)}")
        self.model = Model(hp.num_mels, len(self.vocab))
        world_size = args.world_size * args.ngpus
        if world_size > 1:
            dist.bcast_list_(self.model.parameters(), dist.WORLD)

        # Autodiff gradient manager
        self.gm = autodiff.GradManager().attach(
            self.model.parameters(),
            callbacks=dist.make_allreduce_cb("SUM") if world_size > 1 else None,
        )
        self.global_step = 0
        self.optimizer = mge.optimizer.Adam(self.model.parameters(), lr=hp.lr)

        # load pretrain model
        if args.continue_path:
            ckpt = mge.load(args.continue_path)
            if "model" in ckpt:
                state_dict = ckpt["model"]
                self.model.load_state_dict(state_dict, strict=False)

        self.loss_names = ["total"]
        self.criterion = LabelSmoothingLoss(len(self.vocab), 0, hp.lsm_weight)

    def get_learning_rate(self):
        lr = self.optimizer.param_groups[0]["lr"]
        return lr

    def get_current_losses(self):
        losses = OrderedDict()
        for name in self.loss_names:
            losses[name] = float(getattr(self, "loss_" + name))
        return losses

    def optimize_parameters(self, data):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        text_input, text_output, mel, pos_text, pos_mel, text_length, mel_length = data
        with self.gm:
            hs_pad, hs_mask, pred_pad, pred_mask = self.model.forward(
                mel, mel_length, text_input, text_length
            )
            self.loss_total = self.criterion(pred_pad, text_output)
            self.gm.backward(self.loss_total)
        clip_grad_norm(self.model.parameters(), 1.0)
        self.optimizer.step().clear_grad()


def main():
    os.makedirs(hp.checkpoint_path, exist_ok=True)
    parser = argparse.ArgumentParser()
    parser.add_argument("--continue_path")

    parser.add_argument(
        "-n",
        "--ngpus",
        default=None,
        type=int,
        help="number of GPUs per node (default: None, use all available GPUs)",
    )
    parser.add_argument(
        "--save",
        metavar="DIR",
        default="output",
        help="path to save checkpoint and log",
    )
    parser.add_argument(
        "--epochs",
        default=90,
        type=int,
        help="number of total epochs to run (default: 90)",
    )

    parser.add_argument("-j", "--workers", default=2, type=int)
    parser.add_argument(
        "-p",
        "--print-freq",
        default=20,
        type=int,
        metavar="N",
        help="print frequency (default: 10)",
    )
    parser.add_argument("--dist-addr", default="localhost")
    parser.add_argument("--dist-port", default=23456, type=int)
    parser.add_argument("--world-size", default=1, type=int)
    parser.add_argument("--rank", default=0, type=int)

    args = parser.parse_args()

    # create server if is master
    if args.rank <= 0:
        server = dist.Server(
            port=args.dist_port
        )  # pylint: disable=unused-variable  # noqa: F841

    # get device count
    with multiprocessing.Pool(1) as pool:
        ngpus_per_node, _ = pool.map(megengine.get_device_count, ["gpu", "cpu"])
    if args.ngpus:
        ngpus_per_node = args.ngpus
    # launch processes
    procs = []
    for local_rank in range(ngpus_per_node):
        p = multiprocessing.Process(
            target=worker,
            kwargs=dict(
                rank=args.rank * ngpus_per_node + local_rank,
                world_size=args.world_size * ngpus_per_node,
                ngpus_per_node=ngpus_per_node,
                args=args,
            ),
        )
        p.start()
        procs.append(p)

    # join processes
    for p in procs:
        p.join()


def worker(rank, world_size, ngpus_per_node, args):

    # pylint: disable=too-many-statements
    if rank == 0:
        os.makedirs(os.path.join(args.save, "asr"), exist_ok=True)
        megengine.logger.set_log_file(os.path.join(args.save, "asr", "log.txt"))
    # init process group
    if world_size > 1:
        dist.init_process_group(
            master_ip=args.dist_addr,
            port=args.dist_port,
            world_size=world_size,
            rank=rank,
            device=rank % ngpus_per_node,
            backend="nccl",
        )
        logging.info(
            "init process group rank %d / %d", dist.get_rank(), dist.get_world_size()
        )

    # build dataset
    train_dataloader = build_dataset(args)
    train_queue = iter(train_dataloader)
    steps_per_epoch = 164905 // (world_size * hp.batch_size)
    sess = Session(args)
    ma_rates = [1 - 0.01 ** x for x in range(3)]
    ma_dict = {
        "losses": [
            {k: exponential_ma(rate) for k in sess.loss_names} for rate in ma_rates
        ],
        "ttrain": [exponential_ma(rate) for rate in ma_rates],
        "tdata": [exponential_ma(rate) for rate in ma_rates],
    }
    for epoch in range(1, (hp.epochs + 1) * steps_per_epoch):
        t_minibatch_start = time.time()
        sess.global_step += 1
        if sess.global_step < 400000:
            adjust_learning_rate(sess.optimizer, sess.global_step)
        tdata = time.time() - t_minibatch_start
        data = next(train_queue)
        sess.optimize_parameters(data)
        losses = sess.get_current_losses()
        ttrain = time.time() - t_minibatch_start
        # print(ttrain, tdata)
        update_train_log(sess.loss_names, ma_dict, losses, ttrain, tdata)
        if sess.global_step % hp.log_interval == 0 and rank == 0:
            print_train_log(sess, epoch, epoch, ma_dict, hp.epochs * steps_per_epoch)
        if sess.global_step % hp.save_interval == 0 and rank == 0:
            print("*******************************************")
            mge.save(
                {"model": sess.model.state_dict(), "global_step": sess.global_step},
                os.path.join(
                    hp.checkpoint_path, "checkpoint_%d.pkl" % sess.global_step
                ),
            )
            print("*******************************************")
        if sess.global_step > hp.max_steps:
            exit(1)


def build_dataset(args):

    dataset = AsrDataset()
    train_sampler = data.Infinite(
        RandomSampler(dataset=dataset, batch_size=hp.batch_size)
    )
    dataloader = DataLoader(
        dataset=dataset, sampler=train_sampler, collator=AsrCollator()
    )

    return dataloader


if __name__ == "__main__":
    main()