def train(FLAGS):
    '\n    FLAGS:\n        saveto: str\n        reload: store_true\n        config_path: str\n        pretrain_path: str, default=""\n        model_name: str\n        log_path: str\n    '
    write_log_to_file(os.path.join(FLAGS.log_path, ('%s.log' % time.strftime('%Y%m%d-%H%M%S'))))
    GlobalNames.USE_GPU = FLAGS.use_gpu
    if GlobalNames.USE_GPU:
        CURRENT_DEVICE = 'cpu'
    else:
        CURRENT_DEVICE = 'cuda:0'
    config_path = os.path.abspath(FLAGS.config_path)
    with open(config_path.strip()) as f:
        configs = yaml.load(f)
    INFO(pretty_configs(configs))
    configs = default_configs(configs)
    data_configs = configs['data_configs']
    model_configs = configs['model_configs']
    optimizer_configs = configs['optimizer_configs']
    training_configs = configs['training_configs']
    GlobalNames.SEED = training_configs['seed']
    set_seed(GlobalNames.SEED)
    best_model_prefix = os.path.join(FLAGS.saveto, (FLAGS.model_name + GlobalNames.MY_BEST_MODEL_SUFFIX))
    timer = Timer()
    INFO('Loading data...')
    timer.tic()
    vocab_src = Vocabulary(**data_configs['vocabularies'][0])
    train_batch_size = (training_configs['batch_size'] * max(1, training_configs['update_cycle']))
    train_buffer_size = (training_configs['buffer_size'] * max(1, training_configs['update_cycle']))
    train_bitext_dataset = ZipDataset(TextLineDataset(data_path=data_configs['train_data'][0], vocabulary=vocab_src, max_len=data_configs['max_len'][0]), shuffle=training_configs['shuffle'])
    valid_bitext_dataset = ZipDataset(TextLineDataset(data_path=data_configs['valid_data'][0], vocabulary=vocab_src))
    training_iterator = DataIterator(dataset=train_bitext_dataset, batch_size=train_batch_size, use_bucket=training_configs['use_bucket'], buffer_size=train_buffer_size, batching_func=training_configs['batching_key'])
    valid_iterator = DataIterator(dataset=valid_bitext_dataset, batch_size=training_configs['valid_batch_size'], use_bucket=True, buffer_size=100000, numbering=True)
    INFO('Done. Elapsed time {0}'.format(timer.toc()))
    lrate = optimizer_configs['learning_rate']
    is_early_stop = False
    model_collections = Collections()
    checkpoint_saver = Saver(save_prefix='{0}.ckpt'.format(os.path.join(FLAGS.saveto, FLAGS.model_name)), num_max_keeping=training_configs['num_kept_checkpoints'])
    best_model_saver = BestKSaver(save_prefix='{0}.best'.format(os.path.join(FLAGS.saveto, FLAGS.model_name)), num_max_keeping=training_configs['num_kept_best_checkpoints'])
    INFO('Building model...')
    timer.tic()
    nmt_model = build_model(n_words=vocab_src.max_n_words, **model_configs)
    INFO(nmt_model)
    params_total = sum([p.numel() for (n, p) in nmt_model.named_parameters()])
    params_with_embedding = sum([p.numel() for (n, p) in nmt_model.named_parameters() if (n.find('embedding') == (- 1))])
    INFO('Total parameters: {}'.format(params_total))
    INFO('Total parameters (excluding word embeddings): {}'.format(params_with_embedding))
    critic = NMTCriterion(label_smoothing=model_configs['label_smoothing'])
    INFO(critic)
    INFO('Done. Elapsed time {0}'.format(timer.toc()))
    if GlobalNames.USE_GPU:
        nmt_model = nmt_model.cuda()
        critic = critic.cuda()
    load_pretrained_model(nmt_model, FLAGS.pretrain_path, exclude_prefix=None, device=CURRENT_DEVICE)
    INFO('Building Optimizer...')
    optim = Optimizer(name=optimizer_configs['optimizer'], model=nmt_model, lr=lrate, grad_clip=optimizer_configs['grad_clip'], optim_args=optimizer_configs['optimizer_params'])
    if (optimizer_configs['schedule_method'] is not None):
        if (optimizer_configs['schedule_method'] == 'loss'):
            scheduler = ReduceOnPlateauScheduler(optimizer=optim, **optimizer_configs['scheduler_configs'])
        elif (optimizer_configs['schedule_method'] == 'noam'):
            scheduler = NoamScheduler(optimizer=optim, **optimizer_configs['scheduler_configs'])
        elif (optimizer_configs['schedule_method'] == 'rsqrt'):
            scheduler = RsqrtScheduler(optimizer=optim, **optimizer_configs['scheduler_configs'])
        else:
            WARN('Unknown scheduler name {0}. Do not use lr_scheduling.'.format(optimizer_configs['schedule_method']))
            scheduler = None
    else:
        scheduler = None
    if (training_configs['moving_average_method'] is not None):
        ma = MovingAverage(moving_average_method=training_configs['moving_average_method'], named_params=nmt_model.named_parameters(), alpha=training_configs['moving_average_alpha'])
    else:
        ma = None
    INFO('Done. Elapsed time {0}'.format(timer.toc()))
    if FLAGS.reload:
        checkpoint_saver.load_latest(model=nmt_model, optim=optim, lr_scheduler=scheduler, collections=model_collections, ma=ma)
    eidx = model_collections.get_collection('eidx', [0])[(- 1)]
    uidx = model_collections.get_collection('uidx', [0])[(- 1)]
    bad_count = model_collections.get_collection('bad_count', [0])[(- 1)]
    oom_count = model_collections.get_collection('oom_count', [0])[(- 1)]
    summary_writer = SummaryWriter(log_dir=FLAGS.log_path)
    cum_samples = 0
    cum_words = 0
    valid_loss = best_valid_loss = float('inf')
    saving_files = []
    timer_for_speed = Timer()
    timer_for_speed.tic()
    INFO('Begin training...')
    while True:
        summary_writer.add_scalar('Epoch', (eidx + 1), uidx)
        training_iter = training_iterator.build_generator()
        training_progress_bar = tqdm(desc=' - (Epc {}, Upd {}) '.format(eidx, uidx), total=len(training_iterator), unit='sents')
        for batch in training_iter:
            uidx += 1
            if ((optimizer_configs['schedule_method'] is not None) and (optimizer_configs['schedule_method'] != 'loss')):
                scheduler.step(global_step=uidx)
            seqs_x = batch
            n_samples_t = len(seqs_x)
            n_words_t = sum((len(s) for s in seqs_x))
            cum_samples += n_samples_t
            cum_words += n_words_t
            train_loss = 0.0
            optim.zero_grad()
            try:
                for (seqs_x_t,) in split_shard(seqs_x, split_size=training_configs['update_cycle']):
                    x = prepare_data(seqs_x_t, cuda=GlobalNames.USE_GPU)
                    loss = compute_forward(model=nmt_model, critic=critic, seqs_x=x, eval=False, normalization=n_samples_t, norm_by_words=training_configs['norm_by_words'])
                    train_loss += (loss / x.size(1))
                optim.step()
            except RuntimeError as e:
                if ('out of memory' in str(e)):
                    print('| WARNING: ran out of memory, skipping batch')
                    oom_count += 1
                    optim.zero_grad()
                else:
                    raise e
            if ((ma is not None) and (eidx >= training_configs['moving_average_start_epoch'])):
                ma.step()
            training_progress_bar.update(n_samples_t)
            training_progress_bar.set_description(' - (Epc {}, Upd {}) '.format(eidx, uidx))
            training_progress_bar.set_postfix_str('TrainLoss: {:.2f}, ValidLoss(best): {:.2f} ({:.2f})'.format(train_loss, valid_loss, best_valid_loss))
            summary_writer.add_scalar('train_loss', scalar_value=train_loss, global_step=uidx)
            if should_trigger_by_steps(uidx, eidx, every_n_step=training_configs['disp_freq']):
                words_per_sec = (cum_words / timer.toc(return_seconds=True))
                sents_per_sec = (cum_samples / timer.toc(return_seconds=True))
                lrate = list(optim.get_lrate())[0]
                summary_writer.add_scalar('Speed(words/sec)', scalar_value=words_per_sec, global_step=uidx)
                summary_writer.add_scalar('Speed(sents/sen)', scalar_value=sents_per_sec, global_step=uidx)
                summary_writer.add_scalar('lrate', scalar_value=lrate, global_step=uidx)
                summary_writer.add_scalar('oom_count', scalar_value=oom_count, global_step=uidx)
                timer.tic()
                cum_words = 0
                cum_samples = 0
            if should_trigger_by_steps(global_step=uidx, n_epoch=eidx, every_n_step=training_configs['loss_valid_freq'], debug=FLAGS.debug):
                if (ma is not None):
                    origin_state_dict = deepcopy(nmt_model.state_dict())
                    nmt_model.load_state_dict(ma.export_ma_params(), strict=False)
                valid_loss = loss_validation(model=nmt_model, critic=critic, valid_iterator=valid_iterator)
                model_collections.add_to_collection('history_losses', valid_loss)
                min_history_loss = np.array(model_collections.get_collection('history_losses')).min()
                summary_writer.add_scalar('loss', valid_loss, global_step=uidx)
                summary_writer.add_scalar('best_loss', min_history_loss, global_step=uidx)
                best_valid_loss = min_history_loss
                if (ma is not None):
                    nmt_model.load_state_dict(origin_state_dict)
                    del origin_state_dict
                if (optimizer_configs['schedule_method'] == 'loss'):
                    scheduler.step(global_step=uidx, metric=best_valid_loss)
                if (valid_loss < best_valid_loss):
                    bad_count = 0
                    if (is_early_stop is False):
                        torch.save(nmt_model.state_dict(), (best_model_prefix + '.final'))
                        model_collections.add_to_collection('uidx', uidx)
                        model_collections.add_to_collection('eidx', eidx)
                        model_collections.add_to_collection('bad_count', bad_count)
                        best_model_saver.save(global_step=uidx, metric=valid_loss, model=nmt_model, optim=optim, lr_scheduler=scheduler, collections=model_collections, ma=ma)
                else:
                    bad_count += 1
                    if ((bad_count >= training_configs['early_stop_patience']) and (eidx > 0)):
                        is_early_stop = True
                        WARN('Early Stop!')
                summary_writer.add_scalar('bad_count', bad_count, uidx)
                INFO('{0} Loss: {1:.2f} lrate: {2:6f} patience: {3}'.format(uidx, valid_loss, lrate, bad_count))
            if should_trigger_by_steps(uidx, eidx, every_n_step=training_configs['save_freq'], debug=FLAGS.debug):
                model_collections.add_to_collection('uidx', uidx)
                model_collections.add_to_collection('eidx', eidx)
                model_collections.add_to_collection('bad_count', bad_count)
                if (not is_early_stop):
                    checkpoint_saver.save(global_step=uidx, model=nmt_model, optim=optim, lr_scheduler=scheduler, collections=model_collections, ma=ma)
        training_progress_bar.close()
        eidx += 1
        if (eidx > training_configs['max_epochs']):
            break
FLAGS:
    saveto: str
    reload: store_true
    config_path: str
    pretrain_path: str, default=""
    model_name: str
    log_path: str
def train(FLAGS):
    '\n    FLAGS:\n        saveto: str\n        reload: store_true\n        config_path: str\n        pretrain_path: str, default=\n        model_name: str\n        log_path: str\n    '
    write_log_to_file(os.path.join(FLAGS.log_path, ('%s.log' % time.strftime('%Y%m%d-%H%M%S'))))
    GlobalNames.USE_GPU = FLAGS.use_gpu
    if GlobalNames.USE_GPU:
        CURRENT_DEVICE = 'cpu'
    else:
        CURRENT_DEVICE = 'cuda:0'
    config_path = os.path.abspath(FLAGS.config_path)
    with open(config_path.strip()) as f:
        configs = yaml.load(f)
    INFO(pretty_configs(configs))
    configs = default_configs(configs)
    data_configs = configs['data_configs']
    model_configs = configs['model_configs']
    optimizer_configs = configs['optimizer_configs']
    training_configs = configs['training_configs']
    GlobalNames.SEED = training_configs['seed']
    set_seed(GlobalNames.SEED)
    best_model_prefix = os.path.join(FLAGS.saveto, (FLAGS.model_name + GlobalNames.MY_BEST_MODEL_SUFFIX))
    timer = Timer()
    INFO('Loading data...')
    timer.tic()
    vocab_src = Vocabulary(**data_configs['vocabularies'][0])
    train_batch_size = (training_configs['batch_size'] * max(1, training_configs['update_cycle']))
    train_buffer_size = (training_configs['buffer_size'] * max(1, training_configs['update_cycle']))
    train_bitext_dataset = ZipDataset(TextLineDataset(data_path=data_configs['train_data'][0], vocabulary=vocab_src, max_len=data_configs['max_len'][0]), shuffle=training_configs['shuffle'])
    valid_bitext_dataset = ZipDataset(TextLineDataset(data_path=data_configs['valid_data'][0], vocabulary=vocab_src))
    training_iterator = DataIterator(dataset=train_bitext_dataset, batch_size=train_batch_size, use_bucket=training_configs['use_bucket'], buffer_size=train_buffer_size, batching_func=training_configs['batching_key'])
    valid_iterator = DataIterator(dataset=valid_bitext_dataset, batch_size=training_configs['valid_batch_size'], use_bucket=True, buffer_size=100000, numbering=True)
    INFO('Done. Elapsed time {0}'.format(timer.toc()))
    lrate = optimizer_configs['learning_rate']
    is_early_stop = False
    model_collections = Collections()
    checkpoint_saver = Saver(save_prefix='{0}.ckpt'.format(os.path.join(FLAGS.saveto, FLAGS.model_name)), num_max_keeping=training_configs['num_kept_checkpoints'])
    best_model_saver = BestKSaver(save_prefix='{0}.best'.format(os.path.join(FLAGS.saveto, FLAGS.model_name)), num_max_keeping=training_configs['num_kept_best_checkpoints'])
    INFO('Building model...')
    timer.tic()
    nmt_model = build_model(n_words=vocab_src.max_n_words, **model_configs)
    INFO(nmt_model)
    params_total = sum([p.numel() for (n, p) in nmt_model.named_parameters()])
    params_with_embedding = sum([p.numel() for (n, p) in nmt_model.named_parameters() if (n.find('embedding') == (- 1))])
    INFO('Total parameters: {}'.format(params_total))
    INFO('Total parameters (excluding word embeddings): {}'.format(params_with_embedding))
    critic = NMTCriterion(label_smoothing=model_configs['label_smoothing'])
    INFO(critic)
    INFO('Done. Elapsed time {0}'.format(timer.toc()))
    if GlobalNames.USE_GPU:
        nmt_model = nmt_model.cuda()
        critic = critic.cuda()
    load_pretrained_model(nmt_model, FLAGS.pretrain_path, exclude_prefix=None, device=CURRENT_DEVICE)
    INFO('Building Optimizer...')
    optim = Optimizer(name=optimizer_configs['optimizer'], model=nmt_model, lr=lrate, grad_clip=optimizer_configs['grad_clip'], optim_args=optimizer_configs['optimizer_params'])
    if (optimizer_configs['schedule_method'] is not None):
        if (optimizer_configs['schedule_method'] == 'loss'):
            scheduler = ReduceOnPlateauScheduler(optimizer=optim, **optimizer_configs['scheduler_configs'])
        elif (optimizer_configs['schedule_method'] == 'noam'):
            scheduler = NoamScheduler(optimizer=optim, **optimizer_configs['scheduler_configs'])
        elif (optimizer_configs['schedule_method'] == 'rsqrt'):
            scheduler = RsqrtScheduler(optimizer=optim, **optimizer_configs['scheduler_configs'])
        else:
            WARN('Unknown scheduler name {0}. Do not use lr_scheduling.'.format(optimizer_configs['schedule_method']))
            scheduler = None
    else:
        scheduler = None
    if (training_configs['moving_average_method'] is not None):
        ma = MovingAverage(moving_average_method=training_configs['moving_average_method'], named_params=nmt_model.named_parameters(), alpha=training_configs['moving_average_alpha'])
    else:
        ma = None
    INFO('Done. Elapsed time {0}'.format(timer.toc()))
    if FLAGS.reload:
        checkpoint_saver.load_latest(model=nmt_model, optim=optim, lr_scheduler=scheduler, collections=model_collections, ma=ma)
    eidx = model_collections.get_collection('eidx', [0])[(- 1)]
    uidx = model_collections.get_collection('uidx', [0])[(- 1)]
    bad_count = model_collections.get_collection('bad_count', [0])[(- 1)]
    oom_count = model_collections.get_collection('oom_count', [0])[(- 1)]
    summary_writer = SummaryWriter(log_dir=FLAGS.log_path)
    cum_samples = 0
    cum_words = 0
    valid_loss = best_valid_loss = float('inf')
    saving_files = []
    timer_for_speed = Timer()
    timer_for_speed.tic()
    INFO('Begin training...')
    while True:
        summary_writer.add_scalar('Epoch', (eidx + 1), uidx)
        training_iter = training_iterator.build_generator()
        training_progress_bar = tqdm(desc=' - (Epc {}, Upd {}) '.format(eidx, uidx), total=len(training_iterator), unit='sents')
        for batch in training_iter:
            uidx += 1
            if ((optimizer_configs['schedule_method'] is not None) and (optimizer_configs['schedule_method'] != 'loss')):
                scheduler.step(global_step=uidx)
            seqs_x = batch
            n_samples_t = len(seqs_x)
            n_words_t = sum((len(s) for s in seqs_x))
            cum_samples += n_samples_t
            cum_words += n_words_t
            train_loss = 0.0
            optim.zero_grad()
            try:
                for (seqs_x_t,) in split_shard(seqs_x, split_size=training_configs['update_cycle']):
                    x = prepare_data(seqs_x_t, cuda=GlobalNames.USE_GPU)
                    loss = compute_forward(model=nmt_model, critic=critic, seqs_x=x, eval=False, normalization=n_samples_t, norm_by_words=training_configs['norm_by_words'])
                    train_loss += (loss / x.size(1))
                optim.step()
            except RuntimeError as e:
                if ('out of memory' in str(e)):
                    print('| WARNING: ran out of memory, skipping batch')
                    oom_count += 1
                    optim.zero_grad()
                else:
                    raise e
            if ((ma is not None) and (eidx >= training_configs['moving_average_start_epoch'])):
                ma.step()
            training_progress_bar.update(n_samples_t)
            training_progress_bar.set_description(' - (Epc {}, Upd {}) '.format(eidx, uidx))
            training_progress_bar.set_postfix_str('TrainLoss: {:.2f}, ValidLoss(best): {:.2f} ({:.2f})'.format(train_loss, valid_loss, best_valid_loss))
            summary_writer.add_scalar('train_loss', scalar_value=train_loss, global_step=uidx)
            if should_trigger_by_steps(uidx, eidx, every_n_step=training_configs['disp_freq']):
                words_per_sec = (cum_words / timer.toc(return_seconds=True))
                sents_per_sec = (cum_samples / timer.toc(return_seconds=True))
                lrate = list(optim.get_lrate())[0]
                summary_writer.add_scalar('Speed(words/sec)', scalar_value=words_per_sec, global_step=uidx)
                summary_writer.add_scalar('Speed(sents/sen)', scalar_value=sents_per_sec, global_step=uidx)
                summary_writer.add_scalar('lrate', scalar_value=lrate, global_step=uidx)
                summary_writer.add_scalar('oom_count', scalar_value=oom_count, global_step=uidx)
                timer.tic()
                cum_words = 0
                cum_samples = 0
            if should_trigger_by_steps(global_step=uidx, n_epoch=eidx, every_n_step=training_configs['loss_valid_freq'], debug=FLAGS.debug):
                if (ma is not None):
                    origin_state_dict = deepcopy(nmt_model.state_dict())
                    nmt_model.load_state_dict(ma.export_ma_params(), strict=False)
                valid_loss = loss_validation(model=nmt_model, critic=critic, valid_iterator=valid_iterator)
                model_collections.add_to_collection('history_losses', valid_loss)
                min_history_loss = np.array(model_collections.get_collection('history_losses')).min()
                summary_writer.add_scalar('loss', valid_loss, global_step=uidx)
                summary_writer.add_scalar('best_loss', min_history_loss, global_step=uidx)
                best_valid_loss = min_history_loss
                if (ma is not None):
                    nmt_model.load_state_dict(origin_state_dict)
                    del origin_state_dict
                if (optimizer_configs['schedule_method'] == 'loss'):
                    scheduler.step(global_step=uidx, metric=best_valid_loss)
                if (valid_loss < best_valid_loss):
                    bad_count = 0
                    if (is_early_stop is False):
                        torch.save(nmt_model.state_dict(), (best_model_prefix + '.final'))
                        model_collections.add_to_collection('uidx', uidx)
                        model_collections.add_to_collection('eidx', eidx)
                        model_collections.add_to_collection('bad_count', bad_count)
                        best_model_saver.save(global_step=uidx, metric=valid_loss, model=nmt_model, optim=optim, lr_scheduler=scheduler, collections=model_collections, ma=ma)
                else:
                    bad_count += 1
                    if ((bad_count >= training_configs['early_stop_patience']) and (eidx > 0)):
                        is_early_stop = True
                        WARN('Early Stop!')
                summary_writer.add_scalar('bad_count', bad_count, uidx)
                INFO('{0} Loss: {1:.2f} lrate: {2:6f} patience: {3}'.format(uidx, valid_loss, lrate, bad_count))
            if should_trigger_by_steps(uidx, eidx, every_n_step=training_configs['save_freq'], debug=FLAGS.debug):
                model_collections.add_to_collection('uidx', uidx)
                model_collections.add_to_collection('eidx', eidx)
                model_collections.add_to_collection('bad_count', bad_count)
                if (not is_early_stop):
                    checkpoint_saver.save(global_step=uidx, model=nmt_model, optim=optim, lr_scheduler=scheduler, collections=model_collections, ma=ma)
        training_progress_bar.close()
        eidx += 1
        if (eidx > training_configs['max_epochs']):
            break