Python model.model_fn() Examples

The following are 2 code examples of model.model_fn(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module model , or try the search function .
Example #1
Source File: experiment.py    From MemTrack with MIT License 6 votes vote down vote up 
def experiment():

    train_input_fn = generate_input_fn(
        is_train=True,
        tfrecords_path=config.tfrecords_path,
        batch_size=config.batch_size,
        time_step=config.time_step)

    eval_input_fn = generate_input_fn(
        is_train=False,
        tfrecords_path=config.tfrecords_path,
        batch_size=config.batch_size_eval,
        time_step=config.time_step_eval)

    estimator = Estimator(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        model_fn=model_fn)

    estimator.train()
Example #2
Source File: create_pb.py    From light-head-rcnn with MIT License 5 votes vote down vote up 
def export_savedmodel():
    config = tf.ConfigProto()
    config.gpu_options.visible_device_list = GPU_TO_USE
    run_config = tf.estimator.RunConfig()
    run_config = run_config.replace(
        model_dir=params['model_dir'],
        session_config=config
    )
    params['nms_max_output_size'] = NMS_MAX_OUTPUT_SIZE
    estimator = tf.estimator.Estimator(model_fn, params=params, config=run_config)

    def serving_input_receiver_fn():
        raw_images = tf.placeholder(dtype=tf.uint8, shape=[BATCH_SIZE, None, None, 3], name='images')
        w, h = tf.shape(raw_images)[2], tf.shape(raw_images)[1]

        with tf.device('/gpu:0'):

            images = tf.to_float(raw_images)
            if RESIZE:
                images = tf.squeeze(images, 0)
                images = resize_keeping_aspect_ratio(images, MIN_DIMENSION, MAX_DIMENSION)
                images = tf.expand_dims(images, 0)

            features = {
                'images': (1.0/255.0) * images,
                'images_size': tf.stack([w, h])
            }
        return tf.estimator.export.ServingInputReceiver(features, {'images': raw_images})

    shutil.rmtree(OUTPUT_FOLDER, ignore_errors=True)
    os.mkdir(OUTPUT_FOLDER)
    estimator.export_savedmodel(OUTPUT_FOLDER, serving_input_receiver_fn)