コード例 #1
0
def _setup_context(args, exit_stack):
    if args.use_cpu:
        context = mx.cpu()
    else:
        num_gpus = get_num_gpus()
        check_condition(
            num_gpus >= 1,
            "No GPUs found, consider running on the CPU with --use-cpu "
            "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
            "binary isn't on the path).")
        check_condition(
            len(args.device_ids) == 1,
            "cannot run on multiple devices for now")
        gpu_id = args.device_ids[0]
        if args.disable_device_locking:
            # without locking and a negative device id we just take the first device
            gpu_id = 0
        else:
            if gpu_id < 0:
                # get a single (!) gpu id automatically:
                gpu_ids = exit_stack.enter_context(
                    acquire_gpus([-1], lock_dir=args.lock_dir))
                gpu_id = gpu_ids[0]
        context = mx.gpu(gpu_id)
    return context
コード例 #2
0
ファイル: rpc_server.py プロジェクト: smarthi/sockeye-rpc
def main():
    params = argparse.ArgumentParser(description='Translate CLI')
    arguments.add_translate_cli_args(params)
    arguments.add_bpe_args(params)
    args = params.parse_args()

    with ExitStack() as exit_stack:
        if args.use_cpu:
            context = mx.cpu()
        else:
            num_gpus = get_num_gpus()
            check_condition(num_gpus >= 1,
                            "No GPUs found, consider running on the CPU with --use-cpu "
                            "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
                            "binary isn't on the path).")
            check_condition(len(args.device_ids) == 1, "cannot run on multiple devices for now")
            gpu_id = args.device_ids[0]
            if args.disable_device_locking:
                if gpu_id < 0:
                    # without locking and a negative device id we just take the first device
                    gpu_id = 0
            else:
                gpu_ids = exit_stack.enter_context(acquire_gpus([gpu_id], lock_dir=args.lock_dir))
                gpu_id = gpu_ids[0]

            context = mx.gpu(gpu_id)

        models, source_vocabs, target_vocab = inference.load_models(
            context=context,
            max_input_len=args.max_input_len,
            beam_size=args.beam_size,
            batch_size=args.batch_size,
            model_folders=args.models,
            checkpoints=args.checkpoints,
            softmax_temperature=args.softmax_temperature,
            max_output_length_num_stds=args.max_output_length_num_stds,
            decoder_return_logit_inputs=args.restrict_lexicon is not None,
            cache_output_layer_w_b=args.restrict_lexicon is not None)

        translator = inference.Translator(context=context,
                                          ensemble_mode=args.ensemble_mode,
                                          bucket_source_width=args.bucket_width,
                                          length_penalty=inference.LengthPenalty(args.length_penalty_alpha,
                                                                                 args.length_penalty_beta),
                                          models=models,
                                          source_vocabs=source_vocabs,
                                          target_vocab=target_vocab,
                                          restrict_lexicon=None,
                                          store_beam=False,
                                          strip_unknown_words=args.strip_unknown_words)

        logger.info('Parsing vocabulary')
        sys.stdin = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8')
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', write_through=True, line_buffering=True)

        opened_vocab = codecs.open(args.bpe_vocabulary.name, encoding='utf-8')
        bpe_filtered_vocab = read_vocabulary(opened_vocab, args.bpe_vocabulary_threshold)
        bpe_merges = -1  # Apply all merge operations.
        bpe_separator = '@@'  # Use default BPE separator.
        bpe_glossaries = None # No excluded words.
        bpe = BPE(args.bpe_codes, bpe_merges, bpe_separator, bpe_filtered_vocab, bpe_glossaries)

        logger.info('Starting RPC server.')
        rpc_server = SockeyeRpcServer(translator, bpe)
        rpc_server.serve()
コード例 #3
0
def main():

    params = argparse.ArgumentParser(
        description='Translate from STDIN to STDOUT')
    params.add_argument(
        '--model-prefixes',
        '-m',
        required=False,
        nargs='+',
        help='model prefix(es). Use multiple for ensemble decoding. ' +
        'Model prefix determines config, best epoch params and vocab files.')
    params.add_argument(
        '--epochs',
        '-e',
        required=False,
        default=None,
        type=int,
        nargs='+',
        help=
        'If not given, chooses best epochs/checkpoints for each model. If specified, must have the same '
        + 'length as --model-prefix and be integer')

    params.add_argument(
        '--max-input-len',
        '-n',
        type=int,
        default=None,
        help='Maximum sentence length. Default: value from trained model(s).')
    params.add_argument(
        '--output-type',
        default='translation',
        choices=["translation", "align_plot", "align_text"],
        help=
        'Either print the translation or visualize the alignment. Default: translation'
    )
    params.add_argument('--align-plot-prefix',
                        default="align",
                        help='Prefix used when plotting the alignment.')
    params.add_argument('--log-level',
                        default=logging.INFO,
                        type=int,
                        choices=[logging.INFO, logging.WARN, logging.DEBUG])
    params.add_argument('--beam-size',
                        '-b',
                        type=int,
                        default=1,
                        help='beam size. If == 1, greedy decode')
    params.add_argument(
        '--ensemble-mode',
        type=str,
        default='linear',
        choices=['linear', 'log_linear'],
        help=
        'Ensemble mode: linear or log-linear interpolation of model predictions. Default: linear'
    )
    params.add_argument(
        '--softmax-temperature',
        type=float,
        default=None,
        required=False,
        help=
        'Controls peakiness of model predictions. Values < 1.0 produce peaked predictions, '
        + 'values > 1.0 produce smoothed distributions.')
    params = arguments.add_device_args(params)
    args = params.parse_args()

    args.model_prefixes = ['model/']
    assert args.beam_size > 0, "Beam size must be 1 or greater."
    if args.epochs is not None:
        assert len(args.epochs) == len(
            args.model_prefixes), "must provide epochs for each model"

    #sockeye.utils.setup_logging(args.log_level)
    logging.basicConfig(filename='test.log', level=logging.INFO)

    logging.info("Command: %s", " ".join(sys.argv))
    logging.info("Arguments: %s", args)

    with ExitStack() as exit_stack:
        if args.use_cpu:
            context = mx.cpu()
        else:
            num_gpus = get_num_gpus()
            assert num_gpus > 0, "No GPUs found, consider running on the CPU with --use-cpu " \
                                 "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi " \
                                 "binary isn't on the path)."
            assert len(
                args.device_ids) == 1, "cannot run on multiple devices for now"
            gpu_id = args.device_ids[0]
            if gpu_id < 0:
                # get an automatic gpu id:
                gpu_id = exit_stack.enter_context(acquire_gpu())
            context = mx.gpu(gpu_id)

        translator = sockeye.inference.Translator(
            context, args.ensemble_mode,
            *sockeye.inference.load_models(context, args.max_input_len,
                                           args.beam_size, args.model_prefixes,
                                           args.epochs,
                                           args.softmax_temperature))
        ############ CHANGE HERE ########################
        sample_file = open('train_question_token.txt', 'r')
        encoder_file = open('train_question_encoder.txt', "w")
        #################################################

        for i, line in enumerate(sample_file, 1):
            trans_input = translator.make_input(i, line)
            source, source_length, bucket_key = translator._get_inference_input(
                trans_input.tokens)
            encoded_source, _, _, _, _ = translator.models[0].run_encoder(
                source, source_length, bucket_key)
            last_slice_source = mx.ndarray.mean(encoded_source,
                                                axis=1,
                                                keepdims=True)
            last_slice_source = last_slice_source.reshape((-1, ))
            encoder_file.write(
                " ".join(map(str, last_slice_source.asnumpy())) + "\n")

        encoder_file.close()
コード例 #4
0
def main():
    params = argparse.ArgumentParser(description='CLI to train sockeye sequence-to-sequence models.')
    arguments.add_train_cli_args(params)
    args = params.parse_args()

    # seed the RNGs
    np.random.seed(args.seed)
    random.seed(args.seed)
    mx.random.seed(args.seed)

    if args.use_fused_rnn:
        check_condition(not args.use_cpu, "GPU required for FusedRNN cells")

    check_condition(args.optimized_metric == C.BLEU or args.optimized_metric in args.metrics,
                    "Must optimize either BLEU or one of tracked metrics (--metrics)")

    # Checking status of output folder, resumption, etc.
    # Create temporary logger to console only
    logger = setup_main_logger(__name__, file_logging=False, console=not args.quiet)

    output_folder = os.path.abspath(args.output)
    resume_training = False
    training_state_dir = os.path.join(output_folder, C.TRAINING_STATE_DIRNAME)
    if os.path.exists(output_folder):
        if args.overwrite_output:
            logger.info("Removing existing output folder %s.", output_folder)
            shutil.rmtree(output_folder)
            os.makedirs(output_folder)
        elif os.path.exists(training_state_dir):
            with open(os.path.join(output_folder, C.ARGS_STATE_NAME), "r") as fp:
                old_args = json.load(fp)
            arg_diffs = _dict_difference(vars(args), old_args) | _dict_difference(old_args, vars(args))
            # Remove args that may differ without affecting the training.
            arg_diffs -= set(C.ARGS_MAY_DIFFER)
            # allow different device-ids provided their total count is the same
            if 'device_ids' in arg_diffs and len(old_args['device_ids']) == len(vars(args)['device_ids']):
                arg_diffs.discard('device_ids')
            if not arg_diffs:
                resume_training = True
            else:
                # We do not have the logger yet
                logger.error("Mismatch in arguments for training continuation.")
                logger.error("Differing arguments: %s.", ", ".join(arg_diffs))
                sys.exit(1)
        elif os.path.exists(os.path.join(output_folder, C.PARAMS_BEST_NAME)):
            logger.error("Refusing to overwrite model folder %s as it seems to contain a trained model.", output_folder)
            sys.exit(1)
        else:
            logger.info("The output folder %s already exists, but no training state or parameter file was found. "
                        "Will start training from scratch.", output_folder)
    else:
        os.makedirs(output_folder)

    logger = setup_main_logger(__name__,
                               file_logging=True,
                               console=not args.quiet, path=os.path.join(output_folder, C.LOG_NAME))
    log_sockeye_version(logger)
    log_mxnet_version(logger)
    logger.info("Command: %s", " ".join(sys.argv))
    logger.info("Arguments: %s", args)
    with open(os.path.join(output_folder, C.ARGS_STATE_NAME), "w") as fp:
        json.dump(vars(args), fp)

    with ExitStack() as exit_stack:
        # context
        if args.use_cpu:
            logger.info("Device: CPU")
            context = [mx.cpu()]
        else:
            num_gpus = get_num_gpus()
            check_condition(num_gpus >= 1,
                            "No GPUs found, consider running on the CPU with --use-cpu "
                            "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
                            "binary isn't on the path).")
            if args.disable_device_locking:
                context = expand_requested_device_ids(args.device_ids)
            else:
                context = exit_stack.enter_context(acquire_gpus(args.device_ids, lock_dir=args.lock_dir))
            logger.info("Device(s): GPU %s", context)
            context = [mx.gpu(gpu_id) for gpu_id in context]

        # load existing or create vocabs
        if resume_training:
            vocab_source = vocab.vocab_from_json_or_pickle(os.path.join(output_folder, C.VOCAB_SRC_NAME))
            vocab_target = vocab.vocab_from_json_or_pickle(os.path.join(output_folder, C.VOCAB_TRG_NAME))
        else:
            num_words_source, num_words_target = args.num_words
            word_min_count_source, word_min_count_target = args.word_min_count

            # if the source and target embeddings are tied we build a joint vocabulary:
            if args.weight_tying and C.WEIGHT_TYING_SRC in args.weight_tying_type \
                    and C.WEIGHT_TYING_TRG in args.weight_tying_type:
                vocab_source = vocab_target = _build_or_load_vocab(args.source_vocab,
                                                                   [args.source, args.target],
                                                                   num_words_source,
                                                                   word_min_count_source)
            else:
                vocab_source = _build_or_load_vocab(args.source_vocab, [args.source],
                                                    num_words_source, word_min_count_source)
                vocab_target = _build_or_load_vocab(args.target_vocab, [args.target],
                                                    num_words_target, word_min_count_target)

            # write vocabularies
            vocab.vocab_to_json(vocab_source, os.path.join(output_folder, C.VOCAB_SRC_NAME) + C.JSON_SUFFIX)
            vocab.vocab_to_json(vocab_target, os.path.join(output_folder, C.VOCAB_TRG_NAME) + C.JSON_SUFFIX)

        vocab_source_size = len(vocab_source)
        vocab_target_size = len(vocab_target)
        logger.info("Vocabulary sizes: source=%d target=%d", vocab_source_size, vocab_target_size)

        # create data iterators
        max_seq_len_source, max_seq_len_target = args.max_seq_len
        batch_num_devices = 1 if args.use_cpu else sum(-di if di < 0 else 1 for di in args.device_ids)
        train_iter, eval_iter, config_data = data_io.get_training_data_iters(source=os.path.abspath(args.source),
                                                                             target=os.path.abspath(args.target),
                                                                             validation_source=os.path.abspath(
                                                                                 args.validation_source),
                                                                             validation_target=os.path.abspath(
                                                                                 args.validation_target),
                                                                             vocab_source=vocab_source,
                                                                             vocab_target=vocab_target,
                                                                             vocab_source_path=args.source_vocab,
                                                                             vocab_target_path=args.target_vocab,
                                                                             batch_size=args.batch_size,
                                                                             batch_by_words=args.batch_type == C.BATCH_TYPE_WORD,
                                                                             batch_num_devices=batch_num_devices,
                                                                             fill_up=args.fill_up,
                                                                             max_seq_len_source=max_seq_len_source,
                                                                             max_seq_len_target=max_seq_len_target,
                                                                             bucketing=not args.no_bucketing,
                                                                             bucket_width=args.bucket_width)

        # learning rate scheduling
        learning_rate_half_life = none_if_negative(args.learning_rate_half_life)
        # TODO: The loading for continuation of the scheduler is done separately from the other parts
        if not resume_training:
            lr_scheduler_instance = lr_scheduler.get_lr_scheduler(args.learning_rate_scheduler_type,
                                                                  args.checkpoint_frequency,
                                                                  learning_rate_half_life,
                                                                  args.learning_rate_reduce_factor,
                                                                  args.learning_rate_reduce_num_not_improved,
                                                                  args.learning_rate_schedule,
                                                                  args.learning_rate_warmup)
        else:
            with open(os.path.join(training_state_dir, C.SCHEDULER_STATE_NAME), "rb") as fp:
                lr_scheduler_instance = pickle.load(fp)

        # model configuration
        num_embed_source, num_embed_target = args.num_embed
        encoder_num_layers, decoder_num_layers = args.num_layers

        encoder_embed_dropout, decoder_embed_dropout = args.embed_dropout
        encoder_rnn_dropout_inputs, decoder_rnn_dropout_inputs = args.rnn_dropout_inputs
        encoder_rnn_dropout_states, decoder_rnn_dropout_states = args.rnn_dropout_states
        if encoder_embed_dropout > 0 and encoder_rnn_dropout_inputs > 0:
            logger.warning("Setting encoder RNN AND source embedding dropout > 0 leads to "
                           "two dropout layers on top of each other.")
        if decoder_embed_dropout > 0 and decoder_rnn_dropout_inputs > 0:
            logger.warning("Setting encoder RNN AND source embedding dropout > 0 leads to "
                           "two dropout layers on top of each other.")
        encoder_rnn_dropout_recurrent, decoder_rnn_dropout_recurrent = args.rnn_dropout_recurrent
        if encoder_rnn_dropout_recurrent > 0 or decoder_rnn_dropout_recurrent > 0:
            check_condition(args.rnn_cell_type == C.LSTM_TYPE,
                            "Recurrent dropout without memory loss only supported for LSTMs right now.")

        encoder_transformer_preprocess, decoder_transformer_preprocess = args.transformer_preprocess
        encoder_transformer_postprocess, decoder_transformer_postprocess = args.transformer_postprocess

        config_conv = None
        if args.encoder == C.RNN_WITH_CONV_EMBED_NAME:
            config_conv = encoder.ConvolutionalEmbeddingConfig(num_embed=num_embed_source,
                                                               max_filter_width=args.conv_embed_max_filter_width,
                                                               num_filters=args.conv_embed_num_filters,
                                                               pool_stride=args.conv_embed_pool_stride,
                                                               num_highway_layers=args.conv_embed_num_highway_layers,
                                                               dropout=args.conv_embed_dropout)

        if args.encoder in (C.TRANSFORMER_TYPE, C.TRANSFORMER_WITH_CONV_EMBED_TYPE):
            config_encoder = transformer.TransformerConfig(
                model_size=args.transformer_model_size,
                attention_heads=args.transformer_attention_heads,
                feed_forward_num_hidden=args.transformer_feed_forward_num_hidden,
                num_layers=encoder_num_layers,
                vocab_size=vocab_source_size,
                dropout_attention=args.transformer_dropout_attention,
                dropout_relu=args.transformer_dropout_relu,
                dropout_prepost=args.transformer_dropout_prepost,
                weight_tying=args.weight_tying and C.WEIGHT_TYING_SRC in args.weight_tying_type,
                positional_encodings=not args.transformer_no_positional_encodings,
                preprocess_sequence=encoder_transformer_preprocess,
                postprocess_sequence=encoder_transformer_postprocess,
                conv_config=config_conv)
        else:
            config_encoder = encoder.RecurrentEncoderConfig(
                vocab_size=vocab_source_size,
                num_embed=num_embed_source,
                embed_dropout=encoder_embed_dropout,
                rnn_config=rnn.RNNConfig(cell_type=args.rnn_cell_type,
                                         num_hidden=args.rnn_num_hidden,
                                         num_layers=encoder_num_layers,
                                         dropout_inputs=encoder_rnn_dropout_inputs,
                                         dropout_states=encoder_rnn_dropout_states,
                                         dropout_recurrent=encoder_rnn_dropout_recurrent,
                                         residual=args.rnn_residual_connections,
                                         first_residual_layer=args.rnn_first_residual_layer,
                                         forget_bias=args.rnn_forget_bias),
                conv_config=config_conv,
                reverse_input=args.rnn_encoder_reverse_input)

        decoder_weight_tying = args.weight_tying and C.WEIGHT_TYING_TRG in args.weight_tying_type \
                               and C.WEIGHT_TYING_SOFTMAX in args.weight_tying_type

        if args.decoder == C.TRANSFORMER_TYPE:
            config_decoder = transformer.TransformerConfig(
                model_size=args.transformer_model_size,
                attention_heads=args.transformer_attention_heads,
                feed_forward_num_hidden=args.transformer_feed_forward_num_hidden,
                num_layers=decoder_num_layers,
                vocab_size=vocab_target_size,
                dropout_attention=args.transformer_dropout_attention,
                dropout_relu=args.transformer_dropout_relu,
                dropout_prepost=args.transformer_dropout_prepost,
                weight_tying=decoder_weight_tying,
                positional_encodings=not args.transformer_no_positional_encodings,
                preprocess_sequence=decoder_transformer_preprocess,
                postprocess_sequence=decoder_transformer_postprocess,
                conv_config=None)

        else:
            attention_num_hidden = args.rnn_num_hidden if not args.attention_num_hidden else args.attention_num_hidden
            config_coverage = None
            if args.attention_type == C.ATT_COV:
                config_coverage = coverage.CoverageConfig(type=args.attention_coverage_type,
                                                          num_hidden=args.attention_coverage_num_hidden,
                                                          layer_normalization=args.layer_normalization)
            config_attention = attention.AttentionConfig(type=args.attention_type,
                                                         num_hidden=attention_num_hidden,
                                                         input_previous_word=args.attention_use_prev_word,
                                                         rnn_num_hidden=args.rnn_num_hidden,
                                                         layer_normalization=args.layer_normalization,
                                                         config_coverage=config_coverage,
                                                         num_heads=args.attention_mhdot_heads)
            config_decoder = decoder.RecurrentDecoderConfig(
                vocab_size=vocab_target_size,
                max_seq_len_source=max_seq_len_source,
                num_embed=num_embed_target,
                rnn_config=rnn.RNNConfig(cell_type=args.rnn_cell_type,
                                         num_hidden=args.rnn_num_hidden,
                                         num_layers=decoder_num_layers,
                                         dropout_inputs=decoder_rnn_dropout_inputs,
                                         dropout_states=decoder_rnn_dropout_states,
                                         dropout_recurrent=decoder_rnn_dropout_recurrent,
                                         residual=args.rnn_residual_connections,
                                         first_residual_layer=args.rnn_first_residual_layer,
                                         forget_bias=args.rnn_forget_bias),
                attention_config=config_attention,
                embed_dropout=decoder_embed_dropout,
                hidden_dropout=args.rnn_decoder_hidden_dropout,
                weight_tying=decoder_weight_tying,
                state_init=args.rnn_decoder_state_init,
                context_gating=args.rnn_context_gating,
                layer_normalization=args.layer_normalization,
                attention_in_upper_layers=args.attention_in_upper_layers)

        config_loss = loss.LossConfig(type=args.loss,
                                      vocab_size=vocab_target_size,
                                      normalize=args.normalize_loss,
                                      smoothed_cross_entropy_alpha=args.smoothed_cross_entropy_alpha)

        model_config = model.ModelConfig(config_data=config_data,
                                         max_seq_len_source=max_seq_len_source,
                                         max_seq_len_target=max_seq_len_target,
                                         vocab_source_size=vocab_source_size,
                                         vocab_target_size=vocab_target_size,
                                         config_encoder=config_encoder,
                                         config_decoder=config_decoder,
                                         config_loss=config_loss,
                                         lexical_bias=args.lexical_bias,
                                         learn_lexical_bias=args.learn_lexical_bias,
                                         weight_tying=args.weight_tying,
                                         weight_tying_type=args.weight_tying_type if args.weight_tying else None)
        model_config.freeze()

        # create training model
        training_model = training.TrainingModel(config=model_config,
                                                context=context,
                                                train_iter=train_iter,
                                                fused=args.use_fused_rnn,
                                                bucketing=not args.no_bucketing,
                                                lr_scheduler=lr_scheduler_instance)

        # We may consider loading the params in TrainingModule, for consistency
        # with the training state saving
        if resume_training:
            logger.info("Found partial training in directory %s. Resuming from saved state.", training_state_dir)
            training_model.load_params_from_file(os.path.join(training_state_dir, C.TRAINING_STATE_PARAMS_NAME))
        elif args.params:
            logger.info("Training will initialize from parameters loaded from '%s'", args.params)
            training_model.load_params_from_file(args.params)

        lexicon_array = lexicon.initialize_lexicon(args.lexical_bias,
                                                   vocab_source, vocab_target) if args.lexical_bias else None

        weight_initializer = initializer.get_initializer(args.weight_init, args.weight_init_scale,
                                                         args.rnn_h2h_init, lexicon=lexicon_array)

        optimizer = args.optimizer
        optimizer_params = {'wd': args.weight_decay,
                            "learning_rate": args.initial_learning_rate}
        if lr_scheduler_instance is not None:
            optimizer_params["lr_scheduler"] = lr_scheduler_instance
        clip_gradient = none_if_negative(args.clip_gradient)
        if clip_gradient is not None:
            optimizer_params["clip_gradient"] = clip_gradient
        if args.momentum is not None:
            optimizer_params["momentum"] = args.momentum
        if args.normalize_loss:
            # When normalize_loss is turned on we normalize by the number of non-PAD symbols in a batch which implicitly
            # already contains the number of sentences and therefore we need to disable rescale_grad.
            optimizer_params["rescale_grad"] = 1.0
        else:
            # Making MXNet module API's default scaling factor explicit
            optimizer_params["rescale_grad"] = 1.0 / args.batch_size
        logger.info("Optimizer: %s", optimizer)
        logger.info("Optimizer Parameters: %s", optimizer_params)

        # Handle options that override training settings
        max_updates = args.max_updates
        max_num_checkpoint_not_improved = args.max_num_checkpoint_not_improved
        min_num_epochs = args.min_num_epochs
        # Fixed training schedule always runs for a set number of updates
        if args.learning_rate_schedule:
            max_updates = sum(num_updates for (_, num_updates) in args.learning_rate_schedule)
            max_num_checkpoint_not_improved = -1
            min_num_epochs = 0

        monitor_bleu = args.monitor_bleu
        # Turn on BLEU monitoring when the optimized metric is BLEU and it hasn't been enabled yet
        if args.optimized_metric == C.BLEU and monitor_bleu == 0:
            logger.info("You chose BLEU as the optimized metric, will turn on BLEU monitoring during training. "
                        "To control how many validation sentences are used for calculating bleu use "
                        "the --monitor-bleu argument.")
            monitor_bleu = -1

        training_model.fit(train_iter, eval_iter,
                           output_folder=output_folder,
                           max_params_files_to_keep=args.keep_last_params,
                           metrics=args.metrics,
                           initializer=weight_initializer,
                           max_updates=max_updates,
                           checkpoint_frequency=args.checkpoint_frequency,
                           optimizer=optimizer, optimizer_params=optimizer_params,
                           optimized_metric=args.optimized_metric,
                           max_num_not_improved=max_num_checkpoint_not_improved,
                           min_num_epochs=min_num_epochs,
                           monitor_bleu=monitor_bleu,
                           use_tensorboard=args.use_tensorboard,
                           mxmonitor_pattern=args.monitor_pattern,
                           mxmonitor_stat_func=args.monitor_stat_func)
