Created
October 15, 2017 17:59
-
-
Save Goddard/187e2ec3348a29966b03b0ce91b8e6b4 to your computer and use it in GitHub Desktop.
main network file for semantic segmentation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os.path | |
import tensorflow as tf | |
import helper | |
import freeze | |
import warnings | |
from distutils.version import LooseVersion | |
import project_tests as tests | |
# Check TensorFlow Version | |
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer. You are using {}'.format(tf.__version__) | |
print('TensorFlow Version: {}'.format(tf.__version__)) | |
# Check for a GPU | |
if not tf.test.gpu_device_name(): | |
warnings.warn('No GPU found. Please use a GPU to train your neural network.') | |
else: | |
print('Default GPU Device: {}'.format(tf.test.gpu_device_name())) | |
def load_vgg(sess, vgg_path): | |
""" | |
Load Pretrained VGG Model into TensorFlow. | |
:param sess: TensorFlow Session | |
:param vgg_path: Path to vgg folder, containing "variables/" and "saved_model.pb" | |
:return: Tuple of Tensors from VGG model (image_input, keep_prob, layer3_out, layer4_out, layer7_out) | |
""" | |
# DONE: Implement function | |
# Use tf.saved_model.loader.load to load the model and weights | |
vgg_tag = 'vgg16' | |
vgg_input_tensor_name = 'image_input:0' | |
vgg_keep_prob_tensor_name = 'keep_prob:0' | |
vgg_layer3_out_tensor_name = 'layer3_out:0' | |
vgg_layer4_out_tensor_name = 'layer4_out:0' | |
vgg_layer7_out_tensor_name = 'layer7_out:0' | |
tf.saved_model.loader.load(sess, [vgg_tag], vgg_path) | |
graph = tf.get_default_graph() | |
input_w = graph.get_tensor_by_name(vgg_input_tensor_name) | |
probabilities = graph.get_tensor_by_name(vgg_keep_prob_tensor_name) | |
layer3 = graph.get_tensor_by_name(vgg_layer3_out_tensor_name) | |
layer4 = graph.get_tensor_by_name(vgg_layer4_out_tensor_name) | |
layer7 = graph.get_tensor_by_name(vgg_layer7_out_tensor_name) | |
return input_w, probabilities, layer3, layer4, layer7 | |
tests.test_load_vgg(load_vgg, tf) | |
def layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes): | |
""" | |
Create the layers for a fully convolutional network. Build skip-layers using the vgg layers. | |
:param vgg_layer7_out: TF Tensor for VGG Layer 3 output | |
:param vgg_layer4_out: TF Tensor for VGG Layer 4 output | |
:param vgg_layer3_out: TF Tensor for VGG Layer 7 output | |
:param num_classes: Number of classes to classify | |
:return: The Tensor for the last layer of output | |
""" | |
# DONE: Implement function | |
# FCN16 | |
# divide output stride in half by predicting from 16 pixel stride layer | |
# add 1x1 convo layer on top of pool4 this adds additonal class predictions | |
# fuse this output with conv7 at stride 32 by adding 2x upsampling layer and summing both predictions | |
# initialize the 2x upsampling to bi-linear interpolations finally upsample stride 16 predictions to image size | |
# FCN32 | |
# pool4 layer params are zero-initialized | |
# FCN16 | |
# fusing predictions from pool3 with a 2x upsampling of predictions fused from pool4 and conv7 | |
# input, class or features so road or not road, kernel size 1 since it is only a 1x1 convolution, padding, and regularization | |
layer_7_1x1 = tf.layers.conv2d(vgg_layer7_out, num_classes, 1, padding='same', | |
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3)) | |
# layer_4_1x1 = tf.layers.conv2d(vgg_layer4_out, num_classes, 1, padding='same', | |
# kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3)) | |
#layer_3_1x1 = tf.layers.conv2d(vgg_layer3_out, num_classes, 1, padding='same', | |
# kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3)) | |
# upscale/ add features | |
output = tf.layers.conv2d_transpose(layer_7_1x1, 512, 4, strides=(2, 2), padding='same', | |
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3)) | |
# add layers together or "skip connections" | |
output = tf.add(output, vgg_layer4_out) | |
# upscale / reduce features | |
output = tf.layers.conv2d_transpose(output, 256, 4, strides=(2, 2), padding='same', | |
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3)) | |
# add layers together or "skip connections" | |
output = tf.add(output, vgg_layer3_out) | |
# upscale / reduce features | |
output = tf.layers.conv2d_transpose(output, num_classes, 16, strides=(8, 8), padding='same', | |
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3), name='my_output') | |
return output | |
tests.test_layers(layers) | |
def optimize(nn_last_layer, correct_label, learning_rate, num_classes): | |
""" | |
Build the TensorFLow loss and optimizer operations. | |
:param nn_last_layer: TF Tensor of the last layer in the neural network | |
:param correct_label: TF Placeholder for the correct label image | |
:param learning_rate: TF Placeholder for the learning rate | |
:param num_classes: Number of classes to classify | |
:return: Tuple of (logits, train_op, cross_entropy_loss) | |
""" | |
# TODO: Implement function | |
logits = tf.reshape(nn_last_layer, (-1, num_classes)) | |
# reg_ws = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES) | |
cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=correct_label)) | |
# cross_entropy_loss += tf.reduce_sum(reg_ws) | |
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate) | |
training_operation = optimizer.minimize(cross_entropy_loss) | |
return logits, training_operation, cross_entropy_loss | |
tests.test_optimize(optimize) | |
def train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, input_image, | |
correct_label, keep_prob, learning_rate): | |
""" | |
Train neural network and print out the loss during training. | |
:param sess: TF Session | |
:param epochs: Number of epochs | |
:param batch_size: Batch size | |
:param get_batches_fn: Function to get batches of training data. Call using get_batches_fn(batch_size) | |
:param train_op: TF Operation to train the neural network | |
:param cross_entropy_loss: TF Tensor for the amount of loss | |
:param input_image: TF Placeholder for input images | |
:param correct_label: TF Placeholder for label images | |
:param keep_prob: TF Placeholder for dropout keep probability | |
:param learning_rate: TF Placeholder for learning rate | |
""" | |
# DONE: Implement function | |
sess.run(tf.global_variables_initializer()) | |
for epoch in range(epochs): | |
for (image, label) in get_batches_fn(batch_size): | |
feed = {input_image: image, | |
correct_label: label, | |
keep_prob: 0.3, | |
learning_rate: 0.0003} | |
_, loss = sess.run([train_op, cross_entropy_loss], feed_dict=feed) | |
print("Loss: {} at Epoch {}/{}".format(loss, epoch, epochs)) | |
tests.test_train_nn(train_nn) | |
def run(): | |
num_classes = 2 | |
image_shape = (160, 576) | |
data_dir = './data' | |
runs_dir = './runs' | |
tests.test_for_kitti_dataset(data_dir) | |
# Download pretrained vgg model | |
helper.maybe_download_pretrained_vgg(data_dir) | |
# OPTIONAL: Train and Inference on the cityscapes dataset instead of the Kitti dataset. | |
# You'll need a GPU with at least 10 teraFLOPS to train on. | |
# https://www.cityscapes-dataset.com/ | |
epochs = 27 | |
batch_size = 7 | |
with tf.Session() as sess: | |
# Path to vgg model | |
vgg_path = os.path.join(data_dir, 'vgg') | |
# Create function to get batches | |
get_batches_fn = helper.gen_batch_function(os.path.join(data_dir, 'data_road/training'), image_shape) | |
# OPTIONAL: Augment Images for better results | |
# https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network | |
correct_label = tf.placeholder(tf.int32) | |
learning_rate = tf.placeholder(tf.float32) | |
# TODO: Build NN using load_vgg, layers, and optimize function | |
input_image, keep_prob, layer3_out, layer4_out, layer7_out = load_vgg(sess, vgg_path) | |
layer_output = layers(layer3_out, layer4_out, layer7_out, num_classes) | |
logits, train_op, cross_entropy_loss = optimize(layer_output, correct_label, learning_rate, num_classes) | |
# TODO: Train NN using the train_nn function | |
sess.run(tf.global_variables_initializer()) | |
train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, input_image, | |
correct_label, keep_prob, learning_rate) | |
for i in tf.get_default_graph().get_operations(): | |
print(i.name) | |
exit() | |
# TODO: Save inference data using helper.save_inference_samples | |
helper.save_inference_samples(runs_dir, data_dir, sess, image_shape, logits, keep_prob, input_image) | |
# OPTIONAL: Apply the trained model to a video | |
saver = tf.train.Saver() | |
save_path = os.path.join(runs_dir, 'model.ckpt') | |
saver.save(sess, save_path) | |
print('Saved at : {}'.format(save_path)) | |
freeze.freeze_graph(runs_dir, "my_output") | |
if __name__ == '__main__': | |
run() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment