models/research/ptn/pretrain_rotator.py

# Copyright 2017 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Contains training plan for the Rotator model (Pretraining in NIPS16)."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import numpy as np
from six.moves import xrange
import tensorflow as tf

from tensorflow import app

import model_rotator as model

flags = tf.app.flags
slim = tf.contrib.slim

flags.DEFINE_string('inp_dir', '',
                    'Directory path containing the input data (tfrecords).')
flags.DEFINE_string(
    'dataset_name', 'shapenet_chair',
    'Dataset name that is to be used for training and evaluation.')
flags.DEFINE_integer('z_dim', 512, '')
flags.DEFINE_integer('a_dim', 3, '')
flags.DEFINE_integer('f_dim', 64, '')
flags.DEFINE_integer('fc_dim', 1024, '')
flags.DEFINE_integer('num_views', 24, 'Num of viewpoints in the input data.')
flags.DEFINE_integer('image_size', 64,
                     'Input images dimension (pixels) - width & height.')
flags.DEFINE_integer('step_size', 1, 'Steps to take for rotation in pretraining.')
flags.DEFINE_integer('batch_size', 32, 'Batch size for training.')
flags.DEFINE_string('encoder_name', 'ptn_encoder',
                    'Name of the encoder network being used.')
flags.DEFINE_string('decoder_name', 'ptn_im_decoder',
                    'Name of the decoder network being used.')
flags.DEFINE_string('rotator_name', 'ptn_rotator',
                    'Name of the rotator network being used.')
# Save options
flags.DEFINE_string('checkpoint_dir', '/tmp/ptn_train/',
                    'Directory path for saving trained models and other data.')
flags.DEFINE_string('model_name', 'deeprotator_pretrain',
                    'Name of the model used in naming the TF job. Must be different for each run.')
flags.DEFINE_string('init_model', None,
                    'Checkpoint path of the model to initialize with.')
flags.DEFINE_integer('save_every', 1000,
                     'Average period of steps after which we save a model.')
# Optimization
flags.DEFINE_float('image_weight', 10, 'Weighting factor for image loss.')
flags.DEFINE_float('mask_weight', 1, 'Weighting factor for mask loss.')
flags.DEFINE_float('learning_rate', 0.0001, 'Learning rate.')
flags.DEFINE_float('weight_decay', 0.001, 'Weight decay parameter while training.')
flags.DEFINE_float('clip_gradient_norm', 0, 'Gradient clim norm, leave 0 if no gradient clipping.')
flags.DEFINE_integer('max_number_of_steps', 320000, 'Maximum number of steps for training.')
# Summary
flags.DEFINE_integer('save_summaries_secs', 15, 'Seconds interval for dumping TF summaries.')
flags.DEFINE_integer('save_interval_secs', 60 * 5, 'Seconds interval to save models.')
# Distribution
flags.DEFINE_string('master', '', 'The address of the tensorflow master if running distributed.')
flags.DEFINE_bool('sync_replicas', False, 'Whether to sync gradients between replicas for optimizer.')
flags.DEFINE_integer('worker_replicas', 1, 'Number of worker replicas (train tasks).')
flags.DEFINE_integer('backup_workers', 0, 'Number of backup workers.')
flags.DEFINE_integer('ps_tasks', 0, 'Number of ps tasks.')
flags.DEFINE_integer('task', 0,
                     'Task identifier flag to be set for each task running in distributed manner. Task number 0 '
                     'will be chosen as the chief.')

FLAGS = flags.FLAGS


