train_distributed.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from datetime import datetime
import os.path
import re
import time

import numpy as np
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

from train import *

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string('train_dir', '.\\distributed',
                           """Directory where to write event logs """
                           """and checkpoint.""")
tf.app.flags.DEFINE_integer('max_steps', 1000000,
                            """Number of batches to run.""")
tf.app.flags.DEFINE_integer('num_gpus', 4,
                            """How many GPUs to use.""")
tf.app.flags.DEFINE_integer('batch_size', 32,
                            """Batch size for mini-batch method.""")
tf.app.flags.DEFINE_boolean('log_device_placement', False,
                            """Whether to log device placement.""")
tf.app.flags.DEFINE_boolean('load_with_distribution', True,
                            """Whether to load the model from models trained with distribution.""")

batch_size = FLAGS.batch_size

def tower_loss(scope, input_images, loss='mse', hps=HParams):
  """Calculate the total loss on a single tower
  Args:
    scope: unique prefix string identifying the tower, e.g. 'tower_0'
    input_images: a mini-batch of images as input tensor
  Returns:
     Tensor of shape [] containing the total loss for a batch of data
  """
  features = build_encoder(input_images=input_images, hps=hps)
  quant_code = my_quant(features)
  recon_images = build_decoder(code=quant_code, hps=hps)

  entropy_rate = entropy_estimate(quant_code)
  im_shape = input_images.get_shape().as_list()
  tf.summary.scalar('entropy_rate', entropy_rate)
  bpp = entropy_rate / (batch_size * im_shape[1] * im_shape[2])
  tf.summary.scalar('bpp', bpp)

  clipped_recon_images = recon_images
  clipped_recon_images = tf.minimum(tf.maximum(clipped_recon_images, 0.0), 1.0)
  tf.summary.image('clipped_recon_image', clipped_recon_images, 1)

  mse_loss = tf.reduce_mean(tf.square(clipped_recon_images - input_images))
  if loss == 'nlp':
    nlp_loss_weight = nlp_loss_weight_rgb(input_images, clipped_recon_images)
    nlp_loss_weight_mse = tf.reduce_mean(nlp_loss_weight)
    tf.summary.scalar('nlp_loss', nlp_loss_weight_mse)
  else:
    tf.summary.scalar('mse_loss', mse_loss)

  psnr = tf.cond(tf.equal(mse_loss, 0), lambda: tf.constant(100, dtype=tf.float32),
               lambda: 10 * (tf.log(1 * 1 / mse_loss) / np.log(10)))
  tf.summary.scalar('psnr', psnr)

  tradeoff_loss = 2e-6 * entropy_rate + hps['beta'] * (nlp_loss_weight_mse if loss == 'nlp' else mse_loss)
  return tradeoff_loss, psnr, clipped_recon_images

def average_gradients(tower_grads):
  """Calculate the average gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
  """
  average_grads = []
  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads = []
    for g, _ in grad_and_vars:
      # Add 0 dimension to the gradients to represent the tower.
      expanded_g = tf.expand_dims(g, 0)

      # Append on a 'tower' dimension which we will average over below.
      grads.append(expanded_g)

    # Average over the 'tower' dimension.
    grad = tf.concat(grads, 0)
    grad = tf.reduce_mean(grad, 0)

    # Keep in mind that the Variables are redundant because they are shared
    # across towers. So .. we will just return the first tower's pointer to
    # the Variable.
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
  return average_grads

