-
Notifications
You must be signed in to change notification settings - Fork 0
/
monodepth_main.py
314 lines (253 loc) · 16.3 KB
/
monodepth_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
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
# Copyright UCL Business plc 2017. Patent Pending. All rights reserved.
#
# The MonoDepth Software is licensed under the terms of the UCLB ACP-A licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.
#
# For any other use of the software not covered by the UCLB ACP-A Licence,
# please contact [email protected]
from __future__ import absolute_import, division, print_function
# only keep warnings and errors
import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2' #默认的显示等级,显示所有信息 1 -> 2
os.environ['CUDA_VISIBLE_DEVICES'] = '1' # 指定GPU训练
"""
os.environ["TF_CPP_MIN_LOG_LEVEL"]='2' # 只显示 warning 和 Error
os.environ["TF_CPP_MIN_LOG_LEVEL"]='3' # 只显示 Error
"""
import numpy as np
import argparse
import re
import time
import tensorflow as tf
import tensorflow.contrib.slim as slim #tensorflow辅助工具,用来简化代码
from tensorflow.python import debug as tf_debug
from monodepth_model import *
from monodepth_dataloader import *
from average_gradients import *
parser = argparse.ArgumentParser(description='Monodepth TensorFlow implementation.') # 参数的定义
# parser.add_argument('--mode', type=str, help='train or test', default='train')
# parser.add_argument('--model_name', type=str, help='model name', default='monodepth')
# parser.add_argument('--encoder', type=str, help='type of encoder, vgg or resnet50', default='vgg')
# parser.add_argument('--dataset', type=str, help='dataset to train on, kitti, or cityscapes', default='kitti')
# parser.add_argument('--data_path', type=str, help='path to the data', required=True)
# parser.add_argument('--filenames_file', type=str, help='path to the filenames text file', required=True)
# parser.add_argument('--input_height', type=int, help='input height', default=256)
# parser.add_argument('--input_width', type=int, help='input width', default=512)
# parser.add_argument('--batch_size', type=int, help='batch size', default=8) # 8->32
# parser.add_argument('--num_epochs', type=int, help='number of epochs', default=50)
# parser.add_argument('--learning_rate', type=float, help='initial learning rate', default=1e-4)
# parser.add_argument('--lr_loss_weight', type=float, help='left-right consistency weight', default=1.0)
# parser.add_argument('--alpha_image_loss', type=float, help='weight between SSIM and L1 in the image loss', default=0.85)
# parser.add_argument('--disp_gradient_loss_weight', type=float, help='disparity smoothness weigth', default=0.1)
# parser.add_argument('--do_stereo', help='if set, will train the stereo model', action='store_true')
# parser.add_argument('--wrap_mode', type=str, help='bilinear sampler wrap mode, edge or border', default='border')
# parser.add_argument('--use_deconv', help='if set, will use transposed convolutions', action='store_true')
# parser.add_argument('--num_gpus', type=int, help='number of GPUs to use for training', default=1) # 1->0
# parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=8) #8->16
# parser.add_argument('--output_directory', type=str, help='output directory for test disparities, if empty outputs to checkpoint folder', default='')
# parser.add_argument('--log_directory', type=str, help='directory to save checkpoints and summaries', default='')
# parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', default='')
# parser.add_argument('--retrain', help='if used with checkpoint_path, will restart training from step zero', action='store_true')
# parser.add_argument('--full_summary', help='if set, will keep more data for each summary. Warning: the file can become very large', action='store_true')
parser.add_argument('--mode', type=str, help='train or test', default='train')
parser.add_argument('--model_name', type=str, help='model name', default='train')
parser.add_argument('--encoder', type=str, help='type of encoder, vgg or resnet50', default='vgg')
parser.add_argument('--dataset', type=str, help='dataset to train on, kitti, or cityscapes', default='kitti')
parser.add_argument('--data_path', type=str, help='path to the data', default='/data/tion/kitti/')
parser.add_argument('--filenames_file', type=str, help='path to the filenames text file', default='/data/tion/my_data/all_npy_png.txt')
parser.add_argument('--input_height', type=int, help='input height', default=256)
parser.add_argument('--input_width', type=int, help='input width', default=512)
parser.add_argument('--batch_size', type=int, help='batch size', default=4)
parser.add_argument('--num_epochs', type=int, help='number of epochs', default=50)
parser.add_argument('--learning_rate', type=float, help='initial learning rate', default=1e-5)
parser.add_argument('--lr_loss_weight', type=float, help='left-right consistency weight', default=1.0)
parser.add_argument('--alpha_image_loss', type=float, help='weight between SSIM and L1 in the image loss', default=0.85)
parser.add_argument('--disp_gradient_loss_weight', type=float, help='disparity smoothness weigth', default=0.1)
parser.add_argument('--do_stereo', help='if set, will train the stereo model', action='store_true')
parser.add_argument('--wrap_mode', type=str, help='bilinear sampler wrap mode, edge or border', default='border')
parser.add_argument('--use_deconv', help='if set, will use transposed convolutions', action='store_true')
parser.add_argument('--num_gpus', type=int, help='number of GPUs to use for training', default=1)
parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=1) # 4->1
parser.add_argument('--output_directory', type=str, help='output directory for test disparities, if empty outputs to checkpoint folder', default='')
parser.add_argument('--log_directory', type=str, help='directory to save checkpoints and summaries', default='/data/tion/tmp/my_loss_vgg_kitti_lr_batch4/')
parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', default='')
parser.add_argument('--retrain', help='if used with checkpoint_path, will restart training from step zero', action='store_true')
parser.add_argument('--full_summary', help='if set, will keep more data for each summary. Warning: the file can become very large', action='store_true')
args = parser.parse_args()
def post_process_disparity(disp):
_, h, w = disp.shape
l_disp = disp[0,:,:]
r_disp = np.fliplr(disp[1,:,:])
m_disp = 0.5 * (l_disp + r_disp)
l, _ = np.meshgrid(np.linspace(0, 1, w), np.linspace(0, 1, h))
l_mask = 1.0 - np.clip(20 * (l - 0.05), 0, 1)
r_mask = np.fliplr(l_mask)
return r_mask * l_disp + l_mask * r_disp + (1.0 - l_mask - r_mask) * m_disp
def count_text_lines(file_path): # 计算文本文件行数 即训练数据量 一行存放一条训练数据
f = open(file_path, 'r')
lines = f.readlines()
f.close()
return len(lines) # 返回行数 txt 文件中是照片的路径一共 29000 行 每一行都是一个双目图像 两张图的路径
def train(params):
"""Training loop."""
with tf.Graph().as_default(), tf.device('/cpu:0'): # 创建graph
global_step = tf.Variable(0, trainable=False)
# OPTIMIZER
num_training_samples = count_text_lines(args.filenames_file)
steps_per_epoch = np.ceil(num_training_samples / params.batch_size).astype(np.int32)
num_total_steps = params.num_epochs * steps_per_epoch
start_learning_rate = args.learning_rate
# 根据步长改变学习率， 前3/5部分使用的是0.0001，后面每过1/5，将学习率除2
boundaries = [np.int32((3/5) * num_total_steps), np.int32((4/5) * num_total_steps)]
values = [args.learning_rate, args.learning_rate / 2, args.learning_rate / 4]
learning_rate = tf.train.piecewise_constant(global_step, boundaries, values) # 当走到一定步长时更改学习率
opt_step = tf.train.AdamOptimizer(learning_rate)
print("total number of samples: {}".format(num_training_samples))
print("total number of steps: {}".format(num_total_steps))
dataloader = MonodepthDataloader(args.data_path, args.filenames_file, params, args.dataset, args.mode)
left = dataloader.left_image_batch
right = dataloader.right_image_batch
# split for each gpu
left_splits = tf.split(left, args.num_gpus, 0)
right_splits = tf.split(right, args.num_gpus, 0)
# ==================== coordinates
left_coord = dataloader.left_coord_batch
right_coord = dataloader.right_coord_batch
# split for each gpu
left_coord_splits = tf.split(left_coord, args.num_gpus, 0)
right_coord_splits = tf.split(right_coord, args.num_gpus, 0)
# ====================
tower_grads = []
tower_losses = []
reuse_variables = None
with tf.variable_scope(tf.get_variable_scope()):
for i in range(args.num_gpus):
# with tf.device('/gpu:%d' % i):
with tf.device('/gpu:1'):
model = MonodepthModel(params, args.mode, left_splits[i], right_splits[i], left_coord_splits[i], right_coord_splits[i], reuse_variables, i)
# model = MonodepthModel(params, args.mode, left_splits[i], right_splits[i],reuse_variables, i)
loss = model.total_loss
tower_losses.append(loss)
reuse_variables = True
grads = opt_step.compute_gradients(loss)
tower_grads.append(grads)
grads = average_gradients(tower_grads)
apply_gradient_op = opt_step.apply_gradients(grads, global_step=global_step)
total_loss = tf.reduce_mean(tower_losses)
tf.summary.scalar('learning_rate', learning_rate, ['model_0'])
tf.summary.scalar('total_loss', total_loss, ['model_0'])
summary_op = tf.summary.merge_all('model_0')
# SESSION 通过tf.ConfigProto(allow_soft_placement=True) 来让Tensorflow自己选择可用设备
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# sess = tf_debug.LocalCLIDebugWrapperSession(sess)
# SAVER
summary_writer = tf.summary.FileWriter(args.log_directory + '/' + args.model_name, sess.graph)
train_saver = tf.train.Saver()
# COUNT PARAMS
total_num_parameters = 0
for variable in tf.trainable_variables():
total_num_parameters += np.array(variable.get_shape().as_list()).prod()
print("number of trainable parameters: {}".format(total_num_parameters))
# INIT
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
# LOAD CHECKPOINT IF SET
if args.checkpoint_path != '':
train_saver.restore(sess, args.checkpoint_path.split(".")[0])
if args.retrain:
sess.run(global_step.assign(0))
# GO!
start_step = global_step.eval(session=sess)
start_time = time.time()
for step in range(start_step, num_total_steps):
before_op_time = time.time()
_, loss_value, disp_gradient, lr_loss, image_loss, coord_image_loss, coord_lr_loss, coord_smooth_loss = sess.run(
[apply_gradient_op, total_loss, model.disp_gradient_loss, model.lr_loss,model.image_loss,
model.coord_image_loss, model.coord_lr_loss, model.coord_smoothness_loss
])
# _, loss_value= sess.run([apply_gradient_op, total_loss])
duration = time.time() - before_op_time
if step and step % 100 == 0: # 100 -> 200
# ============================= test
print('--------------------------------------')
print('disp_gradient lr_loss image_loss sum')
print('%.8f | %.8f | %.8f | %.8f' % (
disp_gradient * 0.1, lr_loss, image_loss, (disp_gradient * 0.1 + lr_loss + image_loss)))
print('%.8f | %.8f | %.8f | %.8f' % (
coord_smooth_loss * 0.1, coord_lr_loss, coord_image_loss,
(coord_smooth_loss * 0.1 + coord_lr_loss + coord_image_loss)))
# =============================
examples_per_sec = params.batch_size / duration
time_sofar = (time.time() - start_time) / 3600
training_time_left = (num_total_steps / step - 1.0) * time_sofar
print_string = 'batch {:>6} | examples/s: {:4.2f} | loss: {:.5f} | time elapsed: {:.2f}h | time left: {:.2f}h'
print(print_string.format(step, examples_per_sec, loss_value, time_sofar, training_time_left))
summary_str = sess.run(summary_op)
summary_writer.add_summary(summary_str, global_step=step)
if step and step % 10000 == 0:
train_saver.save(sess, args.log_directory + '/' + args.model_name + '/model', global_step=step)
train_saver.save(sess, args.log_directory + '/' + args.model_name + '/model', global_step=num_total_steps)
def test(params):
"""Test function."""
dataloader = MonodepthDataloader(args.data_path, args.filenames_file, params, args.dataset, args.mode)
left = dataloader.left_image_batch
right = dataloader.right_image_batch
model = MonodepthModel(params, args.mode, left, right)
# SESSION
config = tf.ConfigProto(allow_soft_placement=True)
sess = tf.Session(config=config)
# SAVER
train_saver = tf.train.Saver()
# INIT
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
# RESTORE
if args.checkpoint_path == '':
restore_path = tf.train.latest_checkpoint(args.log_directory + '/' + args.model_name)
else:
restore_path = args.checkpoint_path.split(".")[0]
train_saver.restore(sess, restore_path)
num_test_samples = count_text_lines(args.filenames_file)
print('now testing {} files'.format(num_test_samples))
disparities = np.zeros((num_test_samples, params.height, params.width), dtype=np.float32)
disparities_pp = np.zeros((num_test_samples, params.height, params.width), dtype=np.float32)
for step in range(num_test_samples):
disp = sess.run(model.disp_left_est[0])
disparities[step] = disp[0].squeeze()
disparities_pp[step] = post_process_disparity(disp.squeeze())
print('done.')
print('writing disparities.')
if args.output_directory == '':
output_directory = os.path.dirname(args.checkpoint_path)
else:
output_directory = args.output_directory
np.save(output_directory + '/disparities.npy', disparities)
np.save(output_directory + '/disparities_pp.npy', disparities_pp)
print('done.')
def main(_):
params = monodepth_parameters(
encoder=args.encoder,
height=args.input_height,
width=args.input_width,
batch_size=args.batch_size,
num_threads=args.num_threads,
num_epochs=args.num_epochs,
do_stereo=args.do_stereo,
wrap_mode=args.wrap_mode,
use_deconv=args.use_deconv,
alpha_image_loss=args.alpha_image_loss,
disp_gradient_loss_weight=args.disp_gradient_loss_weight,
lr_loss_weight=args.lr_loss_weight,
full_summary=args.full_summary)
if args.mode == 'train':
train(params)
elif args.mode == 'test':
test(params)
if __name__ == '__main__':
tf.app.run()