def run(rank, world_size, args):
    '\n    Args:\n      rank:\n        It is a value between 0 and `world_size-1`, which is\n        passed automatically by `mp.spawn()` in :func:`main`.\n        The node with rank 0 is responsible for saving checkpoint.\n      world_size:\n        Number of GPUs for DDP training.\n      args:\n        The return value of get_parser().parse_args()\n    '
    params = get_params()
    params.update(vars(args))
    fix_random_seed(42)
    if (world_size > 1):
        setup_dist(rank, world_size, params.master_port)
    setup_logger(f'{params.exp_dir}/log/log-train')
    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
    lexicon = Lexicon(params.lang_dir)
    max_phone_id = max(lexicon.tokens)
    device = torch.device('cpu')
    if torch.cuda.is_available():
        device = torch.device('cuda', rank)
    graph_compiler = CtcTrainingGraphCompiler(lexicon=lexicon, device=device)
    model = TdnnLstm(num_features=params.feature_dim, num_classes=(max_phone_id + 1), subsampling_factor=params.subsampling_factor)
    checkpoints = load_checkpoint_if_available(params=params, model=model)
    model.to(device)
    if (world_size > 1):
        model = DDP(model, device_ids=[rank])
    optimizer = optim.AdamW(model.parameters(), lr=params.lr, weight_decay=params.weight_decay)
    scheduler = StepLR(optimizer, step_size=8, gamma=0.1)
    if checkpoints:
        optimizer.load_state_dict(checkpoints['optimizer'])
        scheduler.load_state_dict(checkpoints['scheduler'])
    librispeech = LibriSpeechAsrDataModule(args)
    train_dl = librispeech.train_dataloaders()
    valid_dl = librispeech.valid_dataloaders()
    for epoch in range(params.start_epoch, params.num_epochs):
        train_dl.sampler.set_epoch(epoch)
        if (epoch > params.start_epoch):
            logging.info(f'epoch {epoch}, lr: {scheduler.get_last_lr()[0]}')
        if (tb_writer is not None):
            tb_writer.add_scalar('train/lr', scheduler.get_last_lr()[0], params.batch_idx_train)
            tb_writer.add_scalar('train/epoch', epoch, params.batch_idx_train)
        params.cur_epoch = epoch
        train_one_epoch(params=params, model=model, optimizer=optimizer, graph_compiler=graph_compiler, train_dl=train_dl, valid_dl=valid_dl, tb_writer=tb_writer, world_size=world_size)
        scheduler.step()
        save_checkpoint(params=params, model=model, optimizer=optimizer, scheduler=scheduler, rank=rank)
    logging.info('Done!')
    if (world_size > 1):
        torch.distributed.barrier()
        cleanup_dist()
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()
def run(rank, world_size, args):
    '\n    Args:\n      rank:\n        It is a value between 0 and `world_size-1`, which is\n        passed automatically by `mp.spawn()` in :func:`main`.\n        The node with rank 0 is responsible for saving checkpoint.\n      world_size:\n        Number of GPUs for DDP training.\n      args:\n        The return value of get_parser().parse_args()\n    '
    params = get_params()
    params.update(vars(args))
    fix_random_seed(42)
    if (world_size > 1):
        setup_dist(rank, world_size, params.master_port)
    setup_logger(f'{params.exp_dir}/log/log-train')
    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
    lexicon = Lexicon(params.lang_dir)
    max_phone_id = max(lexicon.tokens)
    device = torch.device('cpu')
    if torch.cuda.is_available():
        device = torch.device('cuda', rank)
    graph_compiler = CtcTrainingGraphCompiler(lexicon=lexicon, device=device)
    model = TdnnLstm(num_features=params.feature_dim, num_classes=(max_phone_id + 1), subsampling_factor=params.subsampling_factor)
    checkpoints = load_checkpoint_if_available(params=params, model=model)
    model.to(device)
    if (world_size > 1):
        model = DDP(model, device_ids=[rank])
    optimizer = optim.AdamW(model.parameters(), lr=params.lr, weight_decay=params.weight_decay)
    scheduler = StepLR(optimizer, step_size=8, gamma=0.1)
    if checkpoints:
        optimizer.load_state_dict(checkpoints['optimizer'])
        scheduler.load_state_dict(checkpoints['scheduler'])
    librispeech = LibriSpeechAsrDataModule(args)
    train_dl = librispeech.train_dataloaders()
    valid_dl = librispeech.valid_dataloaders()
    for epoch in range(params.start_epoch, params.num_epochs):
        train_dl.sampler.set_epoch(epoch)
        if (epoch > params.start_epoch):
            logging.info(f'epoch {epoch}, lr: {scheduler.get_last_lr()[0]}')
        if (tb_writer is not None):
            tb_writer.add_scalar('train/lr', scheduler.get_last_lr()[0], params.batch_idx_train)
            tb_writer.add_scalar('train/epoch', epoch, params.batch_idx_train)
        params.cur_epoch = epoch
        train_one_epoch(params=params, model=model, optimizer=optimizer, graph_compiler=graph_compiler, train_dl=train_dl, valid_dl=valid_dl, tb_writer=tb_writer, world_size=world_size)
        scheduler.step()
        save_checkpoint(params=params, model=model, optimizer=optimizer, scheduler=scheduler, rank=rank)
    logging.info('Done!')
    if (world_size > 1):
        torch.distributed.barrier()
        cleanup_dist()