def build_model():
  """Train CIFAR-10 for a number of steps."""
  with tf.Graph().as_default(), tf.device('/cpu:0'):
    # Create a variable to count the number of train() calls. This equals the
    # number of batches processed * FLAGS.num_gpus.
    global_step = tf.contrib.framework.get_or_create_global_step()
    # global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)

    dir = os.path.join('..', 'data', 'Raise_6k')
    train_image_batch, validation_image_batch = input_pipeline(dir, 3, batch_size * FLAGS.num_gpus)
    select_data = tf.placeholder(dtype=bool, shape=[], name='select_data')
    image_batch = tf.cond(
      select_data,
      lambda: train_image_batch,
      lambda: validation_image_batch
    )
    towers_image_batch = tf.split(image_batch, num_or_size_splits=FLAGS.num_gpus, axis=0)

    opt = tf.train.AdamOptimizer(HParams['learning_rate'])
    towers_grads = []
    towers_losses = []
    towers_psnr = []
    towers_recon_batch = []
    with tf.variable_scope('build_towers'):
      for i in xrange(FLAGS.num_gpus):
        with tf.device('/gpu:%d' % i):
          with tf.name_scope('%s_%d' % ('tower_', i)) as scope:
            loss, psnr, recon_batch = tower_loss(scope, towers_image_batch[i])
            summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)

            tf.get_variable_scope().reuse_variables()

            grads = opt.compute_gradients(loss)
            vars = [v for _, v in grads]
            grads = [g for g, _ in grads]
            clipped_grads, norm = tf.clip_by_global_norm(grads, 5)
            grads = [(g, v) for g, v in zip(clipped_grads, vars)]

            towers_grads.append(grads)
            towers_losses.append(loss)
            towers_psnr.append(psnr)
            towers_recon_batch.append(recon_batch)

    grads = average_gradients(towers_grads)
    loss = tf.reduce_mean(tf.stack(towers_losses))
    psnr = tf.reduce_mean(tf.stack(towers_psnr))
    recon_batch = tf.concat(towers_recon_batch, axis=0)
    for grad, var in grads:
      if grad is not None:
        summaries.append(tf.summary.histogram(var.op.name + '/gradients', grad))
    for var in tf.trainable_variables():
      summaries.append(tf.summary.histogram(var.op.name, var))

    train_step = opt.apply_gradients(grads, global_step=global_step)

    summary_op = tf.summary.merge(summaries)
    merged = tf.summary.merge_all()

    ms_ssim = tf.placeholder(shape=[], dtype=tf.float32, name='ms_ssim')
    ms_ssim_db = -10 * tf.log(1 - ms_ssim) / np.log(10)
    summary_ms_ssim = tf.summary.scalar('ms_ssim', ms_ssim)
    summary_ms_ssim_db = tf.summary.scalar('ms_ssim_db', ms_ssim_db)
    merged_ms_ssim = tf.summary.merge([summary_ms_ssim, summary_ms_ssim_db])

    init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    # allow_soft_placement must be set to True to build towers on GPU
    # if some of the ops do not have GPU implementations.
    sess = tf.Session(config=tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=FLAGS.log_device_placement))
    sess.run(init)

    if FLAGS.load_with_distribution:
      restore_saver = saving_saver = tf.train.Saver()
      module_dir = os.path.join(FLAGS.train_dir, 'checkPoint', 'model_no_lambda_beta_%d_%d'
                                % (HParams['beta'], HParams['codec_dim']))
    else:
      train_vars = tf.trainable_variables()
      train_vars_names = [v.name for v in train_vars]
      restore_dict = {n[n.find('/')+1:n.find(':')]: v for n, v in zip(train_vars_names, train_vars)}
      #restore_dict['global_step'] = global_step
      restore_saver = tf.train.Saver(restore_dict)
      module_dir = os.path.join('checkPoint', 'model_no_lambda_beta_%d_%d' % (HParams['beta'], HParams['codec_dim']))
      saving_saver = tf.train.Saver()
    if not os.path.exists(module_dir):
      os.makedirs(module_dir)
    module_file = tf.train.latest_checkpoint(module_dir)
    if module_file != None:
      restore_saver.restore(sess, module_file)

    # Start the queue runners.
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    summary_writer_train = tf.summary.FileWriter(
      os.path.join(FLAGS.train_dir, 'log', 'log_no_lambda_beta_%d_%d' % (HParams['beta'], HParams['codec_dim']),
                   'train'), sess.graph)
    summary_writer_test = tf.summary.FileWriter(
      os.path.join(FLAGS.train_dir, 'log', 'log_no_lambda_beta_%d_%d' % (HParams['beta'], HParams['codec_dim']),
                   'test'), sess.graph)

    start_time = time.time()
    for step in xrange(FLAGS.max_steps):
      print(step)

      summary, _, psnr_val, actual_step, recon_batch_np, image_batch_np = sess.run(
        [merged, train_step, psnr, global_step, recon_batch, image_batch],
        feed_dict={select_data: True})

      if actual_step % 100 == 0 or step == 0:
        duration = time.time() - start_time
        print(actual_step, ' steps, fps %2f, psnr %f' % (duration /(100 * FLAGS.num_gpus), psnr_val))
        start_time = time.time()
        summary_writer_train.add_summary(summary, actual_step)
        #ms_ssim_np = MultiScaleSSIM(recon_images_np, input_images, max_val=1.0)
        #summary_ms_ssim_val = sess.run(merged_ms_ssim, feed_dict={ms_ssim: ms_ssim_np})
        #summary_writer_train.add_summary(summary_ms_ssim_val, actual_step)
      if actual_step % 2000 == 0 or step == 0:
        checkpoint_path = os.path.join(FLAGS.train_dir, 'checkPoint', 'model_no_lambda_beta_%d_%d/model'%(HParams['beta'], HParams['codec_dim']))
        saving_saver.save(sess, checkpoint_path, global_step=actual_step)

def main(_):
  build_model()

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