def main(_):
  train_dir = os.path.join(FLAGS.checkpoint_dir, FLAGS.model_name, 'train')
  save_image_dir = os.path.join(train_dir, 'images')
  if not os.path.exists(train_dir):
    os.makedirs(train_dir)
  if not os.path.exists(save_image_dir):
    os.makedirs(save_image_dir)

  g = tf.Graph()
  with g.as_default():
    with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
      global_step = slim.get_or_create_global_step()
      ##########
      ## data ##
      ##########
      train_data = model.get_inputs(
          FLAGS.inp_dir,
          FLAGS.dataset_name,
          'train',
          FLAGS.batch_size,
          FLAGS.image_size,
          is_training=True)
      inputs = model.preprocess(train_data, FLAGS.step_size)
      ###########
      ## model ##
      ###########
      model_fn = model.get_model_fn(FLAGS, is_training=True)
      outputs = model_fn(inputs)
      ##########
      ## loss ##
      ##########
      task_loss = model.get_loss(inputs, outputs, FLAGS)
      regularization_loss = model.get_regularization_loss(
          ['encoder', 'rotator', 'decoder'], FLAGS)
      loss = task_loss + regularization_loss
      ###############
      ## optimizer ##
      ###############
      optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
      if FLAGS.sync_replicas:
        optimizer = tf.train.SyncReplicasOptimizer(
            optimizer,
            replicas_to_aggregate=FLAGS.workers_replicas - FLAGS.backup_workers,
            total_num_replicas=FLAGS.worker_replicas)

      ##############
      ## train_op ##
      ##############
      train_op = model.get_train_op_for_scope(
          loss, optimizer, ['encoder', 'rotator', 'decoder'], FLAGS)
      ###########
      ## saver ##
      ###########
      saver = tf.train.Saver(max_to_keep=np.minimum(5,
                                                    FLAGS.worker_replicas + 1))

      if FLAGS.task == 0:
        val_data = model.get_inputs(
            FLAGS.inp_dir,
            FLAGS.dataset_name,
            'val',
            FLAGS.batch_size,
            FLAGS.image_size,
            is_training=False)
        val_inputs = model.preprocess(val_data, FLAGS.step_size)
        # Note: don't compute loss here
        reused_model_fn = model.get_model_fn(
            FLAGS, is_training=False, reuse=True)
        val_outputs = reused_model_fn(val_inputs)
        with tf.device(tf.DeviceSpec(device_type='CPU')):
          if FLAGS.step_size == 1:
            vis_input_images = val_inputs['images_0'] * 255.0
            vis_output_images = val_inputs['images_1'] * 255.0
            vis_pred_images = val_outputs['images_1'] * 255.0
            vis_pred_masks = (val_outputs['masks_1'] * (-1) + 1) * 255.0
          else:
            rep_times = int(np.ceil(32.0 / float(FLAGS.step_size)))
            vis_list_1 = []
            vis_list_2 = []
            vis_list_3 = []
            vis_list_4 = []
            for j in xrange(rep_times):
              for k in xrange(FLAGS.step_size):
                vis_input_image = val_inputs['images_0'][j],
                vis_output_image = val_inputs['images_%d' % (k + 1)][j]
                vis_pred_image = val_outputs['images_%d' % (k + 1)][j]
                vis_pred_mask = val_outputs['masks_%d' % (k + 1)][j]
                vis_list_1.append(tf.expand_dims(vis_input_image, 0))
                vis_list_2.append(tf.expand_dims(vis_output_image, 0))
                vis_list_3.append(tf.expand_dims(vis_pred_image, 0))
                vis_list_4.append(tf.expand_dims(vis_pred_mask, 0))

            vis_list_1 = tf.reshape(
                tf.stack(vis_list_1), [
                    rep_times * FLAGS.step_size, FLAGS.image_size,
                    FLAGS.image_size, 3
                ])
            vis_list_2 = tf.reshape(
                tf.stack(vis_list_2), [
                    rep_times * FLAGS.step_size, FLAGS.image_size,
                    FLAGS.image_size, 3
                ])
            vis_list_3 = tf.reshape(
                tf.stack(vis_list_3), [
                    rep_times * FLAGS.step_size, FLAGS.image_size,
                    FLAGS.image_size, 3
                ])
            vis_list_4 = tf.reshape(
                tf.stack(vis_list_4), [
                    rep_times * FLAGS.step_size, FLAGS.image_size,
                    FLAGS.image_size, 1
                ])

            vis_input_images = vis_list_1 * 255.0
            vis_output_images = vis_list_2 * 255.0
            vis_pred_images = vis_list_3 * 255.0
            vis_pred_masks = (vis_list_4 * (-1) + 1) * 255.0

          write_disk_op = model.write_disk_grid(
              global_step=global_step,
              summary_freq=FLAGS.save_every,
              log_dir=save_image_dir,
              input_images=vis_input_images,
              output_images=vis_output_images,
              pred_images=vis_pred_images,
              pred_masks=vis_pred_masks)
        with tf.control_dependencies([write_disk_op]):
          train_op = tf.identity(train_op)

      #############
      ## init_fn ##
      #############
      init_fn = model.get_init_fn(['encoder, ' 'rotator', 'decoder'], FLAGS)

      ##############
      ## training ##
      ##############
      slim.learning.train(
          train_op=train_op,
          logdir=train_dir,
          init_fn=init_fn,
          master=FLAGS.master,
          is_chief=(FLAGS.task == 0),
          number_of_steps=FLAGS.max_number_of_steps,
          saver=saver,
          save_summaries_secs=FLAGS.save_summaries_secs,
          save_interval_secs=FLAGS.save_interval_secs)


if __name__ == '__main__':
  app.run()