コード例 #5
0
ファイル: train.py プロジェクト: mengjiexu/sockeye
def main():
    params = argparse.ArgumentParser(
        description='CLI to train sockeye sequence-to-sequence models.')
    arguments.add_io_args(params)
    arguments.add_model_parameters(params)
    arguments.add_training_args(params)
    arguments.add_device_args(params)
    args = params.parse_args()

    # seed the RNGs
    np.random.seed(args.seed)
    random.seed(args.seed)
    mx.random.seed(args.seed)

    if args.use_fused_rnn:
        check_condition(not args.use_cpu, "GPU required for FusedRNN cells")

    if args.rnn_residual_connections:
        check_condition(args.rnn_num_layers > 2,
                        "Residual connections require at least 3 RNN layers")

    check_condition(
        args.optimized_metric == C.BLEU
        or args.optimized_metric in args.metrics,
        "Must optimize either BLEU or one of tracked metrics (--metrics)")

    # Checking status of output folder, resumption, etc.
    # Create temporary logger to console only
    logger = setup_main_logger(__name__,
                               file_logging=False,
                               console=not args.quiet)
    output_folder = os.path.abspath(args.output)
    resume_training = False
    training_state_dir = os.path.join(output_folder, C.TRAINING_STATE_DIRNAME)
    if os.path.exists(output_folder):
        if args.overwrite_output:
            logger.info("Removing existing output folder %s.", output_folder)
            shutil.rmtree(output_folder)
            os.makedirs(output_folder)
        elif os.path.exists(training_state_dir):
            with open(os.path.join(output_folder, C.ARGS_STATE_NAME),
                      "r") as fp:
                old_args = json.load(fp)
            arg_diffs = _dict_difference(
                vars(args), old_args) | _dict_difference(old_args, vars(args))
            # Remove args that may differ without affecting the training.
            arg_diffs -= set(C.ARGS_MAY_DIFFER)
            # allow different device-ids provided their total count is the same
            if 'device_ids' in arg_diffs and len(
                    old_args['device_ids']) == len(vars(args)['device_ids']):
                arg_diffs.discard('device_ids')
            if not arg_diffs:
                resume_training = True
            else:
                # We do not have the logger yet
                logger.error(
                    "Mismatch in arguments for training continuation.")
                logger.error("Differing arguments: %s.", ", ".join(arg_diffs))
                sys.exit(1)
        else:
            logger.error("Refusing to overwrite existing output folder %s.",
                         output_folder)
            sys.exit(1)
    else:
        os.makedirs(output_folder)

    logger = setup_main_logger(__name__,
                               file_logging=True,
                               console=not args.quiet,
                               path=os.path.join(output_folder, C.LOG_NAME))
    log_sockeye_version(logger)
    logger.info("Command: %s", " ".join(sys.argv))
    logger.info("Arguments: %s", args)
    with open(os.path.join(output_folder, C.ARGS_STATE_NAME), "w") as fp:
        json.dump(vars(args), fp)

    with ExitStack() as exit_stack:
        # context
        if args.use_cpu:
            logger.info("Device: CPU")
            context = [mx.cpu()]
        else:
            num_gpus = get_num_gpus()
            check_condition(
                num_gpus >= 1,
                "No GPUs found, consider running on the CPU with --use-cpu "
                "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi "
                "binary isn't on the path).")
            if args.disable_device_locking:
                context = expand_requested_device_ids(args.device_ids)
            else:
                context = exit_stack.enter_context(
                    acquire_gpus(args.device_ids, lock_dir=args.lock_dir))
            logger.info("Device(s): GPU %s", context)
            context = [mx.gpu(gpu_id) for gpu_id in context]

        # load existing or create vocabs
        if resume_training:
            vocab_source = vocab.vocab_from_json_or_pickle(
                os.path.join(output_folder, C.VOCAB_SRC_NAME))
            vocab_target = vocab.vocab_from_json_or_pickle(
                os.path.join(output_folder, C.VOCAB_TRG_NAME))
        else:
            num_words_source = args.num_words if args.num_words_source is None else args.num_words_source
            vocab_source = _build_or_load_vocab(args.source_vocab, args.source,
                                                num_words_source,
                                                args.word_min_count)
            vocab.vocab_to_json(
                vocab_source,
                os.path.join(output_folder, C.VOCAB_SRC_NAME) + C.JSON_SUFFIX)

            num_words_target = args.num_words if args.num_words_target is None else args.num_words_target
            vocab_target = _build_or_load_vocab(args.target_vocab, args.target,
                                                num_words_target,
                                                args.word_min_count)
            vocab.vocab_to_json(
                vocab_target,
                os.path.join(output_folder, C.VOCAB_TRG_NAME) + C.JSON_SUFFIX)

        vocab_source_size = len(vocab_source)
        vocab_target_size = len(vocab_target)
        logger.info("Vocabulary sizes: source=%d target=%d", vocab_source_size,
                    vocab_target_size)

        config_data = data_io.DataConfig(
            os.path.abspath(args.source), os.path.abspath(args.target),
            os.path.abspath(args.validation_source),
            os.path.abspath(args.validation_target), args.source_vocab,
            args.target_vocab)

        # create data iterators
        max_seq_len_source = args.max_seq_len if args.max_seq_len_source is None else args.max_seq_len_source
        max_seq_len_target = args.max_seq_len if args.max_seq_len_target is None else args.max_seq_len_target
        train_iter, eval_iter = data_io.get_training_data_iters(
            source=config_data.source,
            target=config_data.target,
            validation_source=config_data.validation_source,
            validation_target=config_data.validation_target,
            vocab_source=vocab_source,
            vocab_target=vocab_target,
            batch_size=args.batch_size,
            fill_up=args.fill_up,
            max_seq_len_source=max_seq_len_source,
            max_seq_len_target=max_seq_len_target,
            bucketing=not args.no_bucketing,
            bucket_width=args.bucket_width)

        # learning rate scheduling
        learning_rate_half_life = none_if_negative(
            args.learning_rate_half_life)
        # TODO: The loading for continuation of the scheduler is done separately from the other parts
        if not resume_training:
            lr_scheduler_instance = lr_scheduler.get_lr_scheduler(
                args.learning_rate_scheduler_type, args.checkpoint_frequency,
                learning_rate_half_life, args.learning_rate_reduce_factor,
                args.learning_rate_reduce_num_not_improved)
        else:
            with open(os.path.join(training_state_dir, C.SCHEDULER_STATE_NAME),
                      "rb") as fp:
                lr_scheduler_instance = pickle.load(fp)

        # model configuration
        num_embed_source = args.num_embed if args.num_embed_source is None else args.num_embed_source
        num_embed_target = args.num_embed if args.num_embed_target is None else args.num_embed_target

        config_rnn = rnn.RNNConfig(cell_type=args.rnn_cell_type,
                                   num_hidden=args.rnn_num_hidden,
                                   num_layers=args.rnn_num_layers,
                                   dropout=args.dropout,
                                   residual=args.rnn_residual_connections,
                                   forget_bias=args.rnn_forget_bias)

        config_conv = None
        if args.encoder == C.RNN_WITH_CONV_EMBED_NAME:
            config_conv = encoder.ConvolutionalEmbeddingConfig(
                num_embed=num_embed_source,
                max_filter_width=args.conv_embed_max_filter_width,
                num_filters=args.conv_embed_num_filters,
                pool_stride=args.conv_embed_pool_stride,
                num_highway_layers=args.conv_embed_num_highway_layers,
                dropout=args.dropout)

        config_encoder = encoder.RecurrentEncoderConfig(
            vocab_size=vocab_source_size,
            num_embed=num_embed_source,
            rnn_config=config_rnn,
            conv_config=config_conv)

        config_decoder = decoder.RecurrentDecoderConfig(
            vocab_size=vocab_target_size,
            num_embed=num_embed_target,
            rnn_config=config_rnn,
            dropout=args.dropout,
            weight_tying=args.weight_tying,
            context_gating=args.context_gating,
            layer_normalization=args.layer_normalization)

        attention_num_hidden = args.rnn_num_hidden if not args.attention_num_hidden else args.attention_num_hidden
        config_coverage = None
        if args.attention_type == "coverage":
            config_coverage = coverage.CoverageConfig(
                type=args.attention_coverage_type,
                num_hidden=args.attention_coverage_num_hidden,
                layer_normalization=args.layer_normalization)
        config_attention = attention.AttentionConfig(
            type=args.attention_type,
            num_hidden=attention_num_hidden,
            input_previous_word=args.attention_use_prev_word,
            rnn_num_hidden=config_rnn.num_hidden,
            layer_normalization=args.layer_normalization,
            config_coverage=config_coverage)

        config_loss = loss.LossConfig(
            type=args.loss,
            vocab_size=vocab_target_size,
            normalize=args.normalize_loss,
            smoothed_cross_entropy_alpha=args.smoothed_cross_entropy_alpha)

        model_config = model.ModelConfig(
            config_data=config_data,
            max_seq_len=max_seq_len_source,
            vocab_source_size=vocab_source_size,
            vocab_target_size=vocab_target_size,
            config_encoder=config_encoder,
            config_decoder=config_decoder,
            config_attention=config_attention,
            config_loss=config_loss,
            lexical_bias=args.lexical_bias,
            learn_lexical_bias=args.learn_lexical_bias)
        model_config.freeze()

        # create training model
        training_model = training.TrainingModel(
            config=model_config,
            context=context,
            train_iter=train_iter,
            fused=args.use_fused_rnn,
            bucketing=not args.no_bucketing,
            lr_scheduler=lr_scheduler_instance)

        # We may consider loading the params in TrainingModule, for consistency
        # with the training state saving
        if resume_training:
            logger.info(
                "Found partial training in directory %s. Resuming from saved state.",
                training_state_dir)
            training_model.load_params_from_file(
                os.path.join(training_state_dir, C.TRAINING_STATE_PARAMS_NAME))
        elif args.params:
            logger.info(
                "Training will initialize from parameters loaded from '%s'",
                args.params)
            training_model.load_params_from_file(args.params)

        lexicon_array = lexicon.initialize_lexicon(
            args.lexical_bias, vocab_source,
            vocab_target) if args.lexical_bias else None

        weight_initializer = initializer.get_initializer(args.rnn_h2h_init,
                                                         lexicon=lexicon_array)

        optimizer = args.optimizer
        optimizer_params = {
            'wd': args.weight_decay,
            "learning_rate": args.initial_learning_rate
        }
        if lr_scheduler_instance is not None:
            optimizer_params["lr_scheduler"] = lr_scheduler_instance
        clip_gradient = none_if_negative(args.clip_gradient)
        if clip_gradient is not None:
            optimizer_params["clip_gradient"] = clip_gradient
        if args.momentum is not None:
            optimizer_params["momentum"] = args.momentum
        if args.normalize_loss:
            # When normalize_loss is turned on we normalize by the number of non-PAD symbols in a batch which implicitly
            # already contains the number of sentences and therefore we need to disable rescale_grad.
            optimizer_params["rescale_grad"] = 1.0
        else:
            # Making MXNet module API's default scaling factor explicit
            optimizer_params["rescale_grad"] = 1.0 / args.batch_size
        logger.info("Optimizer: %s", optimizer)
        logger.info("Optimizer Parameters: %s", optimizer_params)

        training_model.fit(
            train_iter,
            eval_iter,
            output_folder=output_folder,
            max_params_files_to_keep=args.keep_last_params,
            metrics=args.metrics,
            initializer=weight_initializer,
            max_updates=args.max_updates,
            checkpoint_frequency=args.checkpoint_frequency,
            optimizer=optimizer,
            optimizer_params=optimizer_params,
            optimized_metric=args.optimized_metric,
            max_num_not_improved=args.max_num_checkpoint_not_improved,
            min_num_epochs=args.min_num_epochs,
            monitor_bleu=args.monitor_bleu,
            use_tensorboard=args.use_tensorboard)
