img_input = tf.placeholder(tf.float32, shape=(1, 3, 240, 240, 4))
la_input = tf.placeholder(tf.int32, shape=(1, 3, 240, 240, 1))
is_training = tf.placeholder(tf.bool)
folder = glob.glob(datasetDir + '*.tfrecord')
assert len(folder) == 15
im_queues = {}
for f in folder:
imname = f.split("/")[-1].split(".")[0]
batch = load_batch.get_batch_cmc(f)
im_queues[imname] = batch
net = model_cmc.Model()
logits, _ = net.net(img_input, is_training)
logits = [tf.argmax(l, axis=3) for l in logits]
saver = tf.train.Saver()
sess = tf.Session()
dir = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)
saver.restore(sess, dir)
sess.run(tf.local_variables_initializer())
print("Model restore!")
num_class = 5
hist = np.zeros((num_class, num_class))
out_slices = []
la_slices = []
def main(_):
if not FLAGS.dataset_dir:
raise ValueError(
'You must supply the dataset directory with --dataset_dir')
tf.logging.set_verbosity(tf.logging.DEBUG)
with tf.Graph().as_default():
deploy_config = model_deploy.DeploymentConfig(
num_clones=FLAGS.num_clones,
clone_on_cpu=FLAGS.clone_on_cpu,
replica_id=FLAGS.task,
num_replicas=FLAGS.worker_replicas,
num_ps_tasks=FLAGS.num_ps_tasks)
# Create global_step
with tf.device(deploy_config.variables_device()):
global_step = tf.train.get_or_create_global_step()
net = model_cmc.Model()
with tf.device(deploy_config.inputs_device()):
if (FLAGS.model == 'unet'):
batch_queue = \
load_batch.get_batch(FLAGS.dataset_dir,
FLAGS.num_readers,
FLAGS.batch_size,
None,
FLAGS,
file_pattern = FLAGS.file_pattern,
is_training = True,
shuffe = FLAGS.shuffle_data)
elif (FLAGS.model == 'patch'):
batch_queue = \
load_batch_patch.get_batch(FLAGS.dataset_dir,
FLAGS.num_readers,
FLAGS.batch_size,
None,
FLAGS,
file_pattern = FLAGS.file_pattern,
is_training = True,
shuffe = FLAGS.shuffle_data)
elif (FLAGS.model == 'cmc'):
batch_queue = \
load_batch_cmc.get_batch(FLAGS.dataset_dir,
FLAGS.num_readers,
FLAGS.batch_size,
None,
FLAGS,
file_pattern = FLAGS.file_pattern,
is_training = True,
shuffe = FLAGS.shuffle_data)
# =================================================================== #
# Define the model running on every GPU.
# =================================================================== #
print("Batch_loading_successful")
def clone_fn(batch_queue):
batch_shape = [1] * 3
b_image, label = batch_queue
logits, end_points = net.net(b_image)
# Add loss function.
loss, mean_iou = net.weighted_losses(logits, label)
return end_points, mean_iou
summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
clones = model_deploy.create_clones(deploy_config, clone_fn,
[batch_queue])
first_clone_scope = deploy_config.clone_scope(0)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS,
first_clone_scope)
end_points, mean_iou = clones[0].outputs
update_ops.append(mean_iou[1])
#for end_point in end_points:
# x = end_points[end_point]
# summaries.add(tf.summary.histogram('activations/' + end_point, x))
for loss in tf.get_collection('EXTRA_LOSSES', first_clone_scope):
summaries.add(tf.summary.scalar(loss.op.name, loss))
#
#for variable in slim.get_model_variables():
# summaries.add(tf.summary.histogram(variable.op.name, variable))
#################################
# Configure the moving averages #
#################################
if FLAGS.moving_average_decay:
moving_average_variables = slim.get_model_variables()
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
else:
moving_average_variables, variable_averages = None, None
#########################################
# Configure the optimization procedure. #
#########################################
with tf.device(deploy_config.optimizer_device()):
learning_rate = tf_utils.configure_learning_rate(
FLAGS, FLAGS.num_samples, global_step)
optimizer = tf_utils.configure_optimizer(FLAGS, learning_rate)
summaries.add(tf.summary.scalar('learning_rate', learning_rate))
if FLAGS.fine_tune:
gradient_multipliers = pickle.load(
open('nets/multiplier_300.pkl', 'rb'))
else:
gradient_multipliers = None
if FLAGS.moving_average_decay:
# Update ops executed locally by trainer.
update_ops.append(
variable_averages.apply(moving_average_variables))
# Variables to train.
variables_to_train = tf_utils.get_variables_to_train(FLAGS)
# and returns a train_tensor and summary_op
total_loss, clones_gradients = model_deploy.optimize_clones(
clones, optimizer, var_list=variables_to_train)
# Add total_loss to summary.
summaries.add(tf.summary.scalar('total_loss', total_loss))
if gradient_multipliers:
with ops.name_scope('multiply_grads'):
clones_gradients = slim.learning.multiply_gradients(
clones_gradients, gradient_multipliers)
if FLAGS.clip_gradient_norm > 0:
with ops.name_scope('clip_grads'):
clones_gradients = slim.learning.clip_gradient_norms(
clones_gradients, FLAGS.clip_gradient_norm)
# Create gradient updates.
grad_updates = optimizer.apply_gradients(clones_gradients,
global_step=global_step)
update_ops.append(grad_updates)
update_op = tf.group(*update_ops)
train_tensor = control_flow_ops.with_dependencies([update_op],
total_loss,
name='train_op')
#train_tensor = slim.learning.create_train_op(total_loss, optimizer, gradient_multipliers=gradient_multipliers)
# Add the summaries from the first clone. These contain the summaries
# created by model_fn and either optimize_clones() or _gather_clone_loss().
summaries |= set(
tf.get_collection(tf.GraphKeys.SUMMARIES, first_clone_scope))
# Merge all summaries together.
summary_op = tf.summary.merge(list(summaries), name='summary_op')
# =================================================================== #
# Kicks off the training.
# =================================================================== #
def train_step_fn(session, *args, **kwargs):
# visualizer = Beholder(session=session, logdir=FLAGS.train_dir)
total_loss, should_stop = train_step(session, *args, **kwargs)
if train_step_fn.step % FLAGS.validation_check == 0:
_mean_iou = session.run(train_step_fn.mean_iou)
print('evaluation step %d - loss = %.4f mean_iou = %.2f%%' %\
(train_step_fn.step, total_loss, _mean_iou ))
# evaluated_tensors = session.run([end_points['conv4'], end_points['up1']])
# example_frame = session.run(end_points['up2'])
# visualizer.update(arrays=evaluated_tensors, frame=example_frame)
train_step_fn.step += 1
return [total_loss, should_stop]
train_step_fn.step = 0
train_step_fn.end_points = end_points
train_step_fn.mean_iou = mean_iou[0]
gpu_options = tf.GPUOptions(
per_process_gpu_memory_fraction=FLAGS.gpu_memory_fraction,
allocator_type="BFC")
config = tf.ConfigProto(
gpu_options=gpu_options,
log_device_placement=False,
allow_soft_placement=True,
inter_op_parallelism_threads=0,
intra_op_parallelism_threads=1,
)
saver = tf.train.Saver(max_to_keep=5,
keep_checkpoint_every_n_hours=1.0,
write_version=2,
pad_step_number=False)
slim.learning.train(train_tensor,
logdir=FLAGS.train_dir,
master='',
is_chief=True,
train_step_fn=train_step_fn,
init_fn=tf_utils.get_init_fn(FLAGS),
summary_op=summary_op,
number_of_steps=FLAGS.max_number_of_steps,
log_every_n_steps=FLAGS.log_every_n_steps,
save_summaries_secs=FLAGS.save_summaries_secs,
saver=saver,
save_interval_secs=FLAGS.save_interval_secs,
session_config=config,
sync_optimizer=None)