-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
166 lines (136 loc) · 7.55 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import os.path
import tensorflow as tf
import helper
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)
"""
# TODO: Implement function #Do as mentioned in 5:36 - 7:58 of the video
# 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)
image_input = tf.get_default_graph().get_tensor_by_name(vgg_input_tensor_name)
keep_prob = tf.get_default_graph().get_tensor_by_name(vgg_keep_prob_tensor_name)
vgg_layer3_out = tf.get_default_graph().get_tensor_by_name(vgg_layer3_out_tensor_name)
vgg_layer4_out = tf.get_default_graph().get_tensor_by_name(vgg_layer4_out_tensor_name)
vgg_layer7_out = tf.get_default_graph().get_tensor_by_name(vgg_layer7_out_tensor_name)
return image_input, keep_prob, vgg_layer3_out, vgg_layer4_out, vgg_layer7_out
tests.test_load_vgg(load_vgg, tf)
def layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes):
# Concept explained at 9:30, implementation 10:41
"""
Create the layers for a fully convolutional network. Build skip-layers using the vgg layers.
:param vgg_layer3_out: TF Tensor for VGG Layer 3 output
:param vgg_layer4_out: TF Tensor for VGG Layer 4 output
:param vgg_layer7_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
"""
interm_kernels = 4
# TODO: Implement function
kinit= tf.random_normal_initializer(stddev=0.01)
kreg = tf.contrib.layers.l2_regularizer(1e-3)
# First have 1x1 convolution for layers 3 and 4 which will be added while upsampling
layer3_to_add = tf.layers.conv2d(vgg_layer3_out, num_classes, 1, padding='same', kernel_initializer=kinit, kernel_regularizer=kreg)
layer4_to_add = tf.layers.conv2d(vgg_layer4_out, interm_kernels, 1, padding='same', kernel_initializer=kinit, kernel_regularizer=kreg)
# Prepare layer 7
layer7_to_upsample = tf.layers.conv2d(vgg_layer7_out, interm_kernels, 1, padding='same', kernel_initializer=kinit, kernel_regularizer=kreg)
layer7_upsampled = tf.layers.conv2d_transpose(layer7_to_upsample, interm_kernels, 4, 2, padding='same', kernel_initializer=kinit, kernel_regularizer=kreg)
# Add layer4 info to upsampled layer7. Upsample again
layer4_to_upsample = tf.add(layer7_upsampled, layer4_to_add)
layer4_upsampled = tf.layers.conv2d_transpose(layer4_to_upsample, num_classes, 4, 2, padding='same', kernel_initializer=kinit, kernel_regularizer=kreg)
# Add layer3 info
layer3_to_upsample = tf.add(layer4_upsampled, layer3_to_add)
# Deconv to match the image size
nn_last_layer = tf.layers.conv2d_transpose(layer3_to_upsample, num_classes, 16, 8, padding='same', kernel_initializer=kinit, kernel_regularizer=kreg)
return nn_last_layer
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))
correct_label = tf.reshape(correct_label, (-1, num_classes))
cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=correct_label))
train_op = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy_loss)
return logits, train_op, 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
"""
# TODO: Implement function
sess.run(tf.global_variables_initializer())
print("Training:Start")
print("#####################")
for i in range(epochs):
print('Epoch ' + str(i))
for image, label in get_batches_fn(batch_size):
_, loss = sess.run([train_op, cross_entropy_loss], feed_dict={input_image: image, correct_label: label,
keep_prob: 0.7, learning_rate: 0.0005})
print("BatchLoss: = {:.3f}".format(loss))
pass
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)
epochs = 100
batch_size = 10
# 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/
correct_label = tf.placeholder(tf.int32, [None, image_shape[0], image_shape[1], num_classes], name='correct_label')
learning_rate = tf.placeholder(tf.float32, name='learning_rate')
vgg_path = os.path.join(data_dir, 'vgg')
get_batches_fn = helper.gen_batch_function(os.path.join(data_dir, 'data_road/training'), image_shape)
with tf.Session() as sess:
image_input, keep_prob, vgg_layer3_out, vgg_layer4_out, vgg_layer7_out = load_vgg(sess, vgg_path)
nn_last_layer = layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes)
logits, train_op, cross_entropy_loss = optimize(nn_last_layer, correct_label, learning_rate, num_classes)
# Invoke training
train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, image_input, correct_label, keep_prob, learning_rate)
# Save the results on images
helper.save_inference_samples(runs_dir, data_dir, sess, image_shape, logits, keep_prob, image_input)
if __name__ == '__main__':
run()