import tensorflow as tf
from tensorflow.python.ops import gen_logging_ops
from tensorflow.python.framework import ops as _ops
import time
import shutil
import datetime
import os
import cv2
import pickle
import numpy as np
from tensorflow.contrib import slim
import sys
sys.path.append(os.getcwd())
from nets import model as model
from utils.tf_records import read_tfrecord_and_decode_into_image_annotation_pair_tensors
from utils.pascal_voc import pascal_segmentation_lut
# Parameter Setting: batch_size 2 gpu_list 2
tf.app.flags.DEFINE_string('model_type', 'refinenet', 'refinenet or sesnet')
tf.app.flags.DEFINE_integer('batch_size', 2, '')
tf.app.flags.DEFINE_integer('train_size', 512, '')
tf.app.flags.DEFINE_float('learning_rate', 0.00001, '')
tf.app.flags.DEFINE_integer('max_steps', 60000, '')
tf.app.flags.DEFINE_float('moving_average_decay', 0.997, '')
tf.app.flags.DEFINE_integer('num_classes', 2, '')
tf.app.flags.DEFINE_string('gpu_list', '0,1', '')
tf.app.flags.DEFINE_string('checkpoint_path', 'checkpoints/', '')
tf.app.flags.DEFINE_string('logs_path', 'logs/', '')
tf.app.flags.DEFINE_boolean('restore', False, 'whether to restore from checkpoint')
tf.app.flags.DEFINE_integer('save_checkpoint_steps', 500, '')
tf.app.flags.DEFINE_integer('save_summary_steps', 10, '')
tf.app.flags.DEFINE_integer('save_image_steps', 10, '')
tf.app.flags.DEFINE_string('training_data_path', 'data/train.tfrecords', '')
tf.app.flags.DEFINE_string('pretrained_model_path', 'data/resnet_v1_101.ckpt', '')
tf.app.flags.DEFINE_integer('decay_steps', 20000, '')
tf.app.flags.DEFINE_float('decay_rate', 0.1, '')
FLAGS = tf.app.flags.FLAGS
def tower_loss(images, annotation, class_labels, reuse_variables=None):
with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
logits = model.model(FLAGS.model_type, images, is_training=True)
pred = tf.argmax(logits, dimension=3)
model_loss = model.loss(annotation, logits, class_labels)
total_loss = tf.add_n([model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
# add summary
if reuse_variables is None:
tf.summary.scalar('model_loss', model_loss)
tf.summary.scalar('total_loss', total_loss)
return total_loss, model_loss, pred, logits
def average_gradients(tower_grads):
average_grads = []
for grad_and_vars in zip(*tower_grads):
grads = []
for g, _ in grad_and_vars:
expanded_g = tf.expand_dims(g, 0)
grads.append(expanded_g)
grad = tf.concat(grads, 0)
grad = tf.reduce_mean(grad, 0)
v = grad_and_vars[0][1]
grad_and_var = (grad, v)
average_grads.append(grad_and_var)
return average_grads
def build_image_summary():
log_image_data = tf.placeholder(tf.uint8, [None, None, 3])
log_image_name = tf.placeholder(tf.string)
log_image = gen_logging_ops .image_summary(log_image_name, tf.expand_dims(log_image_data, 0), max_images=1)
_ops.add_to_collection(_ops.GraphKeys.SUMMARIES, log_image)
return log_image, log_image_data, log_image_name
def main(argv=None):
gpus = range(len(FLAGS.gpu_list.split(',')))
pascal_voc_lut = pascal_segmentation_lut()
class_labels = pascal_voc_lut.keys()
print(class_labels)
with open('data/color_map', 'rb') as f:
color_map = pickle.load(f)
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
now = datetime.datetime.now()
StyleTime = now.strftime("%Y-%m-%d-%H-%M-%S")
os.makedirs(FLAGS.logs_path + StyleTime)
if not os.path.exists(FLAGS.checkpoint_path):
os.makedirs(FLAGS.checkpoint_path)
else:
if not FLAGS.restore:
if os.path.exists(FLAGS.checkpoint_path):
shutil.rmtree(FLAGS.checkpoint_path)
os.makedirs(FLAGS.checkpoint_path)
filename_queue = tf.train.string_input_producer([FLAGS.training_data_path], num_epochs=1000)
image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(filename_queue)
image_train_size = [FLAGS.train_size, FLAGS.train_size]
annotation_train_size = [FLAGS.train_size // 4, FLAGS.train_size // 4]
resized_image = tf.image.resize_images(image, image_train_size, method=1)
resized_annotation = tf.image.resize_images(annotation, annotation_train_size, method=1)
resized_annotation = tf.squeeze(resized_annotation)
image_batch, annotation_batch = tf.train.shuffle_batch([resized_image, resized_annotation],
batch_size=FLAGS.batch_size * len(gpus), capacity=1000,
num_threads=4,
min_after_dequeue=500)
# split
input_images_split = tf.split(image_batch, len(gpus))
input_segs_split = tf.split(annotation_batch, len(gpus))
learning_rate = tf.Variable(FLAGS.learning_rate, trainable=False)
global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)
# add summary
tf.summary.scalar('learning_rate', learning_rate)
opt = tf.train.AdamOptimizer(learning_rate)
tower_grads = []
reuse_variables = None
for i, gpu_id in enumerate(gpus):
with tf.device('/gpu:%d' % gpu_id):
with tf.name_scope('model_%d' % gpu_id) as scope:
iis = input_images_split[i]
isms = input_segs_split[i]
total_loss, model_loss, output_pred, output_logits = tower_loss(iis, isms, class_labels, reuse_variables)
batch_norm_updates_op = tf.group(*tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope))
reuse_variables = True
grads = opt.compute_gradients(total_loss)
tower_grads.append(grads)
grads = average_gradients(tower_grads)
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
summary_op = tf.summary.merge_all()
log_image, log_image_data, log_image_name = build_image_summary()
# save moving average
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
# batch norm updates
with tf.control_dependencies([variables_averages_op, apply_gradient_op, batch_norm_updates_op]):
train_op = tf.no_op(name='train_op')
saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)
summary_writer = tf.summary.FileWriter(FLAGS.logs_path + StyleTime, tf.get_default_graph())
if FLAGS.pretrained_model_path is not None:
variable_restore_op = slim.assign_from_checkpoint_fn(FLAGS.pretrained_model_path,
slim.get_trainable_variables(),
ignore_missing_vars=True)
global_vars_init_op = tf.global_variables_initializer()
local_vars_init_op = tf.local_variables_initializer()
init = tf.group(local_vars_init_op, global_vars_init_op)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)
config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
restore_step = 0
if FLAGS.restore:
sess.run(init)
print('continue training from previous checkpoint')
ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
restore_step = int(ckpt.split('.')[0].split('_')[-1])
saver.restore(sess, ckpt)
else:
sess.run(init)
if FLAGS.pretrained_model_path is not None:
variable_restore_op(sess)
start = time.time()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
try:
while not coord.should_stop():
for step in range(restore_step, FLAGS.max_steps):
if step != 0 and step % FLAGS.decay_steps == 0:
sess.run(tf.assign(learning_rate, learning_rate.eval() * FLAGS.decay_rate))
ml, tl, _ = sess.run([model_loss, total_loss, train_op])
if np.isnan(tl):
print('Loss diverged, stop training')
break
if step % 10 == 0:
avg_time_per_step = (time.time() - start) / 10
start = time.time()
print('Step {:06d}, model loss {:.4f}, total loss {:.4f}, {:.3f} seconds/step, lr: {:.7f}'). \
format(step, ml, tl, avg_time_per_step, learning_rate.eval())
if (step + 1) % FLAGS.save_checkpoint_steps == 0:
filename = ('SESNet' + '_step_{:d}'.format(step + 1) + '.ckpt')
filename = os.path.join(FLAGS.checkpoint_path, filename)
saver.save(sess, filename)
print('Write model to: {:s}'.format(filename))
if step % FLAGS.save_summary_steps == 0:
_, tl, summary_str = sess.run([train_op, total_loss, summary_op])
summary_writer.add_summary(summary_str, global_step=step)
if step % FLAGS.save_image_steps == 0:
log_image_name_str = ('%06d' % step)
img_split, seg_split, pred = sess.run([iis, isms, output_pred])
img_split = np.squeeze(img_split)[0]
seg_split = np.squeeze(seg_split)[0]
pred = np.squeeze(pred)[0] # why cannot batch_size = 1
# img_split = img_split[0]
# seg_split = seg_split[0]
# pred = pred[0]
# print(np.max(seg_split))
img_split = cv2.resize(img_split,(128,128))
color_seg = np.zeros((seg_split.shape[0], seg_split.shape[1], 3))
for i in range(seg_split.shape[0]):
for j in range(seg_split.shape[1]):
color_seg[i, j, :] = color_map[str(seg_split[i][j])]
color_pred = np.zeros((pred.shape[0], pred.shape[1], 3))
for i in range(pred.shape[0]):
for j in range(pred.shape[1]):
color_pred[i, j, :] = color_map[str(class_labels[pred[i][j]])]
write_img = np.hstack((img_split, color_seg, color_pred))
log_image_summary_op = sess.run(log_image, feed_dict={log_image_name: log_image_name_str, \
log_image_data: write_img})
summary_writer.add_summary(log_image_summary_op, global_step=step)
except tf.errors.OutOfRangeError:
print('finish')
finally:
coord.request_stop()
coord.join(threads)
if __name__ == '__main__':
tf.app.run()