コード例 #6
0
ファイル: test_utils.py プロジェクト: threeleafzerg/sockeye
def test_get_num_gpus():
    assert utils.get_num_gpus() >= 0
コード例 #7
0
def main():
    params = argparse.ArgumentParser(
        description='CLI to train sockeye sequence-to-sequence models.')
    params = arguments.add_io_args(params)
    params = arguments.add_model_parameters(params)
    params = arguments.add_training_args(params)
    params = arguments.add_device_args(params)
    args = params.parse_args()

    # seed the RNGs
    np.random.seed(args.seed)
    random.seed(args.seed)
    mx.random.seed(args.seed)

    if args.use_fused_rnn:
        assert not args.use_cpu, "GPU required for FusedRNN cells"

    if args.rnn_residual_connections:
        assert args.rnn_num_layers > 2, "Residual connections require at least 3 RNN layers"

    assert args.optimized_metric == C.BLEU or args.optimized_metric in args.metrics, \
        "Must optimize either BLEU or one of tracked metrics (--metrics)"

    # Checking status of output folder, resumption, etc.
    # Create temporary logger to console only
    logger = setup_main_logger(__name__, console=not args.quiet)
    output_folder = os.path.abspath(args.output)
    resume_training = False
    training_state_dir = os.path.join(output_folder, C.TRAINING_STATE_DIRNAME)
    if os.path.exists(output_folder):
        if args.overwrite_output:
            logger.info("Removing existing output folder %s.", output_folder)
            shutil.rmtree(output_folder)
            os.makedirs(output_folder)
        elif os.path.exists(training_state_dir):
            with open(os.path.join(output_folder, C.ARGS_STATE_NAME),
                      "r") as fp:
                old_args = json.load(fp)
            arg_diffs = _dict_difference(
                vars(args), old_args) | _dict_difference(old_args, vars(args))
            # Remove args that may differ without affecting the training.
            arg_diffs -= set(C.ARGS_MAY_DIFFER)
            if not arg_diffs:
                resume_training = True
            else:
                # We do not have the logger yet
                logger.error(
                    "Mismatch in arguments for training continuation.")
                logger.error("Differing arguments: %s.", ", ".join(arg_diffs))
                sys.exit(1)
        else:
            logger.error("Refusing to overwrite existing output folder %s.",
                         output_folder)
            sys.exit(1)
    else:
        os.makedirs(output_folder)

    logger = setup_main_logger(__name__,
                               console=not args.quiet,
                               path=os.path.join(output_folder, C.LOG_NAME))

    logger.info("Command: %s", " ".join(sys.argv))
    logger.info("Arguments: %s", args)
    with open(os.path.join(output_folder, C.ARGS_STATE_NAME), "w") as fp:
        json.dump(vars(args), fp)

    with ExitStack() as exit_stack:
        # context
        if args.use_cpu:
            context = [mx.cpu()]
        else:
            num_gpus = get_num_gpus()
            assert num_gpus > 0, "No GPUs found, consider running on the CPU with --use-cpu " \
                                 "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi " \
                                 "binary isn't on the path)."
            context = []
            for gpu_id in args.device_ids:
                if gpu_id < 0:
                    # get an automatic gpu id:
                    gpu_id = exit_stack.enter_context(acquire_gpu())
                context.append(mx.gpu(gpu_id))

        # create vocabs
        vocab_source = _build_or_load_vocab(args.source_vocab, args.source,
                                            args.num_words,
                                            args.word_min_count)
        sockeye.vocab.vocab_to_json(
            vocab_source,
            os.path.join(output_folder, C.VOCAB_SRC_NAME) + C.JSON_SUFFIX)

        vocab_target = _build_or_load_vocab(args.target_vocab, args.target,
                                            args.num_words,
                                            args.word_min_count)
        sockeye.vocab.vocab_to_json(
            vocab_target,
            os.path.join(output_folder, C.VOCAB_TRG_NAME) + C.JSON_SUFFIX)

        vocab_source_size = len(vocab_source)
        vocab_target_size = len(vocab_target)
        logger.info("Vocabulary sizes: source=%d target=%d", vocab_source_size,
                    vocab_target_size)

        data_info = sockeye.data_io.DataInfo(
            os.path.abspath(args.source), os.path.abspath(args.target),
            os.path.abspath(args.validation_source),
            os.path.abspath(args.validation_target), args.source_vocab,
            args.target_vocab)

        # create data iterators
        train_iter, eval_iter = sockeye.data_io.get_training_data_iters(
            source=data_info.source,
            target=data_info.target,
            validation_source=data_info.validation_source,
            validation_target=data_info.validation_target,
            vocab_source=vocab_source,
            vocab_target=vocab_target,
            batch_size=args.batch_size,
            fill_up=args.fill_up,
            max_seq_len=args.max_seq_len,
            bucketing=not args.no_bucketing,
            bucket_width=args.bucket_width)

        # learning rate scheduling
        learning_rate_half_life = none_if_negative(
            args.learning_rate_half_life)
        # TODO: The loading for continuation of the scheduler is done separately from the other parts
        if not resume_training:
            lr_scheduler = sockeye.lr_scheduler.get_lr_scheduler(
                args.learning_rate_scheduler_type, args.checkpoint_frequency,
                learning_rate_half_life, args.learning_rate_reduce_factor,
                args.learning_rate_reduce_num_not_improved)
        else:
            with open(os.path.join(training_state_dir, C.SCHEDULER_STATE_NAME),
                      "rb") as fp:
                lr_scheduler = pickle.load(fp)

        # model configuration
        num_embed_source = args.num_embed if args.num_embed_source is None else args.num_embed_source
        num_embed_target = args.num_embed if args.num_embed_target is None else args.num_embed_target
        attention_num_hidden = args.rnn_num_hidden if not args.attention_num_hidden else args.attention_num_hidden
        model_config = sockeye.model.ModelConfig(
            max_seq_len=args.max_seq_len,
            vocab_source_size=vocab_source_size,
            vocab_target_size=vocab_target_size,
            num_embed_source=num_embed_source,
            num_embed_target=num_embed_target,
            attention_type=args.attention_type,
            attention_num_hidden=attention_num_hidden,
            attention_coverage_type=args.attention_coverage_type,
            attention_coverage_num_hidden=args.attention_coverage_num_hidden,
            attention_use_prev_word=args.attention_use_prev_word,
            dropout=args.dropout,
            rnn_cell_type=args.rnn_cell_type,
            rnn_num_layers=args.rnn_num_layers,
            rnn_num_hidden=args.rnn_num_hidden,
            rnn_residual_connections=args.rnn_residual_connections,
            weight_tying=args.weight_tying,
            context_gating=args.context_gating,
            lexical_bias=args.lexical_bias,
            learn_lexical_bias=args.learn_lexical_bias,
            data_info=data_info,
            loss=args.loss,
            normalize_loss=args.normalize_loss,
            smoothed_cross_entropy_alpha=args.smoothed_cross_entropy_alpha)

        # create training model
        model = sockeye.training.TrainingModel(
            model_config=model_config,
            context=context,
            train_iter=train_iter,
            fused=args.use_fused_rnn,
            bucketing=not args.no_bucketing,
            lr_scheduler=lr_scheduler,
            rnn_forget_bias=args.rnn_forget_bias)

        # We may consider loading the params in TrainingModule, for consistency
        # with the training state saving
        if resume_training:
            logger.info(
                "Found partial training in directory %s. Resuming from saved state.",
                training_state_dir)
            model.load_params_from_file(
                os.path.join(training_state_dir, C.TRAINING_STATE_PARAMS_NAME))
        elif args.params:
            logger.info(
                "Training will initialize from parameters loaded from '%s'",
                args.params)
            model.load_params_from_file(args.params)

        lexicon = sockeye.lexicon.initialize_lexicon(
            args.lexical_bias, vocab_source,
            vocab_target) if args.lexical_bias else None

        initializer = sockeye.initializer.get_initializer(args.rnn_h2h_init,
                                                          lexicon=lexicon)

        optimizer = args.optimizer
        optimizer_params = {
            'wd': args.weight_decay,
            "learning_rate": args.initial_learning_rate
        }
        if lr_scheduler is not None:
            optimizer_params["lr_scheduler"] = lr_scheduler
        clip_gradient = none_if_negative(args.clip_gradient)
        if clip_gradient is not None:
            optimizer_params["clip_gradient"] = clip_gradient
        if args.momentum is not None:
            optimizer_params["momentum"] = args.momentum
        logger.info("Optimizer: %s", optimizer)
        logger.info("Optimizer Parameters: %s", optimizer_params)

        model.fit(train_iter,
                  eval_iter,
                  output_folder=output_folder,
                  metrics=args.metrics,
                  initializer=initializer,
                  max_updates=args.max_updates,
                  checkpoint_frequency=args.checkpoint_frequency,
                  optimizer=optimizer,
                  optimizer_params=optimizer_params,
                  optimized_metric=args.optimized_metric,
                  max_num_not_improved=args.max_num_checkpoint_not_improved,
                  monitor_bleu=args.monitor_bleu,
                  use_tensorboard=args.use_tensorboard)
