train_lambda.py

from encoder import *
from decoder import *
from time import clock
from image_util import *
from msssim import MultiScaleSSIM
import numpy as np
from quantization_op import *
from normalization_op import batch_normalization, batch_denormalization
from nlp import *
from entropy_estimation import *
from tensorflow.python.client import timeline
from image_input_pipeline import *

batch_size = 32

def train(input_images, hps=None, learning_rate=None):
  with tf.name_scope('train'):
    # norm_images, means, stddevs = batch_normalization(input_images)
    with tf.device('/gpu:1'):
      _lambda = tf.get_variable('lambda', shape=[], dtype=tf.float32,
                                initializer=tf.constant_initializer(1.0))
      p_lambda = tf.Print(_lambda, [_lambda])
      features = build_encoder(input_images=input_images, hps=hps)
      lambda_features = p_lambda * features
      quant_code = my_quant(lambda_features)
    # print(quant_code.get_shape().as_list())
    with tf.device('/gpu:2'):
      code = quant_code / p_lambda
      recon_images = build_decoder(code=code, hps=hps)
    # recon_images = batch_denormalization(recon_norm, means, stddevs)

    global_step = tf.contrib.framework.get_or_create_global_step()

    with tf.device('/gpu:3'):
     entropy_rate = tf.reduce_sum(entropy_estimate(lambda_features, p_lambda))
     im_shape = input_images.get_shape().as_list()
     tf.summary.scalar('entropy_rate', entropy_rate)
     tf.summary.scalar('bpp', entropy_rate / (batch_size * im_shape[1] * im_shape[2] * im_shape[3]))

    mse_loss = tf.reduce_mean(tf.square(recon_images - input_images))
    nlp_loss_weight = nlp_loss_weight_rgb(input_images, recon_images)
    nlp_loss_weight_mse = tf.reduce_mean(nlp_loss_weight * tf.square(recon_images - input_images))
    tf.summary.scalar('mse_loss', nlp_loss_weight_mse)
    # 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)

    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)

    if learning_rate is not None:
      # train_var_list = [_lambda] + tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'bit_rate_estimate')
      train_var_list = [_lambda]
      gradients = tf.gradients(1e-6 * entropy_rate + hps['beta'] * nlp_loss_weight_mse, train_var_list)
      # gradients = tf.gradients(nlp_loss_weight_mse, var_list)
      clipped_gradients, norm = tf.clip_by_global_norm(gradients, 5)
      for i in range(len(train_var_list)):
      #  print(var_list[i].name)
        if gradients[i] != None:
          tf.summary.histogram(train_var_list[i].name, train_var_list[i])

      train_method = tf.train.AdamOptimizer(learning_rate)
      train_step = train_method.apply_gradients(zip(clipped_gradients, train_var_list),
                                                global_step=global_step)
      return recon_images, psnr, train_step, global_step

  return recon_images, psnr, quant_code

def build_model(_beta=None):
  with tf.Graph().as_default():
    dir = os.path.join('..', 'data', 'Raise_6k')
    train_image_batch, validation_image_batch = input_pipeline(dir, 3, batch_size)
    select_data = tf.placeholder(dtype=bool, shape=[], name='select_data')
    image_batch = tf.cond(
      select_data,
      lambda: train_image_batch,
      lambda: validation_image_batch
    )
    learning_rate = tf.placeholder(tf.float32, shape=[])
    recon_images, psnr, train_step, global_step = train(image_batch, HParams, learning_rate)

    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])

    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = 0.9
    # config.gpu_options.allow_growth = True

    # sess = tf.Session(config=tf.ConfigProto(log_device_placement=True, allow_soft_placement=True))
    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
    init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    sess.run(init)

    summary_writer_train = tf.summary.FileWriter('log/log_lambda_beta_%d/train/'%HParams['beta'], sess.graph)
    summary_writer_test = tf.summary.FileWriter('log/log_lambda_beta_%d/test/'%HParams['beta'], sess.graph)

    var_list = tf.trainable_variables()
    #for i in range(len(var_list)):
    #  print(var_list[i].name)
    save_saver = tf.train.Saver(var_list + [global_step])
    if not _beta:
      var_list = [v for v in var_list if v.name != 'lambda:0']
      restore_saver = tf.train.Saver(var_list)
    else:
      restore_saver = tf.train.Saver(var_list + [global_step])

    if not os.path.exists("./checkPoint/model_lambda_beta_%d/"%HParams['beta']):
      os.makedirs("./checkPoint/model_lambda_beta_%d/"%HParams['beta'])
    if not os.path.exists("./checkPoint/model_no_lambda/"):
      print("No pre-trained models for fine-tuning.")
      exit()
    if _beta:
      module_file = tf.train.latest_checkpoint("./checkPoint/model_lambda_beta_%d/"%_beta)
    else:
      module_file = tf.train.latest_checkpoint("./checkPoint/model_no_lambda/")
    #saver.restore(sess, './checkPoint/model-385450')
    if module_file != None:
      restore_saver.restore(sess, module_file)
    else:
      print("No pre-trained models for fine-tuning.")
      exit()
    #saver.restore(sess, './checkPoint/')

    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    # saver = tf.train.Saver()
    actual_step = 0
    tic = clock()
    # steps_per_epoch = 200
    # for i in range(steps_per_epoch * 4000):
    for i in range(10000):
      time_0 = clock()
      kappa = 0.8
      tau = 1000
      original_lr = 1e-3
      HParams['learning_rate'] = original_lr * (tau ** kappa) / ((tau + actual_step) ** kappa)
      print(actual_step, ', lr: ', HParams['learning_rate'])
      if i == 0 or actual_step % 100 == 0:
        run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
        run_metadata = tf.RunMetadata()
        summary, _, recon_images_np, psnr_val, actual_step, input_images = sess.run(
          [merged, train_step, recon_images, psnr, global_step, image_batch],
          feed_dict={select_data: True, learning_rate: HParams['learning_rate']},
          options=run_options, run_metadata=run_metadata)
        tl = timeline.Timeline(step_stats=run_metadata.step_stats)
        ctf = tl.generate_chrome_trace_format()
        with open('timeline_lambda.json', 'w') as f:
          f.write(ctf)

        # summary_writer_train.add_run_metadata(run_metadata, 'step %d' % actual_step)
        summary_writer_train.add_summary(summary, actual_step)
      else:
        _, recon_images_np, psnr_val, actual_step, input_images = sess.run(
          [train_step, recon_images, psnr, global_step, image_batch],
          feed_dict={select_data: True, learning_rate: HParams['learning_rate']})
      time_1 = clock()

      # print('Time for training: ', time_1 - time_0)

      if actual_step % 100 == 0 or i == 0:
        toc = clock()
        print(actual_step,
              ' steps, fps %2f, psnr %f' % ((toc - tic) / 100, psnr_val))
        tic = clock()
        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 i == 0:
        save_saver.save(sess, "./checkPoint/model_lambda_beta_%d/model"%HParams['beta'], global_step=actual_step)

      if actual_step % 100 == 0 or i == 0:
        summary, recon_images_np, psnr_val, input_images = sess.run(
          [merged, recon_images, psnr, image_batch], feed_dict={select_data: False})
        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_test.add_summary(summary, actual_step)
        summary_writer_test.add_summary(summary_ms_ssim_val, actual_step)
    sess.close()


def main(_):
  build_model()

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