forked from lppllppl920/EndoscopyDepthEstimation-Pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain.py
494 lines (453 loc) · 28.4 KB
/
train.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
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
'''
Author: Xingtong Liu, Ayushi Sinha, Masaru Ishii, Gregory D. Hager, Austin Reiter, Russell H. Taylor, and Mathias Unberath
Copyright (C) 2019 Johns Hopkins University - All Rights Reserved
You may use, distribute and modify this code under the
terms of the GNU GENERAL PUBLIC LICENSE Version 3 license for non-commercial usage.
You should have received a copy of the GNU GENERAL PUBLIC LICENSE Version 3 license with
'''
import tqdm
import cv2
import numpy as np
from pathlib import Path
import torchsummary
import math
import torch
import random
from tensorboardX import SummaryWriter
import albumentations as albu
import argparse
import datetime
# Local
import models
import losses
import utils
import dataset
import scheduler
if __name__ == '__main__':
cv2.destroyAllWindows()
parser = argparse.ArgumentParser(
description='Self-supervised Depth Estimation on Monocular Endoscopy Dataset -- Train',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--adjacent_range', nargs='+', type=int, required=True,
help='interval range for a pair of video frames')
parser.add_argument('--id_range', nargs='+', type=int, required=True,
help='id range for the training and testing dataset')
parser.add_argument('--input_downsampling', type=float, default=4.0,
help='image downsampling rate')
parser.add_argument('--input_size', nargs='+', type=int, required=True, help='resolution of network input')
parser.add_argument('--batch_size', type=int, default=8, help='batch size for training and testing')
parser.add_argument('--num_workers', type=int, default=8, help='number of workers for input data loader')
parser.add_argument('--num_pre_workers', type=int, default=8,
help='number of workers for preprocessing intermediate data')
parser.add_argument('--dcl_weight', type=float, default=5.0,
help='weight for depth consistency loss in the later training stage')
parser.add_argument('--sfl_weight', type=float, default=20.0, help='weight for sparse flow loss')
parser.add_argument('--max_lr', type=float, default=1.0e-3, help='upper bound learning rate for cyclic lr')
parser.add_argument('--min_lr', type=float, default=1.0e-4, help='lower bound learning rate for cyclic lr')
parser.add_argument('--num_iter', type=int, default=1000, help='number of iterations per epoch')
parser.add_argument('--network_downsampling', type=int, default=64, help='network downsampling of the input image')
parser.add_argument('--inlier_percentage', type=float, default=0.99,
help='percentage of inliers of SfM point clouds (for pruning some outliers)')
parser.add_argument('--validation_interval', type=int, default=1, help='epoch interval for validation')
parser.add_argument('--zero_division_epsilon', type=float, default=1.0e-8, help='epsilon to prevent zero division')
parser.add_argument('--display_interval', type=int, default=10, help='iteration interval for image display')
parser.add_argument('--training_patient_id', nargs='+', type=int, required=True, help='id of the training patient')
parser.add_argument('--testing_patient_id', nargs='+', type=int, required=True, help='id of the testing patient')
parser.add_argument('--validation_patient_id', nargs='+', type=int, required=True,
help='id of the valiadtion patient')
parser.add_argument('--load_intermediate_data', action='store_true', help='whether to load intermediate data')
parser.add_argument('--load_trained_model', action='store_true',
help='whether to load trained student model')
parser.add_argument('--number_epoch', type=int, required=True, help='number of epochs in total')
parser.add_argument('--visibility_overlap', type=int, default=30, help='overlap of point visibility information')
parser.add_argument('--use_hsv_colorspace', action='store_true',
help='convert RGB to hsv colorspace')
parser.add_argument('--training_result_root', type=str, required=True, help='root of the training input and ouput')
parser.add_argument('--training_data_root', type=str, required=True, help='path to the training data')
parser.add_argument('--architecture_summary', action='store_true', help='display the network architecture')
parser.add_argument('--trained_model_path', type=str, default=None,
help='path to the trained student model')
args = parser.parse_args()
# Fix randomness for reproducibility
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
torch.manual_seed(10085)
np.random.seed(10085)
random.seed(10085)
# Hyper-parameters
adjacent_range = args.adjacent_range
input_downsampling = args.input_downsampling
height, width = args.input_size
batch_size = args.batch_size
num_workers = args.num_workers
num_pre_workers = args.num_pre_workers
depth_consistency_weight = args.dcl_weight
sparse_flow_weight = args.sfl_weight
max_lr = args.max_lr
min_lr = args.min_lr
num_iter = args.num_iter
network_downsampling = args.network_downsampling
inlier_percentage = args.inlier_percentage
validation_each = args.validation_interval
depth_scaling_epsilon = args.zero_division_epsilon
depth_warping_epsilon = args.zero_division_epsilon
wsl_epsilon = args.zero_division_epsilon
display_each = args.display_interval
training_patient_id = args.training_patient_id
testing_patient_id = args.testing_patient_id
validation_patient_id = args.validation_patient_id
load_intermediate_data = args.load_intermediate_data
load_trained_model = args.load_trained_model
n_epochs = args.number_epoch
is_hsv = args.use_hsv_colorspace
training_result_root = args.training_result_root
display_architecture = args.architecture_summary
trained_model_path = args.trained_model_path
training_data_root = Path(args.training_data_root)
id_range = args.id_range
visibility_overlap = args.visibility_overlap
currentDT = datetime.datetime.now()
depth_estimation_model_teacher = []
failure_sequences = []
training_transforms = albu.Compose([
# Color augmentation
albu.OneOf([
albu.Compose([
albu.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
albu.RandomGamma(gamma_limit=(80, 120), p=0.5),
albu.HueSaturationValue(hue_shift_limit=30, sat_shift_limit=0, val_shift_limit=0, p=0.5)]),
albu.HueSaturationValue(hue_shift_limit=30, sat_shift_limit=30, val_shift_limit=30, p=0.5)
]),
# Image quality augmentation
albu.OneOf([
albu.Blur(p=0.5),
albu.MedianBlur(p=0.5),
albu.MotionBlur(p=0.5),
albu.JpegCompression(quality_lower=20, quality_upper=100, p=0.5)
]),
# Noise augmentation
albu.OneOf([
albu.GaussNoise(var_limit=(10, 30), p=0.5),
albu.IAAAdditiveGaussianNoise(loc=0, scale=(0.005 * 255, 0.02 * 255), p=0.5)
]),
], p=1.)
log_root = Path(training_result_root) / "depth_estimation_train_run_{}_{}_{}_{}_test_id_{}".format(
currentDT.month,
currentDT.day,
currentDT.hour,
currentDT.minute,
"_".join(testing_patient_id))
if not log_root.exists():
log_root.mkdir()
writer = SummaryWriter(logdir=str(log_root))
print("Tensorboard visualization at {}".format(str(log_root)))
# Get color image filenames
train_filenames, val_filenames, test_filenames = utils.get_color_file_names_by_bag(training_data_root,
training_patient_id=training_patient_id,
validation_patient_id=validation_patient_id,
testing_patient_id=testing_patient_id)
folder_list = utils.get_parent_folder_names(training_data_root, id_range=id_range)
# Build training and validation dataset
train_dataset = dataset.SfMDataset(image_file_names=train_filenames,
folder_list=folder_list,
adjacent_range=adjacent_range, transform=training_transforms,
downsampling=input_downsampling,
network_downsampling=network_downsampling, inlier_percentage=inlier_percentage,
use_store_data=load_intermediate_data,
store_data_root=training_data_root,
phase="train", is_hsv=is_hsv,
num_pre_workers=num_pre_workers, visible_interval=visibility_overlap,
rgb_mode="rgb", num_iter=num_iter)
validation_dataset = dataset.SfMDataset(image_file_names=val_filenames,
folder_list=folder_list,
adjacent_range=adjacent_range,
transform=None,
downsampling=input_downsampling,
network_downsampling=network_downsampling,
inlier_percentage=inlier_percentage,
use_store_data=True,
store_data_root=training_data_root,
phase="validation", is_hsv=is_hsv,
num_pre_workers=num_pre_workers, visible_interval=visibility_overlap,
rgb_mode="rgb", num_iter=None)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
num_workers=num_workers)
validation_loader = torch.utils.data.DataLoader(dataset=validation_dataset, batch_size=batch_size, shuffle=False,
num_workers=batch_size)
depth_estimation_model = models.FCDenseNet57(n_classes=1)
# Initialize the depth estimation network with Kaiming He initialization
depth_estimation_model = utils.init_net(depth_estimation_model, type="kaiming", mode="fan_in",
activation_mode="relu",
distribution="normal")
# Multi-GPU running
depth_estimation_model = torch.nn.DataParallel(depth_estimation_model)
# Summary network architecture
if display_architecture:
torchsummary.summary(depth_estimation_model, input_size=(3, height, width))
# Optimizer
optimizer = torch.optim.SGD(depth_estimation_model.parameters(), lr=max_lr, momentum=0.9)
lr_scheduler = scheduler.CyclicLR(optimizer, base_lr=min_lr, max_lr=max_lr, step_size=num_iter)
# Custom layers
depth_scaling_layer = models.DepthScalingLayer(epsilon=depth_scaling_epsilon)
depth_warping_layer = models.DepthWarpingLayer(epsilon=depth_warping_epsilon)
flow_from_depth_layer = models.FlowfromDepthLayer()
# Loss functions
sparse_flow_loss_function = losses.SparseMaskedL1Loss()
depth_consistency_loss_function = losses.NormalizedDistanceLoss(height=height, width=width)
# Load previous student model, lr scheduler, and so on
if load_trained_model:
if Path(trained_model_path).exists():
print("Loading {:s} ...".format(trained_model_path))
state = torch.load(trained_model_path)
step = state['step']
epoch = state['epoch']
depth_estimation_model.load_state_dict(state['model'])
print('Restored model, epoch {}, step {}'.format(epoch, step))
else:
print("No trained model detected")
raise OSError
else:
epoch = 0
step = 0
for epoch in range(epoch, n_epochs + 1):
# Set the seed correlated to epoch for reproducibility
torch.manual_seed(10086 + epoch)
np.random.seed(10086 + epoch)
random.seed(10086 + epoch)
depth_estimation_model.train()
# Update progress bar
tq = tqdm.tqdm(total=len(train_loader) * batch_size, dynamic_ncols=True, ncols=40)
# Variable initialization
if epoch <= 20:
depth_consistency_weight = 0.1
else:
depth_consistency_weight = args.dcl_weight
for batch, (
colors_1, colors_2, sparse_depths_1, sparse_depths_2, sparse_depth_masks_1, sparse_depth_masks_2,
sparse_flows_1, sparse_flows_2, sparse_flow_masks_1, sparse_flow_masks_2, boundaries, rotations_1_wrt_2,
rotations_2_wrt_1, translations_1_wrt_2, translations_2_wrt_1, intrinsics, folders, file_names) in \
enumerate(train_loader):
# Update learning rate
lr_scheduler.batch_step(batch_iteration=step)
tq.set_description('Epoch {}, lr {}'.format(epoch, lr_scheduler.get_lr()))
with torch.no_grad():
colors_1 = colors_1.cuda()
colors_2 = colors_2.cuda()
sparse_depths_1 = sparse_depths_1.cuda()
sparse_depths_2 = sparse_depths_2.cuda()
sparse_depth_masks_1 = sparse_depth_masks_1.cuda()
sparse_depth_masks_2 = sparse_depth_masks_2.cuda()
sparse_flows_1 = sparse_flows_1.cuda()
sparse_flows_2 = sparse_flows_2.cuda()
sparse_flow_masks_1 = sparse_flow_masks_1.cuda()
sparse_flow_masks_2 = sparse_flow_masks_2.cuda()
boundaries = boundaries.cuda()
rotations_1_wrt_2 = rotations_1_wrt_2.cuda()
rotations_2_wrt_1 = rotations_2_wrt_1.cuda()
translations_1_wrt_2 = translations_1_wrt_2.cuda()
translations_2_wrt_1 = translations_2_wrt_1.cuda()
intrinsics = intrinsics.cuda()
colors_1 = boundaries * colors_1
colors_2 = boundaries * colors_2
# Predicted depth from student model
predicted_depth_maps_1 = depth_estimation_model(colors_1)
predicted_depth_maps_2 = depth_estimation_model(colors_2)
scaled_depth_maps_1, normalized_scale_std_1 = depth_scaling_layer(
[predicted_depth_maps_1, sparse_depths_1, sparse_depth_masks_1])
scaled_depth_maps_2, normalized_scale_std_2 = depth_scaling_layer(
[predicted_depth_maps_2, sparse_depths_2, sparse_depth_masks_2])
# Sparse flow loss
# Flow maps calculated using predicted dense depth maps and camera poses
# should agree with the sparse flows of feature points from SfM
flows_from_depth_1 = flow_from_depth_layer(
[scaled_depth_maps_1, boundaries, translations_1_wrt_2, rotations_1_wrt_2,
intrinsics])
flows_from_depth_2 = flow_from_depth_layer(
[scaled_depth_maps_2, boundaries, translations_2_wrt_1, rotations_2_wrt_1,
intrinsics])
sparse_flow_masks_1 = sparse_flow_masks_1 * boundaries
sparse_flow_masks_2 = sparse_flow_masks_2 * boundaries
sparse_flows_1 = sparse_flows_1 * boundaries
sparse_flows_2 = sparse_flows_2 * boundaries
flows_from_depth_1 = flows_from_depth_1 * boundaries
flows_from_depth_2 = flows_from_depth_2 * boundaries
sparse_flow_loss = sparse_flow_weight * 0.5 * (sparse_flow_loss_function(
[sparse_flows_1, flows_from_depth_1, sparse_flow_masks_1]) + sparse_flow_loss_function(
[sparse_flows_2, flows_from_depth_2, sparse_flow_masks_2]))
# Depth consistency loss
warped_depth_maps_2_to_1, intersect_masks_1 = depth_warping_layer(
[scaled_depth_maps_1, scaled_depth_maps_2, boundaries, translations_1_wrt_2, rotations_1_wrt_2,
intrinsics])
warped_depth_maps_1_to_2, intersect_masks_2 = depth_warping_layer(
[scaled_depth_maps_2, scaled_depth_maps_1, boundaries, translations_2_wrt_1, rotations_2_wrt_1,
intrinsics])
depth_consistency_loss = depth_consistency_weight * 0.5 * (depth_consistency_loss_function(
[scaled_depth_maps_1, warped_depth_maps_2_to_1, intersect_masks_1,
intrinsics]) + depth_consistency_loss_function(
[scaled_depth_maps_2, warped_depth_maps_1_to_2, intersect_masks_2, intrinsics]))
loss = depth_consistency_loss + sparse_flow_loss
if math.isnan(loss.item()) or math.isinf(loss.item()):
optimizer.zero_grad()
loss.backward()
optimizer.zero_grad()
optimizer.step()
continue
else:
optimizer.zero_grad()
loss.backward()
# Prevent one sample from having too much impact on the training
torch.nn.utils.clip_grad_norm_(depth_estimation_model.parameters(), 10.0)
optimizer.step()
if batch == 0:
mean_loss = loss.item()
mean_depth_consistency_loss = depth_consistency_loss.item()
mean_sparse_flow_loss = sparse_flow_loss.item()
else:
mean_loss = (mean_loss * batch + loss.item()) / (batch + 1.0)
mean_depth_consistency_loss = (mean_depth_consistency_loss * batch +
depth_consistency_loss.item()) / (batch + 1.0)
mean_sparse_flow_loss = (mean_sparse_flow_loss * batch + sparse_flow_loss.item()) / (batch + 1.0)
step += 1
tq.update(batch_size)
tq.set_postfix(loss='avg: {:.5f} cur: {:.5f}'.format(mean_loss, loss.item()),
loss_depth_consistency='avg: {:.5f} cur: {:.5f}'.format(
mean_depth_consistency_loss,
depth_consistency_loss.item()),
loss_sparse_flow='avg: {:.5f} cur: {:.5f}'.format(
mean_sparse_flow_loss,
sparse_flow_loss.item()))
writer.add_scalars('Training', {'overall': mean_loss,
'depth_consistency': mean_depth_consistency_loss,
'sparse_flow': mean_sparse_flow_loss}, step)
# Display depth and color at TensorboardX
if batch % display_each == 0:
colors_1_display, pred_depths_1_display, sparse_flows_1_display, dense_flows_1_display = \
utils.display_color_depth_sparse_flow_dense_flow(1, step, writer, colors_1,
scaled_depth_maps_1 * boundaries,
sparse_flows_1, flows_from_depth_1, is_hsv,
phase="Training", is_return_image=True,
color_reverse=True)
colors_2_display, pred_depths_2_display, sparse_flows_2_display, dense_flows_2_display = \
utils.display_color_depth_sparse_flow_dense_flow(2, step, writer, colors_2,
scaled_depth_maps_2 * boundaries,
sparse_flows_2, flows_from_depth_2, is_hsv,
phase="Training", is_return_image=True,
color_reverse=True)
utils.stack_and_display(phase="Training", title="Results (c1, d1, sf1, df1, c2, d2, sf2, df2)",
step=step, writer=writer,
image_list=[colors_1_display, pred_depths_1_display, sparse_flows_1_display,
dense_flows_1_display,
colors_2_display, pred_depths_2_display, sparse_flows_2_display,
dense_flows_2_display])
tq.close()
# Save student model
if epoch % validation_each != 0:
continue
tq = tqdm.tqdm(total=len(validation_loader) * batch_size, dynamic_ncols=True, ncols=40)
tq.set_description('Validation Epoch {}'.format(epoch))
with torch.no_grad():
for batch, (
colors_1, colors_2, sparse_depths_1, sparse_depths_2, sparse_depth_masks_1,
sparse_depth_masks_2, sparse_flows_1,
sparse_flows_2, sparse_flow_masks_1, sparse_flow_masks_2, boundaries, rotations_1_wrt_2,
rotations_2_wrt_1, translations_1_wrt_2, translations_2_wrt_1, intrinsics,
folders, file_names) in enumerate(validation_loader):
colors_1 = colors_1.cuda()
colors_2 = colors_2.cuda()
sparse_depths_1 = sparse_depths_1.cuda()
sparse_depths_2 = sparse_depths_2.cuda()
sparse_depth_masks_1 = sparse_depth_masks_1.cuda()
sparse_depth_masks_2 = sparse_depth_masks_2.cuda()
sparse_flows_1 = sparse_flows_1.cuda()
sparse_flows_2 = sparse_flows_2.cuda()
sparse_flow_masks_1 = sparse_flow_masks_1.cuda()
sparse_flow_masks_2 = sparse_flow_masks_2.cuda()
boundaries = boundaries.cuda()
rotations_1_wrt_2 = rotations_1_wrt_2.cuda()
rotations_2_wrt_1 = rotations_2_wrt_1.cuda()
translations_1_wrt_2 = translations_1_wrt_2.cuda()
translations_2_wrt_1 = translations_2_wrt_1.cuda()
intrinsics = intrinsics.cuda()
colors_1 = boundaries * colors_1
colors_2 = boundaries * colors_2
# Predicted depth from student model
predicted_depth_maps_1 = depth_estimation_model(colors_1)
predicted_depth_maps_2 = depth_estimation_model(colors_2)
scaled_depth_maps_1, normalized_scale_std_1 = depth_scaling_layer(
[torch.abs(predicted_depth_maps_1), sparse_depths_1, sparse_depth_masks_1])
scaled_depth_maps_2, normalized_scale_std_2 = depth_scaling_layer(
[torch.abs(predicted_depth_maps_2), sparse_depths_2, sparse_depth_masks_2])
# Sparse flow loss
flows_from_depth_1 = flow_from_depth_layer(
[scaled_depth_maps_1, boundaries, translations_1_wrt_2, rotations_1_wrt_2,
intrinsics])
flows_from_depth_2 = flow_from_depth_layer(
[scaled_depth_maps_2, boundaries, translations_2_wrt_1, rotations_2_wrt_1,
intrinsics])
sparse_flow_masks_1 = sparse_flow_masks_1 * boundaries
sparse_flow_masks_2 = sparse_flow_masks_2 * boundaries
sparse_flows_1 = sparse_flows_1 * boundaries
sparse_flows_2 = sparse_flows_2 * boundaries
flows_from_depth_1 = flows_from_depth_1 * boundaries
flows_from_depth_2 = flows_from_depth_2 * boundaries
sparse_flow_loss = sparse_flow_weight * 0.5 * (sparse_flow_loss_function(
[sparse_flows_1, flows_from_depth_1, sparse_flow_masks_1]) + sparse_flow_loss_function(
[sparse_flows_2, flows_from_depth_2, sparse_flow_masks_2]))
# Depth consistency loss
warped_depth_maps_2_to_1, intersect_masks_1 = depth_warping_layer(
[scaled_depth_maps_1, scaled_depth_maps_2, boundaries, translations_1_wrt_2, rotations_1_wrt_2,
intrinsics])
warped_depth_maps_1_to_2, intersect_masks_2 = depth_warping_layer(
[scaled_depth_maps_2, scaled_depth_maps_1, boundaries, translations_2_wrt_1, rotations_2_wrt_1,
intrinsics])
depth_consistency_loss = depth_consistency_weight * 0.5 * (depth_consistency_loss_function(
[scaled_depth_maps_1, warped_depth_maps_2_to_1,
intersect_masks_1, intrinsics]) + depth_consistency_loss_function(
[scaled_depth_maps_2, warped_depth_maps_1_to_2, intersect_masks_2, intrinsics]))
loss = depth_consistency_loss + sparse_flow_loss
tq.update(batch_size)
if not np.isnan(loss.item()):
if batch == 0:
mean_loss = loss.item()
mean_depth_consistency_loss = depth_consistency_loss.item()
mean_sparse_flow_loss = sparse_flow_loss.item()
else:
mean_loss = (mean_loss * batch + loss.item()) / (batch + 1.0)
mean_depth_consistency_loss = (mean_depth_consistency_loss * batch +
depth_consistency_loss.item()) / (batch + 1.0)
mean_sparse_flow_loss = (mean_sparse_flow_loss * batch + sparse_flow_loss.item()) / (
batch + 1.0)
# Display depth and color at TensorboardX
if batch % display_each == 0:
colors_1_display, pred_depths_1_display, sparse_flows_1_display, dense_flows_1_display = \
utils.display_color_depth_sparse_flow_dense_flow(1, step, writer, colors_1,
scaled_depth_maps_1 * boundaries,
sparse_flows_1, flows_from_depth_1, is_hsv,
phase="Validation", is_return_image=True,
color_reverse=True)
colors_2_display, pred_depths_2_display, sparse_flows_2_display, dense_flows_2_display = \
utils.display_color_depth_sparse_flow_dense_flow(2, step, writer, colors_2,
scaled_depth_maps_2 * boundaries,
sparse_flows_2, flows_from_depth_2, is_hsv,
phase="Validation", is_return_image=True,
color_reverse=True)
utils.stack_and_display(phase="Validation", title="Results (c1, d1, sf1, df1, c2, d2, sf2, df2)",
step=step, writer=writer,
image_list=[colors_1_display, pred_depths_1_display, sparse_flows_1_display,
dense_flows_1_display,
colors_2_display, pred_depths_2_display, sparse_flows_2_display,
dense_flows_2_display])
# TensorboardX
writer.add_scalars('Validation', {'overall': mean_loss,
'depth_consistency': mean_depth_consistency_loss,
'sparse_flow': mean_sparse_flow_loss}, epoch)
tq.close()
model_path_epoch = log_root / 'checkpoint_model_epoch_{}_validation_{}.pt'.format(epoch,
mean_sparse_flow_loss)
utils.save_model(model=depth_estimation_model, optimizer=optimizer,
epoch=epoch + 1, step=step, model_path=model_path_epoch,
validation_loss=mean_sparse_flow_loss)
writer.export_scalars_to_json(
str(log_root / ("all_scalars_" + str(epoch) + ".json")))
writer.close()