コード例 #8
0
ファイル: translate.py プロジェクト: hagenf/sockeye
def main():
    params = argparse.ArgumentParser(
        description='Translate from STDIN to STDOUT')
    params = arguments.add_inference_args(params)
    params = arguments.add_device_args(params)
    args = params.parse_args()

    logger = setup_main_logger(__name__, file_logging=False)

    assert args.beam_size > 0, "Beam size must be 1 or greater."
    if args.checkpoints is not None:
        assert len(args.checkpoints) == len(
            args.models), "must provide checkpoints for each model"

    logger.info("Command: %s", " ".join(sys.argv))
    logger.info("Arguments: %s", args)

    output_stream = sys.stdout
    output_handler = sockeye.output_handler.get_output_handler(
        args.output_type, output_stream, args.align_plot_prefix,
        args.sure_align_threshold)

    with ExitStack() as exit_stack:
        if args.use_cpu:
            context = mx.cpu()
        else:
            num_gpus = get_num_gpus()
            assert num_gpus > 0, "No GPUs found, consider running on the CPU with --use-cpu " \
                                 "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi " \
                                 "binary isn't on the path)."
            assert len(
                args.device_ids) == 1, "cannot run on multiple devices for now"
            gpu_id = args.device_ids[0]
            if gpu_id < 0:
                # get a gpu id automatically:
                gpu_id = exit_stack.enter_context(acquire_gpu())
            context = mx.gpu(gpu_id)

        translator = sockeye.inference.Translator(
            context, args.ensemble_mode,
            *sockeye.inference.load_models(context, args.max_input_len,
                                           args.beam_size, args.models,
                                           args.checkpoints,
                                           args.softmax_temperature))
        total_time = 0
        i = 0
        for i, line in enumerate(sys.stdin, 1):
            trans_input = translator.make_input(i, line)
            logger.debug(" IN: %s", trans_input)

            tic = time.time()
            trans_output = translator.translate(trans_input)
            trans_wall_time = time.time() - tic
            total_time += trans_wall_time

            logger.debug("OUT: %s", trans_output)
            logger.debug("OUT: time=%.2f", trans_wall_time)

            output_handler.handle(trans_input, trans_output)

        if i != 0:
            logger.info(
                "Processed %d lines. Total time: %.4f sec/sent: %.4f sent/sec: %.4f",
                i, total_time, total_time / i, i / total_time)
        else:
            logger.info("Processed 0 lines.")
