from distutils.version import LooseVersion
import tensorflow as tf
import warnings
from tqdm import trange
import sys
import os.path
import scipy.misc
import shutil
from glob import glob
from collections import deque
import numpy as np
import time
from helpers.tf_variable_summaries import add_variable_summaries
from helpers.visualization_utils import print_segmentation_onto_image, create_split_view
class FCN8s:
def __init__(self, model_load_dir=None, tags=None, vgg16_dir=None, num_classes=None, variables_load_dir=None):
'''
Arguments:
model_load_dir (string, optional): The directory path to a `SavedModel`, i.e. to the directory
that contains a saved FCN-8s model protocol buffer. If a path is provided, the targeted model will
be loaded. If no path is given, the model will be built from scratch on top of a pre-trained,
convolutionalized VGG-16 base network. `model_load_dir` and `vgg16_dir` may not both be `None`.
tags (list, optional): Only relevant if a path to a saved FCN-8s model is given in `model_load_dir`.
A list of strings containing the tags required to load the appropriate metagraph.
vgg16_dir (string, optional): Only relevant if no path to a saved FCN-8s model is given in `model_load_dir`.
The directory that contains a pretrained, convolutionalized VGG-16 model in the form of a protocol buffer.
`model_load_dir` and `vgg16_dir` may not both be `None`.
num_classes (int, optional): Only relevant if no path to a saved FCN-8s model is given in `model_load_dir`.
The number of segmentation classes.
variables_load_dir (string, optional): The path to variables that were saved with `tf.train.Saver`.
Only relevant if `model_load_dir` is `None`.
'''
# Check TensorFlow version
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'This program requires TensorFlow version 1.0 or newer. You are using {}'.format(tf.__version__)
print('TensorFlow Version: {}'.format(tf.__version__))
if (model_load_dir is None) and (vgg16_dir is None or num_classes is None):
raise ValueError("You must provide either both `model_load_dir` and `tags` or both `vgg16_dir` and `num_classes`.")
self.variables_load_dir = variables_load_dir
self.model_load_dir = model_load_dir
self.tags = tags
self.vgg16_dir = vgg16_dir
self.vgg16_tag = 'vgg16'
self.num_classes = num_classes
self.variables_updated = False # Keep track of whether any variable values changed since this model was last saved.
self.eval_dataset = None # Which dataset to use for evaluation during training. Only relevant for training.
# The following lists store data about the metrics being tracked.
# Note that `self.metric_value_tensors` and `self.metric_update_ops` represent
# the metrics being tracked, not the metrics generally available in the model.
self.metric_names = [] # Store the metric names here.
self.metric_values = [] # Store the latest metric evaluations here.
self.best_metric_values = [] # Keep score of the best historical metric values.
self.metric_value_tensors = [] # Store the value tensors from tf.metrics here.
self.metric_update_ops = [] # Store the update ops from tf.metrics here.
self.training_loss = None
self.best_training_loss = 99999999.9
self.sess = tf.Session()
self.g_step = None # The global step
##################################################################
# Load or build the model.
##################################################################
if not model_load_dir is None: # Load the full pre-trained model.
tf.saved_model.loader.load(sess=self.sess, tags=self.tags, export_dir=self.model_load_dir)
graph = tf.get_default_graph()
# Get the input and output ops.
self.image_input = graph.get_tensor_by_name('image_input:0')
self.keep_prob = graph.get_tensor_by_name('keep_prob:0')
self.fcn8s_output = graph.get_tensor_by_name('decoder/fcn8s_output:0')
self.l2_regularization_rate = graph.get_tensor_by_name('l2_regularization_rate:0')
self.labels = graph.get_tensor_by_name('labels_input:0')
self.total_loss = graph.get_tensor_by_name('optimizer/total_loss:0')
self.train_op = graph.get_tensor_by_name('optimizer/train_op:0')
self.learning_rate = graph.get_tensor_by_name('optimizer/learning_rate:0')
self.global_step = graph.get_tensor_by_name('optimizer/global_step:0')
self.softmax_output = graph.get_tensor_by_name('predictor/softmax_output:0')
self.predictions_argmax = graph.get_tensor_by_name('predictor/predictions_argmax:0')
self.mean_loss_value = graph.get_tensor_by_name('metrics/mean_loss_value:0')
self.mean_loss_update_op = graph.get_tensor_by_name('metrics/mean_loss_update_op:0')
self.mean_iou_value = graph.get_tensor_by_name('metrics/mean_iou_value:0')
self.mean_iou_update_op = graph.get_tensor_by_name('metrics/mean_iou_update_op:0')
self.acc_value = graph.get_tensor_by_name('metrics/acc_value:0')
self.acc_update_op = graph.get_tensor_by_name('metrics/acc_update_op:0')
self.metrics_reset_op = graph.get_operation_by_name('metrics/metrics_reset_op')
self.summaries_training = graph.get_tensor_by_name('summaries_training:0')
self.summaries_evaluation = graph.get_tensor_by_name('summaries_evaluation:0')
# For some reason that I don't understand, the local variables belonging to the
# metrics need to be initialized after loading the model.
self.sess.run(self.metrics_reset_op)
else: # Load only the pre-trained VGG-16 encoder and build the rest of the graph from scratch.
# Load the pretrained convolutionalized VGG-16 model as our encoder.
self.image_input, self.keep_prob, self.pool3_out, self.pool4_out, self.fc7_out = self._load_vgg16()
# Build the decoder on top of the VGG-16 encoder.
self.fcn8s_output, self.l2_regularization_rate = self._build_decoder()
# Build the part of the graph that is relevant for the training.
self.labels = tf.placeholder(dtype=tf.int32, shape=[None, None, None, self.num_classes], name='labels_input')
self.total_loss, self.train_op, self.learning_rate, self.global_step = self._build_optimizer()
# Add the prediction outputs.
self.softmax_output, self.predictions_argmax = self._build_predictor()
# Add metrics for evaluation.
self.mean_loss_value, self.mean_loss_update_op, self.mean_iou_value, self.mean_iou_update_op, self.acc_value, self.acc_update_op, self.metrics_reset_op = self._build_metrics()
# Add summary ops for TensorBoard.
self.summaries_training, self.summaries_evaluation = self._build_summary_ops()
# Initialize the global and local (for the metrics) variables.
self.sess.run(tf.global_variables_initializer())
self.sess.run(tf.local_variables_initializer())
# Maybe load variables.
if not variables_load_dir is None:
saver = tf.train.Saver()
saver.restore(self.sess, variables_load_dir)
def _load_vgg16(self):
'''
Loads the pretrained, convolutionalized VGG-16 model into the session.
'''
# 1: Load the model
tf.saved_model.loader.load(sess=self.sess, tags=[self.vgg16_tag], export_dir=self.vgg16_dir)
# 2: Return the tensors of interest
graph = tf.get_default_graph()
vgg16_image_input_tensor_name = 'image_input:0'
vgg16_keep_prob_tensor_name = 'keep_prob:0'
vgg16_pool3_out_tensor_name = 'layer3_out:0'
vgg16_pool4_out_tensor_name = 'layer4_out:0'
vgg16_fc7_out_tensor_name = 'layer7_out:0'
image_input = graph.get_tensor_by_name(vgg16_image_input_tensor_name)
keep_prob = graph.get_tensor_by_name(vgg16_keep_prob_tensor_name)
pool3_out = graph.get_tensor_by_name(vgg16_pool3_out_tensor_name)
pool4_out = graph.get_tensor_by_name(vgg16_pool4_out_tensor_name)
fc7_out = graph.get_tensor_by_name(vgg16_fc7_out_tensor_name)
return image_input, keep_prob, pool3_out, pool4_out, fc7_out
def _build_decoder(self):
'''
Builds the FCN-8s decoder given the pool3, pool4, and fc7 outputs of the VGG-16 encoder.
'''
stddev_1x1 = 0.001 # Standard deviation for the 1x1 kernel initializers
stddev_conv2d_trans = 0.01 # Standard deviation for the convolution transpose kernel initializers
l2_regularization_rate = tf.placeholder(dtype=tf.float32, shape=[], name='l2_regularization_rate') # L2 regularization rate for the kernels
with tf.name_scope('decoder'):
# 1: Append 1x1 convolutions to the three output layers of the encoder to reduce the Number
# of channels to the number of classes.
# The outputs of pool3 and pool4 are being scaled in what the authors of
# the paper call the at-once training approach.
pool3_out_scaled = tf.multiply(self.pool3_out, 0.0001, name='pool3_out_scaled')
pool3_1x1 = tf.layers.conv2d(inputs=pool3_out_scaled,
filters=self.num_classes,
kernel_size=(1, 1),
strides=(1, 1),
padding='same',
kernel_initializer=tf.truncated_normal_initializer(stddev=stddev_1x1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(l2_regularization_rate),
name='pool3_1x1')
pool4_out_scaled = tf.multiply(self.pool4_out, 0.01, name='pool4_out_scaled')
pool4_1x1 = tf.layers.conv2d(inputs=pool4_out_scaled,
filters=self.num_classes,
kernel_size=(1, 1),
strides=(1, 1),
padding='same',
kernel_initializer=tf.truncated_normal_initializer(stddev=stddev_1x1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(l2_regularization_rate),
name='pool4_1x1')
fc7_1x1 = tf.layers.conv2d(inputs=self.fc7_out,
filters=self.num_classes,
kernel_size=(1, 1),
strides=(1, 1),
padding='same',
kernel_initializer=tf.truncated_normal_initializer(stddev=stddev_1x1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(l2_regularization_rate),
name='fc7_1x1')
# 2: Upscale and fuse until we're back at the original image size.
fc7_conv2d_trans = tf.layers.conv2d_transpose(inputs=fc7_1x1,
filters=self.num_classes,
kernel_size=(4, 4),
strides=(2, 2),
padding='same',
kernel_initializer=tf.truncated_normal_initializer(stddev=stddev_conv2d_trans),
kernel_regularizer=tf.contrib.layers.l2_regularizer(l2_regularization_rate),
name='fc7_conv2d_trans')
add_fc7_pool4 = tf.add(fc7_conv2d_trans, pool4_1x1, name='add_fc7_pool4')
fc7_pool4_conv2d_trans = tf.layers.conv2d_transpose(inputs=add_fc7_pool4,
filters=self.num_classes,
kernel_size=(4, 4),
strides=(2, 2),
padding='same',
kernel_initializer=tf.truncated_normal_initializer(stddev=stddev_conv2d_trans),
kernel_regularizer=tf.contrib.layers.l2_regularizer(l2_regularization_rate),
name='fc7_pool4_conv2d_trans')
add_fc7_pool4_pool3 = tf.add(fc7_pool4_conv2d_trans, pool3_1x1, name='add_fc7_pool4_pool3')
fc7_pool4_pool3_conv2d_trans = tf.layers.conv2d_transpose(inputs=add_fc7_pool4_pool3,
filters=self.num_classes,
kernel_size=(16, 16),
strides=(8, 8),
padding='same',
kernel_initializer=tf.truncated_normal_initializer(stddev=stddev_conv2d_trans),
kernel_regularizer=tf.contrib.layers.l2_regularizer(l2_regularization_rate),
name='fc7_pool4_pool3_conv2d_trans')
fcn8s_output = tf.identity(fc7_pool4_pool3_conv2d_trans, name='fcn8s_output')
return fc7_pool4_pool3_conv2d_trans, l2_regularization_rate
def _build_optimizer(self):
'''
Builds the training-relevant part of the graph.
'''
with tf.name_scope('optimizer'):
# Create a training step counter.
global_step = tf.Variable(0, trainable=False, name='global_step')
# Create placeholder for the learning rate.
learning_rate = tf.placeholder(dtype=tf.float32, shape=[], name='learning_rate')
# Compute the regularizatin loss.
regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES) # This is a list of the individual loss values, so we still need to sum them up.
regularization_loss = tf.add_n(regularization_losses, name='regularization_loss') # Scalar
# Compute the total loss.
approximation_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=self.labels, logits=self.fcn8s_output), name='approximation_loss') # Scalar
total_loss = tf.add(approximation_loss, regularization_loss, name='total_loss')
# Compute the gradients and apply them.
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate, name='adam_optimizer')
train_op = optimizer.minimize(total_loss, global_step=global_step, name='train_op')
return total_loss, train_op, learning_rate, global_step
def _build_predictor(self):
'''
Builds the prediction-relevant part of the graph.
'''
with tf.name_scope('predictor'):
softmax_output = tf.nn.softmax(self.fcn8s_output, name='softmax_output')
predictions_argmax = tf.argmax(softmax_output, axis=-1, name='predictions_argmax', output_type=tf.int64)
return softmax_output, predictions_argmax
def _build_metrics(self):
'''
Builds the evaluation-relevant part of the graph, i.e. the metrics operations.
'''
with tf.variable_scope('metrics') as scope:
labels_argmax = tf.argmax(self.labels, axis=-1, name='labels_argmax', output_type=tf.int64)
# 1: Mean loss
mean_loss_value, mean_loss_update_op = tf.metrics.mean(self.total_loss)
mean_loss_value = tf.identity(mean_loss_value, name='mean_loss_value')
mean_loss_update_op = tf.identity(mean_loss_update_op, name='mean_loss_update_op')
# 1: Mean IoU
mean_iou_value, mean_iou_update_op = tf.metrics.mean_iou(labels=labels_argmax,
predictions=self.predictions_argmax,
num_classes=self.num_classes)
mean_iou_value = tf.identity(mean_iou_value, name='mean_iou_value')
mean_iou_update_op = tf.identity(mean_iou_update_op, name='mean_iou_update_op')
# 2: Accuracy
acc_value, acc_update_op = tf.metrics.accuracy(labels=labels_argmax,
predictions=self.predictions_argmax)
acc_value = tf.identity(acc_value, name='acc_value')
acc_update_op = tf.identity(acc_update_op, name='acc_update_op')
# As of version 1.3, TensorFlow's streaming metrics don't have reset operations,
# so we need to create our own as a work-around. Say we want to evaluate
# a metric after every training epoch. If we didn't have
# a way to reset the metric's update op after every evaluation,
# the computed metric value would be the average of the current evaluation
# and all previous evaluations from past epochs, which is obviously not
# what we want.
local_metric_vars = tf.contrib.framework.get_variables(scope=scope, collection=tf.GraphKeys.LOCAL_VARIABLES)
metrics_reset_op = tf.variables_initializer(var_list=local_metric_vars, name='metrics_reset_op')
return (mean_loss_value,
mean_loss_update_op,
mean_iou_value,
mean_iou_update_op,
acc_value,
acc_update_op,
metrics_reset_op)
def _build_summary_ops(self):
'''
Builds the part of the graph that logs summaries for TensorBoard.
'''
graph = tf.get_default_graph()
add_variable_summaries(variable=graph.get_tensor_by_name('pool3_1x1/kernel:0'), scope='pool3_1x1/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('pool3_1x1/bias:0'), scope='pool3_1x1/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('pool4_1x1/kernel:0'), scope='pool4_1x1/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('pool4_1x1/bias:0'), scope='pool4_1x1/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_1x1/kernel:0'), scope='fc7_1x1/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_1x1/bias:0'), scope='fc7_1x1/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_conv2d_trans/kernel:0'), scope='fc7_conv2d_trans/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_conv2d_trans/bias:0'), scope='fc7_conv2d_trans/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_pool4_conv2d_trans/kernel:0'), scope='fc7_pool4_conv2d_trans/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_pool4_conv2d_trans/bias:0'), scope='fc7_pool4_conv2d_trans/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_pool4_pool3_conv2d_trans/kernel:0'), scope='fc7_pool4_pool3_conv2d_trans/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7_pool4_pool3_conv2d_trans/bias:0'), scope='fc7_pool4_pool3_conv2d_trans/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7/weights:0'), scope='fc7/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('fc7/biases:0'), scope='fc7/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('fc6/weights:0'), scope='fc6/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('fc6/biases:0'), scope='fc6/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('conv4_3/filter:0'), scope='conv4_3/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('conv4_3/biases:0'), scope='conv4_3/bias')
add_variable_summaries(variable=graph.get_tensor_by_name('conv3_3/filter:0'), scope='conv3_3/kernel')
add_variable_summaries(variable=graph.get_tensor_by_name('conv3_3/biases:0'), scope='conv3_3/bias')
# Loss and learning rate.
tf.summary.scalar('total_loss', self.total_loss)
tf.summary.scalar('learning_rate', self.learning_rate)
summaries_training = tf.summary.merge_all()
summaries_training = tf.identity(summaries_training, name='summaries_training')
# All metrics.
mean_loss = tf.summary.scalar('mean_loss', self.mean_loss_value)
mean_iou = tf.summary.scalar('mean_iou', self.mean_iou_value)
accuracy = tf.summary.scalar('accuracy', self.acc_value)
summaries_evaluation = tf.summary.merge(inputs=[mean_loss,
mean_iou,
accuracy])
summaries_evaluation = tf.identity(summaries_evaluation, name='summaries_evaluation')
return summaries_training, summaries_evaluation
def _initialize_metrics(self, metrics):
'''
Initializes/resets the metrics before every call to `train` and `evaluate`.
'''
# Reset lists of previous tracked metrics.
self.metric_names = []
self.best_metric_values = []
self.metric_update_ops = []
self.metric_value_tensors = []
# Set the metrics that will be evaluated.
if 'loss' in metrics:
self.metric_names.append('loss')
self.best_metric_values.append(99999999.9)
self.metric_update_ops.append(self.mean_loss_update_op)
self.metric_value_tensors.append(self.mean_loss_value)
if 'mean_iou' in metrics:
self.metric_names.append('mean_iou')
self.best_metric_values.append(0.0)
self.metric_update_ops.append(self.mean_iou_update_op)
self.metric_value_tensors.append(self.mean_iou_value)
if 'accuracy' in metrics:
self.metric_names.append('accuracy')
self.best_metric_values.append(0.0)
self.metric_update_ops.append(self.acc_update_op)
self.metric_value_tensors.append(self.acc_value)
def train(self,
train_generator,
epochs,
steps_per_epoch,
learning_rate_schedule,
keep_prob=0.5,
l2_regularization=0.0,
eval_dataset='train',
eval_frequency=5,
val_generator=None,
val_steps=None,
metrics={},
save_during_training=False,
save_dir=None,
save_best_only=True,
save_tags=['default'],
save_name='',
save_frequency=5,
saver='saved_model',
monitor='loss',
record_summaries=True,
summaries_frequency=10,
summaries_dir=None,
summaries_name=None,
training_loss_display_averaging=3):
'''
Trains the model.
Arguments:
train_generator (generator): A generator that yields batches of images
and associated ground truth images in two separate Numpy arrays.
The images must be a 4D array with format `(batch_size, height, width, channels)`
and the ground truth images must be a 4D array with format
`(batch_size, height, width, num_classes)`, i.e. the ground truth
data must be provided in one-hot format.
epochs (int): The number of epochs to run the training for, where each epoch
consists of `steps_per_epoch` training steps.
steps_per_epoch (int): The number of training steps (i.e. batches processed)
per epoch.
learning_rate_schedule (function): Any function that takes as its sole input
an integer (the global step counter) and returns a float (the learning rate).
keep_prob (float, optional): The keep probability for the two dropout layers
in the VGG-16 encoder network. Defaults to 0.5.
l2_regularization (float, optional): The scaling factor for the L2 regularization
of all decoder kernels. 0 means no regularization at all. This has no effect
on the kernels of the VGG-16 encoder network. Defaults to 0.
eval_dataset (string, optional): Which generator to use for the evaluation
of the model during training. Can be either of 'train' (the train_generator
will be used) or 'val' (the val_generator will be used). Defaults to 'train',
but should be set to 'val' if a validation dataset is available.
eval_frequency (int, optional): The model will be evaluated on `metrics` after every
`eval_frequency` epochs. Defaults to 5.
val_generator (generator, optional): An optional second generator for a second
dataset (validation dataset), works the same way as `train_generator`.
val_steps (int, optional): The number of steps to run `val_generator` for
during evaluation.
metrics (set, optional): The metrics to be evaluated during training. A Python
set containing any subset of `{'loss', 'mean_iou', 'accuracy'}`, which are the
currently available metrics. Defaults to the empty set, meaning that the
model will not be evaluated during training.
save_during_training (bool, optional): Whether or not to save the model periodically
during training, the parameters of which can be set in the subsequent arguments.
Defaults to `False`.
save_dir (string, optional): The full path of the directory to save the model to
during training.
save_best_only (bool, optional): If `True`, the model will only be saved upon
evaluation if the metric defined by `monitor` has improved since it was last
measured before. Can only be `True` if `metrics` is not empty.
save_tags (list, optional): An optional list of tags to save the model metagraph
with in the SavedModel protocol buffer. Defaults to a list only containing
the tag 'default'. At least one tag must be given.
save_name (string, optional): An optional name string to include in the name of
the folder in which the model will be saved during training. Note that what
you pass as the name here will be only part of the folder name. The folder
name also includes a count of the global training step and the values of
any metrics that are being evaluate, although at least the training loss.
It is hence not necessary to pass a name here, each saved model will be
uniquely and descriptively named regardless. Defaults to the empty string.
save_frequency (int, optional): The model will be saved at most after every
`save_frequency` epochs, but possibly less often if `save_best_only` is `True`
and if there was no improvement in the monitored metric. Defaults to 5.
saver (string, optional): Which saver to use when saving the model during training.
Can be either of 'saved_model' in order to use `tf.saved_model` or 'train_saver'
in order to use `tf.train.Saver`. Defaults to `tf.saved_model`. Check the
TensorFlow documentation for details on which saver might be better for your
use case. In general you can't go wrong with either of the two.
monitor (string, optional): The name of the metric that is to be monitored in
order to decide whether the model should be saved. Can be one of
`{'loss', 'mean_iou', 'accuracy'}`, which are the currently available metrics.
Defaults to 'loss'.
record_summaries (bool, optional): Whether or not to record TensorBoard summaries.
Defaults to `True`.
summaries_frequency (int, optional): How often summaries should be logged for
tensors which are updated at every training step. The summaries for such tensors
will be recorded every `summaries_frequency` training steps. Defaults to 10.
summaries_dir (string, optional): The full path of the directory to which to
write the summaries protocol buffers.
summaries_name (string, optional): The name of the summaries buffers.
training_loss_display_averaging (int, optional): During training, the current
training loss is always displayed. Since training on mini-batches has the effect
that the loss might jump from training step to training step, this parameter
allows to average the displayed loss over tha lasst `training_loss_display_averaging`
training steps so that it shows a more representative picture of the actual
current loss. Defaults to 3.
'''
# Check for a GPU
if not tf.test.gpu_device_name():
warnings.warn('No GPU found. Please note that training this network will be unbearably slow without a GPU.')
else:
print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
if not eval_dataset in ['train', 'val']:
raise ValueError("`eval_dataset` must be one of 'train' or 'val', but is '{}'.".format(eval_dataset))
if (eval_dataset == 'val') and ((val_generator is None) or (val_steps is None)):
raise ValueError("When eval_dataset == 'val', a `val_generator` and `val_steps` must be passed.")
for metric in metrics:
if not metric in ['loss', 'mean_iou', 'accuracy']:
raise ValueError("{} is not a valid metric. Valid metrics are ['loss', mean_iou', 'accuracy']".format(metric))
if (not monitor in metrics) and (not monitor == 'loss'):
raise ValueError('You are trying to monitor {}, but it is not in `metrics` and is therefore not being computed.'.format(monitor))
self.eval_dataset = eval_dataset
self.g_step = self.sess.run(self.global_step)
learning_rate = learning_rate_schedule(self.g_step)
self._initialize_metrics(metrics)
# Set up the summary file writers.
if record_summaries:
training_writer = tf.summary.FileWriter(logdir=os.path.join(summaries_dir, summaries_name),
graph=self.sess.graph)
if len(metrics) > 0:
evaluation_writer = tf.summary.FileWriter(logdir=os.path.join(summaries_dir, summaries_name+'_eval'))
for epoch in range(1, epochs+1):
##############################################################
# Run the training for this epoch.
##############################################################
loss_history = deque(maxlen=training_loss_display_averaging)
tr = trange(steps_per_epoch, file=sys.stdout)
tr.set_description('Epoch {}/{}'.format(epoch, epochs))
for train_step in tr:
batch_images, batch_labels = next(train_generator)
if record_summaries and (self.g_step % summaries_frequency == 0):
_, current_loss, self.g_step, training_summary = self.sess.run([self.train_op,
self.total_loss,
self.global_step,
self.summaries_training],
feed_dict={self.image_input: batch_images,
self.labels: batch_labels,
self.learning_rate: learning_rate,
self.keep_prob: keep_prob,
self.l2_regularization_rate: l2_regularization})
training_writer.add_summary(summary=training_summary, global_step=self.g_step)
else:
_, current_loss, self.g_step = self.sess.run([self.train_op,
self.total_loss,
self.global_step],
feed_dict={self.image_input: batch_images,
self.labels: batch_labels,
self.learning_rate: learning_rate,
self.keep_prob: keep_prob,
self.l2_regularization_rate: l2_regularization})
self.variables_updated = True
loss_history.append(current_loss)
losses = np.array(loss_history)
self.training_loss = np.mean(losses)
tr.set_postfix(ordered_dict={'loss': self.training_loss,
'learning rate': learning_rate})
learning_rate = learning_rate_schedule(self.g_step)
##############################################################
# Maybe evaluate the model after this epoch.
##############################################################
if (len(metrics) > 0) and (epoch % eval_frequency == 0):
if eval_dataset == 'train':
data_generator = train_generator
num_batches = steps_per_epoch
description = 'Evaluation on training dataset'
elif eval_dataset == 'val':
data_generator = val_generator
num_batches = val_steps
description = 'Evaluation on validation dataset'
self._evaluate(data_generator=data_generator,
metrics=metrics,
num_batches=num_batches,
l2_regularization=l2_regularization,
description=description)
if record_summaries:
evaluation_summary = self.sess.run(self.summaries_evaluation)
evaluation_writer.add_summary(summary=evaluation_summary, global_step=self.g_step)
##############################################################
# Maybe save the model after this epoch.
##############################################################
if save_during_training and (epoch % save_frequency == 0):
save = False
if save_best_only:
if (monitor == 'loss' and
(not 'loss' in self.metric_names) and
self.training_loss < self.best_training_loss):
save = True
else:
i = self.metric_names.index(monitor)
if (monitor == 'loss') and (self.metric_values[i] < self.best_metric_values[i]):
save = True
elif (monitor in ['accuracry', 'mean_iou']) and (self.metric_values[i] > self.best_metric_values[i]):
save = True
if save:
print('New best {} value, saving model.'.format(monitor))
else:
print('No improvement over previous best {} value, not saving model.'.format(monitor))
else:
save = True
if save:
self.save(model_save_dir=save_dir,
saver=saver,
tags=save_tags,
name=save_name,
include_global_step=True,
include_last_training_loss=True,
include_metrics=(len(self.metric_names) > 0))
##############################################################
# Update the current best metric values.
##############################################################
if self.training_loss < self.best_training_loss:
self.best_training_loss = self.training_loss
if epoch % eval_frequency == 0:
for i, metric_name in enumerate(self.metric_names):
if (metric_name == 'loss') and (self.metric_values[i] < self.best_metric_values[i]):
self.best_metric_values[i] = self.metric_values[i]
elif (metric_name in ['accuracry', 'mean_iou']) and (self.metric_values[i] > self.best_metric_values[i]):
self.best_metric_values[i] = self.metric_values[i]
def _evaluate(self, data_generator, metrics, num_batches, l2_regularization, description='Running evaluation'):
'''
Internal method used by both `evaluate()` and `train()` that performs
the actual evaluation. For the first three arguments, please refer
to the documentation of the public `evaluate()` method.
Arguments:
description (string, optional): A description string that will be prepended
to the progress bar while the evaluation is being processed. During
training, this description is used to clarify whether the evaluation
is being performed on the training or validation dataset.
'''
# Reset all metrics' accumulator variables.
self.sess.run(self.metrics_reset_op)
# Set up the progress bar.
tr = trange(num_batches, file=sys.stdout)
tr.set_description(description)
# Accumulate metrics in batches.
for step in tr:
batch_images, batch_labels = next(data_generator)
self.sess.run(self.metric_update_ops,
feed_dict={self.image_input: batch_images,
self.labels: batch_labels,
self.keep_prob: 1.0,
self.l2_regularization_rate: l2_regularization})
# Compute final metric values.
self.metric_values = self.sess.run(self.metric_value_tensors)
evaluation_results_string = ''
for i, metric_name in enumerate(self.metric_names):
evaluation_results_string += metric_name + ': {:.4f} '.format(self.metric_values[i])
print(evaluation_results_string)
def evaluate(self, data_generator, num_batches, metrics={'loss', 'mean_iou', 'accuracy'}, l2_regularization=0.0, dataset='val'):
'''
Evaluates the model on the given metrics on the data generated by `data_generator`.
Arguments:
data_generator (generator): A generator that yields batches of images
and associated ground truth images in two separate Numpy arrays.
The images must be a 4D array with format `(batch_size, height, width, channels)`
and the ground truth images must be a 4D array with format
`(batch_size, height, width, num_classes)`, i.e. the ground truth
data must be provided in one-hot format. The generator's batch size
has no effect on the outcome of the evaluation.
num_batches (int): The number of batches to evaluate the model on.
Typically this will be the number of batches such that the model
is being evaluated on the whole evaluation dataset.
metrics (set, optional): The metrics to be evaluated. A Python set containing
any subset of `{'loss', 'mean_iou', 'accuracy'}`, which are the
currently available metrics. Defaults to the full set.
dataset (string, optional): Specifies the kind of dataset on which the model
is being evaluated. Should be set to 'train' if the model is being evaluated
on a dataset on which it has also been trained, or 'val' if the model is
being evaluated on a dataset which it has never seen during training.
This argument has no effect on the evaluation of the model, but if you
save the model using `save()` after evaluating it, the model name will
include this value to indicate whether or not the metric values were
achieved on a dataset that has not been used during training. Defaults to 'val'.
'''
for metric in metrics:
if not metric in ['loss', 'mean_iou', 'accuracy']:
raise ValueError("{} is not a valid metric. Valid metrics are ['loss', mean_iou', 'accuracy']".format(metric))
if not dataset in {'train', 'val'}:
raise ValueError("`dataset` must be either 'train' or 'val'.")
self._initialize_metrics(metrics)
self._evaluate(data_generator, metrics, num_batches, l2_regularization, description='Running evaluation')
if dataset == 'val':
self.eval_dataset = 'val'
else:
self.eval_dataset = 'train'
def predict(self, images, argmax=True):
'''
Makes predictions for the input images.
Arguments:
images (array-like): The input image or images. Must be an array-like
object of rank 4. If predicting only one image, encapsulate it in
a Python list.
argmax (bool, optional): If `True`, the model predicts class IDs,
i.e. the last dimension has length 1 and an integer between
zero and `num_classes - 1` for each pixel. Otherwise, the model
outputs the softmax distribution, i.e. the last dimension has
length `num_classes` and contains the probability for each class
for all pixels. Defaults to `True`.
Returns:
The prediction, an array of rank 4 of which the first three dimensions
are identical to the input and the fourth dimension is as described
in `argmax`.
'''
if argmax:
return self.sess.run(self.predictions_argmax,
feed_dict={self.image_input: images,
self.keep_prob: 1.0})
else:
return self.sess.run(self.softmax_output,
feed_dict={self.image_input: images,
self.keep_prob: 1.0})
def predict_and_save(self,
results_dir,
images_dir,
color_map,
resize=False,
image_file_extension='png',
include_unprocessed_image=False,
arrangement='vertical',
overwrite_existing=True):
'''
Makes predictions for all images in a given directory, overlays a copy of the
input images with the respective predictions, and saves the resulting images to disk.
Arguments:
results_dir (string): The directory in which to save the annotated prediction
output images.
images_dir (string): The directory in which the images to be processed are located.
color_map (dictionary): A Python dictionary whose keys are non-negative
integers representing segmentation classes and whose values are 1D tuples
(or lists, Numpy arrays) of length 4 that represent the RGBA color values
in which the respective classes are to be annotated. For example, if the
dictionary contains the key-value pair `{1: (0, 255, 0, 127)}`, then
this means that all pixels in the predicted image segmentation that belong
to segmentation class 1 will be colored in green with 50% transparency
in the input image.
resize (tuple): `False` or a tuple of the form `(image_height, image_width)` that
represents the size to which all images will be resized.
image_file_extension (string, optional): The file extension of the
images in the datasets. Must be identical for all images in all
datasets in `datasets`. Defaults to `png`.
include_unprocessed_image (bool, optional): If `True`, creates split view images
containing both the input image and the overlayed segmented image.
Defaults to `False`.
arrangement (string, optional): Only relevant if `include_unprocessed_image` is `True`.
Determines the arrangement for the split view. Can be either of 'vertical', meaning
the processed and unprocessed images will be above each other, or 'horizontal',
meaning the processed and unprocessed images will be next to each other.
Defaults to 'vertical'.
overwrite_existing (bool, optional): If `True`, overwrites the output directory
in case it already exists.
'''
# Make a directory in which to store the results.
if overwrite_existing and os.path.exists(results_dir):
shutil.rmtree(results_dir)
os.makedirs(results_dir)
image_paths = glob(os.path.join(images_dir, '*.' + image_file_extension))
num_images = len(image_paths)
print('The segmented images will be saved to "{}"'.format(results_dir))
tr = trange(num_images, file=sys.stdout)
tr.set_description('Processing images')
for i in tr:
filepath = image_paths[i]
image = scipy.misc.imread(filepath)
if resize and not np.array_equal(image.shape[:2], resize):
image = scipy.misc.imresize(image, resize)
img_height, img_width, img_ch = image.shape
prediction = self.predict([image], argmax=False)
processed_image = np.asarray(print_segmentation_onto_image(image=image, prediction=prediction, color_map=color_map), dtype=np.uint8)
if include_unprocessed_image:
if arrangement == 'vertical':
output_width = img_width
output_height = 2 * img_height
processed_image = create_split_view(target_size=(output_height, output_width),
images=[processed_image, image],
positions=[(0, 0), (img_height, 0)],
sizes=[(img_height, img_width), (img_height, img_width)])
else:
output_width = 2 * img_width
output_height = img_height
processed_image = create_split_view(target_size=(output_height, output_width),
images=[processed_image, image],
positions=[(0, 0), (0, img_width)],
sizes=[(img_height, img_width), (img_height, img_width)])
scipy.misc.imsave(os.path.join(results_dir, os.path.basename(filepath)), processed_image)
def save(self,
model_save_dir,
saver,
tags=['default'],
name=None,
include_global_step=True,
include_last_training_loss=True,
include_metrics=True,
force_save=False):
'''
Saves the model to disk.
Arguments:
model_save_dir (string): The full path of the directory to which to
save the model.
saver (string, optional): Which saver to use when saving the model during training.
Can be either of 'saved_model' in order to use `tf.saved_model` or 'train_saver'
in order to use `tf.train.Saver`. Defaults to `tf.saved_model`. Check the
TensorFlow documentation for details on which saver might be better for your
use case. In general you can't go wrong with either of the two.
tags (list, optional): An optional list of tags to save the model metagraph
with in the SavedModel protocol buffer. Defaults to a list only containing
the tag 'default'. At least one tag must be given.
name (string, optional): An optional name that will be part of the name of the
saved model's parent directory. Since you have the possibility to include
the global step number and the values of metrics in the model name, giving
an explicit name here is often not necessary.
include_global_step (bool, optional): Whether or not to include the global
step number in the model name. Defaults to `True`.
include_last_training_loss (bool, optional): Whether of not to include the
last training loss value in the model name. Defaults to `True`.
include_metrics (bool, optional): If `True`, the last values of all recorded
metrics will be included in the model name. Defaults to `True`.
force_save (bool, optional): If `True`, force the saver to save the model
even if no variables have changed since saving last. Defaults to `False`.
'''
if (not self.variables_updated) and (not force_save):
print("Abort: Nothing to save, no training has been performed since the model was last saved.")
return
if not saver in {'saved_model', 'train_saver'}:
raise ValueError("Unexpected value for `saver`: Can be either 'saved_model' or 'train_saver', but received '{}'.".format(saver))
if self.training_loss is None:
include_last_training_loss = False
model_name = 'saved_model'
if not name is None:
model_name += '_' + name
if include_global_step:
self.g_step = self.sess.run(self.global_step)
model_name += '_(globalstep-{})'.format(self.g_step)
if include_last_training_loss:
model_name += '_(trainloss-{:.4f})'.format(self.training_loss)
if include_metrics:
if self.eval_dataset == 'val':
model_name += '_(eval_on_val_dataset)'
else:
model_name += '_(eval_on_train_dataset)'
for i in range(len(self.metric_names)):
model_name += '_({}-{:.4f})'.format(self.metric_names[i], self.metric_values[i])
if not (include_global_step or include_last_training_loss or include_metrics) and (name is None):
model_name += '_{}'.format(time.time())
if saver == 'saved_model':
saved_model_builder = tf.saved_model.builder.SavedModelBuilder(os.path.join(model_save_dir, model_name))
saved_model_builder.add_meta_graph_and_variables(sess=self.sess, tags=tags)
saved_model_builder.save()
else:
saver = tf.train.Saver(var_list=None,
reshape=False,
max_to_keep=5,
keep_checkpoint_every_n_hours=10000.0)
saver.save(self.sess,
save_path=os.path.join(model_save_dir, model_name, 'variables'),
write_meta_graph=True,
write_state=True)
self.variables_updated = False
def load_variables(self, path):
'''
Load variable values into the current model. Only works for variables that
were saved with 'train_saver'. See `save()` for details.
'''
saver = tf.train.Saver(var_list=None)
saver.restore(self.sess, path)
def close(self):
'''
Closes the session. This method is important to call when you are done working
with the model in order to release the resources it occupies.
'''
self.sess.close()
print("The session has been closed.")
