Python validate_batch_size_for_multi_gpu示例

示例#1

def main(_):
    os.environ['TF_ENABLE_WINOGRAD_NONFUSED'] = '1'

    if FLAGS.multi_gpu:
        validate_batch_size_for_multi_gpu(FLAGS.batch_size)
        model_function = tf.contrib.estimator.replicate_model_fn(
            model_fn, loss_reduction=tf.losses.Reduction.MEAN)
    else:
        model_function = model_fn

    sess_config = tf.ConfigProto(
        allow_soft_placement=True,
        intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads,
        inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
        gpu_options=tf.GPUOptions(allow_growth=True))

    run_config = tf.estimator.RunConfig(
        session_config=sess_config,
        save_checkpoints_steps=FLAGS.save_checkpoint_steps,
        save_summary_steps=FLAGS.save_summary_steps,
        keep_checkpoint_max=100)

    train_input_fn = functools.partial(input_fn, batch_size=FLAGS.batch_size)

    estimator = tf.estimator.Estimator(model_fn=model_function,
                                       model_dir=FLAGS.job_dir,
                                       config=run_config,
                                       params=FLAGS)

    estimator.train(input_fn=train_input_fn, max_steps=FLAGS.max_steps)

示例#2

def main(_):
  os.environ['TF_ENABLE_WINOGRAD_NONFUSED'] = '1'

  session_config = tf.ConfigProto(
    allow_soft_placement=True,
    intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads,
    inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
    gpu_options=tf.GPUOptions(allow_growth=True))

  run_config = tf.estimator.RunConfig(
    session_config=session_config,
    save_checkpoints_steps=FLAGS.save_checkpoint_steps,
    save_summary_steps=FLAGS.save_summary_steps,
    keep_checkpoint_max=100)

  configs = get_configs_from_pipeline_file(
    'data/faster_rcnn_resnet101_pets.config')
  model_config = configs['model']
  train_config = configs['train_config']
  train_input_config = configs['train_input_config']
  eval_config = configs['eval_config']
  eval_input_config = configs['eval_input_config']

  train_input_fn = create_train_input_fn(train_config, train_input_config,
                                         model_config, 'train')

  eval_input_fn = create_train_input_fn(eval_config, eval_input_config,
                                        model_config, 'val')

  detection_model_fn = functools.partial(
    model_builder.build, model_config=model_config)

  model_fn = create_model_fn(detection_model_fn, configs, FLAGS, False)

  if FLAGS.multi_gpu:
    validate_batch_size_for_multi_gpu(FLAGS.batch_size)
    model_fn = tf.contrib.estimator.replicate_model_fn(
      model_fn,
      loss_reduction=tf.losses.Reduction.MEAN)

  estimator = tf.estimator.Estimator(
    model_fn=model_fn,
    model_dir=FLAGS.job_dir,
    config=run_config,
    params=FLAGS)

  train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn,
                                      max_steps=FLAGS.max_steps)
  eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn, steps=200)

  tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)