コード例 #9
0
def main():
    params = argparse.ArgumentParser(
        description='CLI to train sockeye sequence-to-sequence models.')
    params = arguments.add_io_args(params)
    params = arguments.add_model_parameters(params)
    params = arguments.add_training_args(params)
    params = arguments.add_device_args(params)
    args = params.parse_args()

    # seed the RNGs
    np.random.seed(args.seed)
    random.seed(args.seed)
    mx.random.seed(args.seed)

    if args.use_fused_rnn:
        assert not args.use_cpu, "GPU required for FusedRNN cells"

    if args.rnn_residual_connections:
        assert args.rnn_num_layers > 2, "Residual connections require at least 3 RNN layers"

    assert args.optimized_metric == C.BLEU or args.optimized_metric in args.metrics, \
        "Must optimize either BLEU or one of tracked metrics (--metrics)"

    output_folder = os.path.abspath(args.output)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    logger = setup_main_logger(__name__,
                               console=not args.quiet,
                               path=os.path.join(output_folder, C.LOG_NAME))

    logger.info("Command: %s", " ".join(sys.argv))
    logger.info("Arguments: %s", args)

    with ExitStack() as exit_stack:
        # context
        if args.use_cpu:
            context = [mx.cpu()]
        else:
            num_gpus = get_num_gpus()
            assert num_gpus > 0, "No GPUs found, consider running on the CPU with --use-cpu " \
                                 "(note: check depends on nvidia-smi and this could also mean that the nvidia-smi " \
                                 "binary isn't on the path)."
            context = []
            for gpu_id in args.device_ids:
                if gpu_id < 0:
                    # get an automatic gpu id:
                    gpu_id = exit_stack.enter_context(acquire_gpu())
                context.append(mx.gpu(gpu_id))

        # create vocabs
        vocab_source = _build_or_load_vocab(args.source_vocab, args.source,
                                            args.num_words,
                                            args.word_min_count)
        sockeye.vocab.vocab_to_json(
            vocab_source,
            os.path.join(output_folder, C.VOCAB_SRC_NAME) + C.JSON_SUFFIX)

        vocab_target = _build_or_load_vocab(args.target_vocab, args.target,
                                            args.num_words,
                                            args.word_min_count)
        sockeye.vocab.vocab_to_json(
            vocab_target,
            os.path.join(output_folder, C.VOCAB_TRG_NAME) + C.JSON_SUFFIX)

        vocab_source_size = len(vocab_source)
        vocab_target_size = len(vocab_target)
        logger.info("Vocabulary sizes: source=%d target=%d", vocab_source_size,
                    vocab_target_size)

        data_info = sockeye.data_io.DataInfo(
            os.path.abspath(args.source), os.path.abspath(args.target),
            os.path.abspath(args.validation_source),
            os.path.abspath(args.validation_target), args.source_vocab,
            args.target_vocab)

        # create data iterators
        train_iter, eval_iter = sockeye.data_io.get_training_data_iters(
            source=data_info.source,
            target=data_info.target,
            validation_source=data_info.validation_source,
            validation_target=data_info.validation_target,
            vocab_source=vocab_source,
            vocab_target=vocab_target,
            batch_size=args.batch_size,
            fill_up=args.fill_up,
            max_seq_len=args.max_seq_len,
            bucketing=not args.no_bucketing,
            bucket_width=args.bucket_width)

        # learning rate scheduling
        learning_rate_half_life = none_if_negative(
            args.learning_rate_half_life)
        lr_scheduler = sockeye.lr_scheduler.get_lr_scheduler(
            args.learning_rate_scheduler_type, args.checkpoint_frequency,
            learning_rate_half_life, args.learning_rate_reduce_factor,
            args.learning_rate_reduce_num_not_improved)

        # model configuration
        num_embed_source = args.num_embed if args.num_embed_source is None else args.num_embed_source
        num_embed_target = args.num_embed if args.num_embed_target is None else args.num_embed_target
        attention_num_hidden = args.rnn_num_hidden if not args.attention_num_hidden else args.attention_num_hidden
        model_config = sockeye.model.ModelConfig(
            max_seq_len=args.max_seq_len,
            vocab_source_size=vocab_source_size,
            vocab_target_size=vocab_target_size,
            num_embed_source=num_embed_source,
            num_embed_target=num_embed_target,
            attention_type=args.attention_type,
            attention_num_hidden=attention_num_hidden,
            attention_coverage_type=args.attention_coverage_type,
            attention_coverage_num_hidden=args.attention_coverage_num_hidden,
            attention_use_prev_word=args.attention_use_prev_word,
            dropout=args.dropout,
            rnn_cell_type=args.rnn_cell_type,
            rnn_num_layers=args.rnn_num_layers,
            rnn_num_hidden=args.rnn_num_hidden,
            rnn_residual_connections=args.rnn_residual_connections,
            weight_tying=args.weight_tying,
            context_gating=args.context_gating,
            lexical_bias=args.lexical_bias,
            learn_lexical_bias=args.learn_lexical_bias,
            data_info=data_info,
            loss=args.loss,
            normalize_loss=args.normalize_loss,
            smoothed_cross_entropy_alpha=args.smoothed_cross_entropy_alpha)

        # create training model
        model = sockeye.training.TrainingModel(
            model_config=model_config,
            context=context,
            train_iter=train_iter,
            fused=args.use_fused_rnn,
            bucketing=not args.no_bucketing,
            lr_scheduler=lr_scheduler,
            rnn_forget_bias=args.rnn_forget_bias)

        if args.params:
            model.load_params_from_file(args.params)
            logger.info(
                "Training will continue from parameters loaded from '%s'",
                args.params)

        lexicon = sockeye.lexicon.initialize_lexicon(
            args.lexical_bias, vocab_source,
            vocab_target) if args.lexical_bias else None

        initializer = sockeye.initializer.get_initializer(args.rnn_h2h_init,
                                                          lexicon=lexicon)

        optimizer = args.optimizer
        optimizer_params = {
            'wd': args.weight_decay,
            "learning_rate": args.initial_learning_rate
        }
        if lr_scheduler is not None:
            optimizer_params["lr_scheduler"] = lr_scheduler
        clip_gradient = none_if_negative(args.clip_gradient)
        if clip_gradient is not None:
            optimizer_params["clip_gradient"] = clip_gradient
        if args.momentum is not None:
            optimizer_params["momentum"] = args.momentum
        logger.info("Optimizer: %s", optimizer)
        logger.info("Optimizer Parameters: %s", optimizer_params)

        model.fit(train_iter,
                  eval_iter,
                  output_folder=output_folder,
                  metrics=args.metrics,
                  initializer=initializer,
                  max_updates=args.max_updates,
                  checkpoint_frequency=args.checkpoint_frequency,
                  optimizer=optimizer,
                  optimizer_params=optimizer_params,
                  optimized_metric=args.optimized_metric,
                  max_num_not_improved=args.max_num_checkpoint_not_improved,
                  monitor_bleu=args.monitor_bleu,
                  use_tensorboard=args.use_tensorboard)