def run(self):
self.model.eval()
with torch.no_grad():
for batch_idx, data in enumerate(self.dataloader):
bg, fgbg, mask, depth = data['bg'].to(
self.model.device), data['fgbg'].to(
self.model.device), data['mask'].to(
self.model.device), data['depth'].to(
self.model.device)
mask_pred, depth_pred = self.model(bg, fgbg)
# Calculate loss
if self.criterion1 is not None:
loss1 = self.criterion1(mask_pred, mask)
m_ssim = torch.clamp((1 - ssim(mask_pred, mask)) * 0.5, 0, 1)
loss1 = (0.84 * m_ssim) + (0.16 * loss1)
if self.criterion2 is not None:
loss2 = self.criterion2(depth_pred, depth)
d_ssim = torch.clamp((1 - ssim(depth_pred, depth)) * 0.5, 0, 1)
loss2 = (0.84 * d_ssim) + (0.16 * loss2)
self.loss = 2 * loss1 + loss2
if batch_idx == 0:
inp = fgbg.detach().cpu()
orimp = mask.detach().cpu()
mp = mask_pred.detach().cpu()
oridp = depth.detach().cpu()
dp = depth_pred.detach().cpu()
print(
"First batch in testing fgbg, (mask, predicted mask), (depth, predicted depth)"
)
show(inp[:8, :, :, :], normalize=True)
mdinp = torch.cat([
orimp[:8, :, :, :], mp[:8, :, :, :],
oridp[:8, :, :, :], dp[:8, :, :, :]
],
dim=0)
show(mdinp)
self.stats.add_batch_test_stats(self.loss.item(), 0, len(data))
if self.scheduler and isinstance(
self.scheduler,
torch.optim.lr_scheduler.ReduceLROnPlateau):
#print("hello yes i am ")
self.scheduler.step(self.loss)
def validate(loader, model, epoch, d1, d2, blind, noise_level):
val_psnr = 0
val_ssim = 0
val_l1 = 0
model.train(False)
k1 = model.weight[0].unsqueeze(0).expand(loader.batch_size, -1, -1, -1)
k2 = model.weight[1].unsqueeze(0).expand(loader.batch_size, -1, -1, -1)
d1 = d1.expand(loader.batch_size, -1, -1, -1)
d2 = d2.expand(loader.batch_size, -1, -1, -1)
# pre-create noise levels
if blind:
nls = np.linspace(0.5, noise_level, len(loader))
else:
nls = noise_level * np.ones(len(loader))
with torch.no_grad():
for i, data in tqdm.tqdm(enumerate(loader)):
x, y, k, d = data
x = x.to(device)
y = y.to(device)
k = k.to(device)
d = d.to(device)
nl = nls[i] / 255
y += nl * torch.randn_like(y)
y = y.clamp(0, 1)
hat_x = model(y, k, d, k1, k2, d1, d2)[-1]
hat_x.clamp_(0, 1)
hat_x = utils.crop_valid(hat_x, k)
x = utils.crop_valid(x, k)
y = utils.crop_valid(y, k)
val_psnr += loss.psnr(hat_x, x)
val_ssim += loss.ssim(hat_x, x)
val_l1 += F.l1_loss(hat_x, x).item()
return val_psnr / len(loader), val_ssim / len(loader), val_l1 / len(loader)
def validate(loader, model, epoch, d1, d2, blind, noise_level):
val_psnr = 0
val_ssim = 0
val_l1 = 0
model.train(False)
k1 = model.weight[0].unsqueeze(0).expand(loader.batch_size, -1, -1, -1)
k2 = model.weight[1].unsqueeze(0).expand(loader.batch_size, -1, -1, -1)
d1 = d1.expand(loader.batch_size, -1, -1, -1)
d2 = d2.expand(loader.batch_size, -1, -1, -1)
# pre-create noise levels
if blind:
nls = np.linspace(0.5, noise_level, len(loader))
else:
nls = noise_level * np.ones(len(loader))
with torch.no_grad():
for i, data in tqdm.tqdm(enumerate(loader)):
x, y, mag, ori = data
x = x.to(device)
y = y.to(device)
mag = mag.to(device)
ori = ori.to(device)
ori = (90 - ori).add(360).fmod(180)
labels = utils.get_labels(mag, ori)
ori = ori * np.pi / 180
nl = nls[i] / 255
y += nl * torch.randn_like(y)
y = y.clamp(0, 1)
hat_x = model(y, mag, ori, labels, k1, k2, d1, d2)[-1]
hat_x.clamp_(0, 1)
val_psnr += loss.psnr(hat_x, x)
val_ssim += loss.ssim(hat_x, x)
val_l1 += F.l1_loss(hat_x, x).item()
return val_psnr / len(loader), val_ssim / len(loader), val_l1 / len(loader)
def eval(self):
self.model.eval()
psnr_losses = []
ssim_losses = []
with tqdm(total=len(self.valloader)) as t:
# show batch evaluate process
t.set_description('evaluating...')
for LR, HR in self.valloader:
LR = LR.to(self.device)
HR = HR.to(self.device)
output = self.model(LR)
# # only calculation on y channel or gray
# if self.r_mode == 'RGB' and self.img_channels == 3:
# output = rgb2y_tensor(output)
# HR = rgb2y_tensor(HR)
ssim_loss = ssim(output, HR)
psnr_loss = psnr(output, HR)
# save losses
ssim_losses.append(ssim_loss)
psnr_losses.append(psnr_loss)
t.update()
avg_ssim = torch.stack(ssim_losses, dim=0).mean().item()
avg_psnr = torch.stack(psnr_losses, dim=0).mean().item()
t.set_postfix(avg_psnr=f'{avg_psnr:.010f}',
avg_ssim=f'{avg_ssim:.010f}')
t.set_description('evaluate')
self.model.train()
return avg_ssim, avg_psnr
def LogProgress(model, writer, test_loader, global_step):
with torch.no_grad():
model.eval()
sequential = test_loader
_input, _label = next(iter(sequential))
image = torch.autograd.Variable(_input.float().cuda())
depth = torch.autograd.Variable(_label.float().cuda(non_blocking=True))
# Normalize depth
# depth_n = DepthNorm(depth)
# print()
# if epoch == 0:
# writer.add_image('Train.1.Image', vutils.make_grid(image.data, nrow=6, normalize=True), epoch)
# if epoch == 0:
output = model(image)
# Predict
# output = model(image)
# edge = sobel_(output)
# pred_edge = sobel_(depth)
# Compute the loss
# l_sobel = nn.L1Loss()(edge, pred_edge)
l_depth = nn.L1Loss()(output, depth)
l_ssim = torch.clamp((1 - ssim(output, depth, val_range=10.0)) * 0.5,
0, 1)
# if torch.isnan(l_sobel):
# print(1)
# print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim))
writer.add_scalar('Test/L1', l_depth.item(), global_step)
writer.add_scalar('Test/SSIM', l_ssim.item(), global_step)
# writer.add_scalar('Test/EDGE', l_sobel.item(), global_step)
# normal = Norm(output)
# normal = (normal + 1.0) / 2.
# normal = normal[3, :, :, :]
# normal = normal.detach().cpu().numpy().astype(np.uint8)
# normal = np.transpose(normal, (1, 2, 0))
# print()
# plt.imshow(normal)
# plt.show()
# output = DepthNorm(output)
# writer.add_image('Train.3.Normal', vutils.make_grid(normal.data, nrow=6, normalize=False), epoch)
_fig = colorize(image.data, depth.data, output.data)
# _ = vutils.make_grid(_, nrow=1, padding=0, normalize=False)
# res1 = torch.cat([image[0:1, 3:4].data, depth[0:1].data, output[0:1].data], dim=0)
# xx = colorize(vutils.make_grid(res1, nrow=6, padding=0, normalize=False))
# rgb = image[0:1, :3, :, :].data.cpu().numpy() * 255.
# rgb = rgb.astype(np.uint8).squeeze()
# out = np.concatenate([rgb, xx], axis=2)
# print()
writer.add_image('Train.2.test', _fig.transpose((2, 0, 1)),
global_step)
# writer.add_image('Train.2.Raw', colorize(vutils.make_grid(image[:,3:4,:,:].data, nrow=6, normalize=False)),
# global_step)
# writer.add_image('Train.2.Depth', colorize(vutils.make_grid(depth.data, nrow=6, normalize=False)), global_step)
# writer.add_image('Test.2.Ours', colorize(vutils.make_grid(output.data, nrow=6, normalize=False)), global_step)
# writer.add_image('Test.3.Diff', colorize(vutils.make_grid(torch.abs(output - depth).data, nrow=6, normalize=False)),
# global_step)
del image
del depth
del output
def run_epoch(self, epoch, dataloader, logimage=False, isTrain=True):
# For details see training.
psnr_value = 0
ssim_value = 0
loss_value = 0
if not isTrain:
valid_images = []
for index, all_data in enumerate(dataloader, 0):
self.optimizer.zero_grad()
(
Ft_p,
I0,
IFrame,
I1,
g_I0_F_t_0,
g_I1_F_t_1,
FlowBackWarp_I0_F_1_0,
FlowBackWarp_I1_F_0_1,
F_1_0,
F_0_1,
) = self.slomo(all_data, pred_only=False, isTrain=isTrain)
if (not isTrain) and logimage:
if index % self.args.logimagefreq == 0:
valid_images.append(
255.0
* I0.cpu()[0]
.resize_(1, 1, self.args.data_h, self.args.data_w)
.repeat(1, 3, 1, 1)
)
valid_images.append(
255.0
* IFrame.cpu()[0]
.resize_(1, 1, self.args.data_h, self.args.data_w)
.repeat(1, 3, 1, 1)
)
valid_images.append(
255.0
* I1.cpu()[0]
.resize_(1, 1, self.args.data_h, self.args.data_w)
.repeat(1, 3, 1, 1)
)
valid_images.append(
255.0
* Ft_p.cpu()[0]
.resize_(1, 1, self.args.data_h, self.args.data_w)
.repeat(1, 3, 1, 1)
)
# loss
loss = self.supervisedloss(
Ft_p,
IFrame,
I0,
I1,
g_I0_F_t_0,
g_I1_F_t_1,
FlowBackWarp_I0_F_1_0,
FlowBackWarp_I1_F_0_1,
F_1_0,
F_0_1,
)
if isTrain:
loss.backward()
self.optimizer.step()
self.scheduler.step()
loss_value += loss.item()
# metrics
psnr_value += psnr(Ft_p, IFrame, outputTensor=False)
ssim_value += ssim(Ft_p, IFrame, outputTensor=False)
name_loss = "TrainLoss" if isTrain else "ValLoss"
itr = int(index + epoch * (len(dataloader)))
if self.comet_exp is not None:
self.comet_exp.log_metric(
"PSNR", psnr_value / len(dataloader), step=itr, epoch=epoch
)
self.comet_exp.log_metric(
"SSIM", ssim_value / len(dataloader), step=itr, epoch=epoch
)
self.comet_exp.log_metric(
name_loss, loss_value / len(dataloader), step=itr, epoch=epoch
)
if logimage:
upload_images(
valid_images,
epoch,
exp=self.comet_exp,
im_per_row=4,
rows_per_log=int(len(valid_images) / 4),
)
print(
" Loss: %0.6f Iterations: %4d/%4d ValPSNR: %0.4f ValSSIM: %0.4f "
% (
loss_value / len(dataloader),
index,
len(dataloader),
psnr_value / len(dataloader),
ssim_value / len(dataloader),
)
)
return (
(psnr_value / len(dataloader)),
(ssim_value / len(dataloader)),
(loss_value / len(dataloader)),
)
def main():
# Arguments
parser = argparse.ArgumentParser(description='High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs', default=20, type=int, help='number of total epochs to run')
parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
SAVE_DIR = 'models/191107_mod15'
ifcrop = True
if ifcrop:
HEIGHT = 256
HEIGHT_WITH_RATIO = 240
WIDTH = 320
else:
HEIGHT = 480
WIDTH = 640
with torch.cuda.device(0):
# Create model
# model = Model().cuda()
# print('Model created.')
# =============================================================================
# load saved model
# model = Model_rgbd().cuda()
# model.load_state_dict(torch.load(os.path.join(SAVE_DIR, 'model_overtraining.pth')))
# model.eval()
# print('model loaded for evaluation.')
# =============================================================================
# Create RGB-D model
model = Model_rgbd().cuda()
print('Model created.')
# =============================================================================
# Training parameters
optimizer = torch.optim.Adam( model.parameters(), args.lr, amsgrad=True )
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=1, crop_halfsize=ifcrop)
train_loader_l, test_loader_l = getTranslucentData(batch_size=1, crop_halfsize=ifcrop)
# Test batch is manually enlarged! See getTranslucentData's return.
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(prefix, args.lr, args.epochs, args.bs), flush_secs=30)
# Loss
l1_criterion = nn.L1Loss()
l1_criterion_masked = MaskedL1()
grad_l1_criterion_masked = MaskedL1Grad()
# Hand-craft loss weight of main task
interval1 = 1
interval2 = 2
weight_t1loss = [1] * (10*interval1) + [0] * interval2
weight_txloss = [.0317] * interval1 + [.1] * interval1 + \
[.316] * interval1 + [1] * interval1 + \
[3.16] * interval1 + [10] * interval1 + \
[10] * interval1 + [5.62] * interval1 + \
[3.16] * interval1 + [1.78] * interval1 + \
[0] * interval2
weight_t2loss = [.001] * interval1 + [.00316] * interval1 + \
[.01] * interval1 + [.0316] * interval1 + \
[.1] * interval1 + [.316] * interval1 + \
[1.0] * interval1 + [3.16] * interval1 + \
[10.0] * interval1 + [31.6] * interval1 + \
[100.0] * interval2
if not os.path.exists('%s/img' % SAVE_DIR):
os.makedirs('%s/img' % SAVE_DIR)
# Start training...
for epoch in range(0, 10*interval1 + interval2):
batch_time = AverageMeter()
losses_nyu = AverageMeter()
losses_lucent = AverageMeter()
losses_hole = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
# decide #(iter)
tot_len = min(len(train_loader), len(train_loader_l))
# print(tot_len)
trainiter = iter(train_loader)
trainiter_l = iter(train_loader_l)
for i in range(tot_len):
# print("Iteration "+str(i)+". loop start:")
try:
sample_batched = next(trainiter)
sample_batched_l = next(trainiter_l)
except StopIteration:
print(' (almost) end of iteration.')
continue
# print('in loop.')
# Prepare sample and target
image_nyu = torch.autograd.Variable(sample_batched['image'].cuda())
depth_nyu = torch.autograd.Variable(sample_batched['depth'].cuda(non_blocking=True))
image_raw = torch.autograd.Variable(sample_batched_l['image'].cuda())
mask_raw = torch.autograd.Variable(sample_batched_l['mask'].cuda())
depth_raw = torch.autograd.Variable(sample_batched_l['depth_raw'].cuda())
depth_gt = torch.autograd.Variable(sample_batched_l['depth_truth'].cuda(non_blocking=True))
# if i < 10:
# print('========-=-=')
# print(image_nyu.shape)
# print(depth_nyu.shape)
# print(image_raw.shape)
# print(" " + str(torch.max(image_raw)) + " " + str(torch.min(image_raw)))
# print(mask_raw.shape)
# print(" " + str(torch.max(mask_raw)) + " " + str(torch.min(mask_raw)))
# print(depth_raw.shape)
# print(" " + str(torch.max(depth_raw)) + " " + str(torch.min(depth_raw)))
# print(depth_gt.shape)
# print(" " + str(torch.max(depth_gt)) + " " + str(torch.min(depth_gt)))
N1 = image_nyu.shape[0]
N2 = image_raw.shape[0]
###########################
# (1) Pretext task: depth completion
# if weight_t1loss[epoch] > 0:
# Normalize depth
depth_nyu_n = DepthNorm(depth_nyu)
# Apply random mask to it
rand_index = [random.randint(0, N2-1) for k in range(N1)]
mask_new = mask_raw[rand_index, :, :, :]
depth_nyu_masked = resize2d(depth_nyu_n, (HEIGHT, WIDTH)) * mask_new
# if i < 1:
# print('========')
# print(image_nyu.shape)
# print(" " + str(torch.max(image_raw)) + " " + str(torch.min(image_raw)))
# print(depth_nyu_masked.shape)
# print(" " + str(torch.max(depth_nyu_masked)) + " " + str(torch.min(depth_nyu_masked)))
# Predict
(output_t1, _) = model(image_nyu, depth_nyu_masked)
# print(" (1): " + str(output_task1.shape))
# Calculate Loss and backprop
l_depth_t1 = l1_criterion(output_t1, depth_nyu_n)
l_grad_t1 = l1_criterion(grad_x(output_t1), grad_x(depth_nyu_n)) + l1_criterion(grad_y(output_t1), grad_y(depth_nyu_n))
l_ssim_t1 = torch.clamp((1 - ssim(output_t1, depth_nyu_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
loss_nyu = (0.1 * l_depth_t1) + (1.0 * l_grad_t1) + (1.0 * l_ssim_t1)
# loss_nyu_weighted = weight_t1loss[epoch] * loss_nyu
# https://discuss.pytorch.org/t/freeze-the-learnable-parameters-of-resnet-and-attach-it-to-a-new-network/949
# freeze_weight(model, e_stay=False, e=False, d1_stay=False, d1=True)
# optimizer.zero_grad() # moved to its new position
# loss_nyu_weighted.backward(retain_graph=True)
# optimizer.step()
if i % 150 == 0 or i < 2:
vutils.save_image(DepthNorm(depth_nyu_masked), '%s/img/A_masked_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True)
vutils.save_image(DepthNorm(output_t1), '%s/img/A_out_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True)
save_error_image(DepthNorm(output_t1) - depth_nyu, '%s/img/A_diff_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True, range=(-500, 500))
torch.cuda.empty_cache()
###########################
# (x) Transfer task: /*Fill*/ reconstruct sudo-translucent object
depth_gt_n = DepthNorm(depth_gt)
depth_raw_n = DepthNorm(depth_raw)
# if weight_txloss[epoch] > 0:
# Normalize depth
depth_gt_large = resize2d(depth_gt, (HEIGHT, WIDTH))
object_mask = thresh_mask(depth_gt_large, depth_raw)
# depth_holed = depth_raw * object_mask
depth_holed = blend_depth(depth_raw, depth_gt_large, object_mask)
# print('========')
# print(object_mask.shape)
# print(" " + str(torch.max(object_mask)) + " " + str(torch.min(object_mask)))
# print(depth_holed.shape)
# print(" " + str(torch.max(depth_holed)) + " " + str(torch.min(depth_holed)))
# print(image_raw.shape)
# print(" " + str(torch.max(image_raw)) + " " + str(torch.min(image_raw)))
(output_tx, _) = model(image_raw, DepthNorm(depth_holed))
output_tx_n = DepthNorm(output_tx)
# Calculate Loss and backprop
mask_post = resize2dmask(mask_raw, (int(HEIGHT/2), int(WIDTH/2)))
l_depth_tx = l1_criterion_masked(output_tx, depth_gt_n, mask_post)
l_grad_tx = grad_l1_criterion_masked(output_tx, depth_gt_n, mask_post)
# l_ssim_tx = torch.clamp((1 - ssim(output_tx, depth_nyu_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
loss_hole = (0.1 * l_depth_tx) + (1.0 * l_grad_tx) #+ (0 * l_ssim_tx) ####
# loss_hole_weighted = weight_txloss[epoch] * loss_hole
# for param in model.decoder1.parameters():
# param.requires_grad = False
# for param in model.decoder2.parameters():
# param.requires_grad = True
# freeze_weight(model, d1_stay=False, d1=False, d2_stay=False, d2=True)
# optimizer.zero_grad()
# loss_hole_weighted.backward(retain_graph=True) # https://pytorch.org/docs/stable/autograd.html
# optimizer.step()
if i % 150 == 0 or i < 2:
vutils.save_image(DepthNorm(depth_holed), '%s/img/C_in_%06d.png' % (SAVE_DIR, epoch * 10000 + i),
normalize=True, range=(0, 500))
vutils.save_image(object_mask, '%s/img/C_mask_%06d.png' % (SAVE_DIR, epoch * 10000 + i),
normalize=True, range=(0, 1.5))
vutils.save_image(output_tx_n, '%s/img/C_out_%06d.png' % (SAVE_DIR, epoch * 10000 + i),
normalize=True, range=(0, 500))
save_error_image(output_tx_n - depth_gt, '%s/img/C_zdiff_%06d.png' % (SAVE_DIR, epoch * 10000 + i),
normalize=True, range=(-500, 500))
torch.cuda.empty_cache()
###########################
# (2) Main task: Undistort translucent object
# Predict
(_, output_t2) = model(image_raw, depth_raw_n)
output_t2_n = DepthNorm(output_t2)
# print(" (2): " + str(output.shape))
# Calculate Loss and backprop
l_depth_t2 = l1_criterion_masked(output_t2, depth_gt_n, mask_post)
l_grad_t2 = grad_l1_criterion_masked(output_t2, depth_gt_n, mask_post)
# l_ssim_t2 = torch.clamp((1 - ssim(output_t2, depth_gt_n, val_range=1000.0/10.0)) * 0.5, 0, 1)
loss_lucent = (0.1 * l_depth_t2) + (1.0 * l_grad_t2) # + (0 * l_ssim_t2)
# loss_lucent_weighted = weight_t2loss[epoch] * loss_lucent
# optimizer.zero_grad() # moved to its new position
# loss_lucent_weighted.backward(retain_graph=True)
# optimizer.step()
if i % 150 == 0 or i < 2:
vutils.save_image(depth_raw, '%s/img/B_ln_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(0, 500))
vutils.save_image(depth_gt, '%s/img/B_gt_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(0, 500))
vutils.save_image(output_t2_n, '%s/img/B_out_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(0, 500))
save_error_image(output_t2_n-depth_gt, '%s/img/B_zdiff_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(-500, 500))
if i % 150 == 0 :
o2 = output_t2.cpu().detach().numpy()
o3 = output_t2_n.cpu().detach().numpy()
og = depth_gt.cpu().numpy()
nm = DepthNorm(depth_nyu_masked).cpu().numpy()
ng = depth_nyu.cpu().numpy()
print('> ========')
print("> Output_t2:" + str(np.min(o2)) + "~" + str(np.max(o2)) + " (" + str(np.mean(o2)) +
") // Converted to depth: " + str(np.min(o3)) + "~" + str(np.max(o3)) + " (" + str(np.mean(o3)) + ")")
print("> GT depth : " + str(np.min(og)) + "~" + str(np.max(og)) +
" // NYU GT depth from 0.0~" + str(np.max(nm)) + " to " + str(np.min(ng)) + "~" + str(np.max(ng)) + " (" + str(np.mean(ng)) + ")")
###########################
# (3) Update the network parameters
if i % 150 == 0 or i < 1:
vutils.save_image(mask_post, '%s/img/_mask_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True)
loss = (weight_t1loss[epoch] * loss_nyu) + (weight_t2loss[epoch] * loss_lucent) + (weight_txloss[epoch] * loss_hole) ####
# freeze_weight(model, e_stay=False, e=True, d2_stay=False, d2=False)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# loss = loss_nyu + loss_lucent + loss_hole
# Log losses
losses_nyu.update(loss_nyu.detach().item(), image_nyu.size(0))
losses_lucent.update(loss_lucent.detach().item(), image_raw.size(0))
losses_hole.update(loss_hole.detach().item(), image_raw.size(0))
losses.update(loss.detach().item(), image_nyu.size(0) + image_raw.size(0))
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i))))
# Log progress
niter = epoch*N+i
if i % 15 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f}) ||\t'
'NYU {l1.val:.4f} ({l1.avg:.4f}) [{l1d:.4f} | {l1g:.4f} | {l1s:.4f}]\t'
'LUC {l2.val:.4f} ({l2.avg:.4f}) [{l2d:.4f} | {l2g:.4f}]\t'
'TX {lx.val:.4f} ({lx.avg:.4f}) [{lxd:.4f} | {lxg:.4f}]'
.format(epoch, i, N, batch_time=batch_time, loss=losses, l1=losses_nyu, l1d=l_depth_t1, l1g=l_grad_t1, l1s=l_ssim_t1,
l2=losses_lucent, l2d=l_depth_t2, l2g=l_grad_t2, lx=losses_hole, lxd=l_depth_tx, lxg=l_grad_tx, eta=eta))
# Note that the numbers displayed are pre-weighted.
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 750 == 0:
LogProgress(model, writer, test_loader, test_loader_l, niter, epoch*10000+i, SAVE_DIR, HEIGHT, WIDTH)
path = os.path.join(SAVE_DIR, 'model_overtraining.pth')
torch.save(model.cpu().state_dict(), path) # saving model
model.cuda() # moving model to GPU for further training
del image_nyu, depth_nyu_masked, output_t1, image_raw, depth_raw_n, output_t2
torch.cuda.empty_cache()
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, test_loader_l, niter, epoch*10000+i, SAVE_DIR, HEIGHT, WIDTH)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# all the saves come from https://discuss.pytorch.org/t/how-to-save-a-model-from-a-previous-epoch/20252
torch.save(model.state_dict(), os.path.join(SAVE_DIR, 'epoch-{}.pth'.format(epoch)))
print('Program terminated.')
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=20,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
# Create a new directory to save logs
runs = sorted(glob.glob(os.path.join('logs', 'exp-*')))
prev_run_id = int(runs[-1].split('-')[-1]) if runs else 0
MODEL_LOG_DIR = os.path.join('logs', 'exp-{:03d}'.format(prev_run_id + 1))
CHECKPOINT_DIR = os.path.join(MODEL_LOG_DIR, 'checkpoints')
os.makedirs(CHECKPOINT_DIR)
print('Saving logs to folder: ' +
colored('"{}"'.format(MODEL_LOG_DIR), 'blue'))
# Create model
model = Model()
print('Model created.')
# Enable Multi-GPU training
print("Let's use", torch.cuda.device_count(), "GPUs!")
if torch.cuda.device_count() > 1:
print(
'Multiple GPUs being used, can\'t save model graph to Tensorboard')
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Logging
# writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(prefix, args.lr, args.epochs, args.bs), flush_secs=30)
writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph')
# Loss
l1_criterion = nn.L1Loss()
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 300 == 0:
LogProgress(model, writer, test_loader, niter)
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# Save the model checkpoint every N epochs
saveModelInterval = 2
if (epoch % saveModelInterval) == 0:
filename = os.path.join(
CHECKPOINT_DIR, 'checkpoint-epoch-{:04d}.pth'.format(epoch))
if torch.cuda.device_count() > 1:
model_params = model.module.state_dict(
) # Saving nn.DataParallel model
else:
model_params = model.state_dict()
torch.save(
{
'model_state_dict': model_params,
'optimizer_state_dict': optimizer.state_dict(),
'epoch': epoch
}, filename)
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=config.TRAIN_EPOCH,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=1e-4,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
# Create model
model = Model().cuda()
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
prefix, args.lr, args.epochs, args.bs),
flush_secs=30)
# Loss
l1_criterion = nn.L1Loss()
niter = 0
if config.load_model and config.model_name != '':
model.load_state_dict(torch.load('Checkpoints\\%s' %
config.model_name))
niter = int(config.model_name.split('_')[-3])
print('load success -> %s' % config.model_name)
# Start training...
for epoch in range(0, args.epochs):
if niter % config.save_interval == 0:
if not os.path.exists('Checkpoints\\%s' % config.save_name):
os.makedirs('Checkpoints\\%s' % config.save_name)
save_name = '%s\\net_params_%s_%s.pkl' % (
config.save_name, niter,
datetime.datetime.now().strftime("%Y-%m-%d_%H-%M"))
torch.save(model.state_dict(), 'Checkpoints\\%s' % save_name)
print('save success -> %s' % save_name)
# batch_time = AverageMeter()
loss_l1 = AverageMeter()
loss_ssim = AverageMeter()
loss_edge = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
# end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
edge = sobel_(output)
pred_edge = sobel_(depth_n)
# Compute the loss
l_sobel = nn.L1Loss()(edge, pred_edge)
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
1)
# print('nan', torch.sum(torch.isnan(image)).item())
# print('inf', torch.sum(torch.isinf(image)).item())
loss = (1.0 * l_ssim) + (0.1 * l_depth) + l_sobel
if torch.isnan(l_sobel) or torch.isnan(l_depth) or torch.isnan(
l_ssim) or torch.isinf(l_sobel) or torch.isinf(
l_depth) or torch.isinf(l_ssim):
print(0)
print('nan', torch.sum(torch.isnan(image)).item())
print('inf', torch.sum(torch.isinf(image)).item())
print(torch.isnan(l_sobel), torch.isnan(l_depth),
torch.isnan(l_ssim))
print(torch.isinf(l_sobel), torch.isinf(l_depth),
torch.isinf(l_ssim))
return
# Update step
loss_l1.update(l_depth.data.item(), image.size(0))
loss_ssim.update(l_ssim.data.item(), image.size(0))
loss_edge.update(l_sobel.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
# batch_time.update(time.time() - end)
# end = time.time()
# eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i))))
# Log progress
niter += 1
if i % 5 == 0:
# Print to console
# print('Epoch: [{0}][{1}/{2}]\t'
# 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
# 'ETA {eta}\t'
# 'Loss {loss.val:.4f} ({loss.avg:.4f})'
# .format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta))
# Log to tensorboard
writer.add_scalar('Train/L1', loss_l1.val, niter)
writer.add_scalar('Train/SSIM', loss_ssim.val, niter)
writer.add_scalar('Train/EDGE', loss_edge.val, niter)
if i % 300 == 0:
LogProgress(model, writer, test_loader, niter)
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=20,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
parser.add_argument('--o',
dest='optimizer',
help='training optimizer',
default="adam",
type=str)
parser.add_argument('--lr_decay_step',
dest='lr_decay_step',
help='step to do learning rate decay, unit is epoch',
default=5,
type=int)
parser.add_argument('--lr_decay_gamma',
dest='lr_decay_gamma',
help='learning rate decay ratio',
default=0.1,
type=float)
args = parser.parse_args()
# Create model
# model = Model().cuda()
model = UNet(n_channels=3, n_classes=1, bilinear=True).cuda()
writer_train = SummaryWriter('runs/train_0')
# model.load_state_dict(torch.load(r'models\06-05-2020_15-04-20-n15000-e10-bs4-lr0.0001\weights.epoch9_model.pth'))
print('Model created.')
logger = Logger('./logs')
# hyperparams
lr = args.lr
bs = args.bs
lr_decay_step = args.lr_decay_step
lr_decay_gamma = args.lr_decay_gamma
DOUBLE_BIAS = True
WEIGHT_DECAY = 0.0001
# params
params = []
for key, value in dict(model.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
params += [{'params': [value], 'lr': lr * (DOUBLE_BIAS + 1), \
'weight_decay': 4e-5 and WEIGHT_DECAY or 0}]
else:
params += [{'params': [value], 'lr': lr, 'weight_decay': 4e-5}]
# optimizer
if args.optimizer == "adam":
optimizer = torch.optim.Adam(params,
lr=lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=4e-5)
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'unet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size,
trainCount=68730)
# # Logging
# writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(prefix, args.lr, args.epochs, args.bs), flush_secs=30)
# Loss
# l1_criterion = nn.L1Loss()
# rmse = RMSE()
l1_criterion = nn.L1Loss()
grad_criterion = GradLoss()
normal_criterion = NormalLoss()
# eval_metric = RMSE_log()
now = datetime.datetime.now() # current date and time
runID = now.strftime("%m-%d-%Y_%H-%M-%S") + '-n' + str(
len(train_loader)) + '-e' + str(
args.epochs) + '-bs' + str(batch_size) + '-lr' + str(args.lr)
outputPath = './models/'
runPath = outputPath + runID
pathlib.Path(runPath).mkdir(parents=True, exist_ok=True)
# constants
# grad_factor = 10.
# normal_factor = 10.
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
# l_ssim = torch.clamp((1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=5.0)) * 0.5, 0,
1) # sbasak01
# sbasak01 loss modification
grad_real, grad_fake = imgrad_yx(depth_n), imgrad_yx(output)
grad_loss = grad_criterion(
grad_fake, grad_real) # * grad_factor # * (epoch > 3)
normal_loss = normal_criterion(
grad_fake, grad_real) # * normal_factor # * (epoch > 7)
# loss = (1.0 * l_ssim) + (0.1 * l_depth)
loss = (l_ssim) + (l_depth) + (grad_loss) + (normal_loss)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# ADD TRAINING LOSS TO SummaryWriter
writer_train.add_scalar('LOSS', loss.data.item(), i)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print(
'Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'L1_Loss: {l1_loss:.4f} SSIM_Loss: {ssim_loss:.4f} grad_loss: {gradloss:.4f} normal_loss: {'
'normalloss:.4f} '.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta,
l1_loss=l_depth,
ssim_loss=l_ssim,
gradloss=grad_loss, # sbasak01 loss modification
normalloss=normal_loss # sbasak01 loss modification
))
# # 1. Log scalar values (scalar summary)
# info = {'loss': loss.item()}
#
# for tag, value in info.items():
# logger.scalar_summary(tag, value, i + 1)
#
# # 2. Log values and gradients of the parameters (histogram summary)
# for tag, value in model.named_parameters():
# tag = tag.replace('.', '/')
# logger.histo_summary(tag, value.data.cpu().numpy(), i + 1)
# logger.histo_summary(tag + '/grad', value.grad.data.cpu().numpy(), i + 1)
if i % 100 == 0:
sys.stdout.write("Iterations: %d \r" % (i))
sys.stdout.flush()
# save Model intermediate
path = runPath + '/weights.epoch{0}_model.pth'.format(epoch)
torch.save(model.cpu().state_dict(), path) # saving model
model.cuda()
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=50,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=6, type=int, help='batch size')
parser.add_argument('--o',
dest='optimizer',
help='training optimizer',
default="adam",
type=str)
parser.add_argument('--lr_decay_step',
dest='lr_decay_step',
help='step to do learning rate decay, unit is epoch',
default=5,
type=int)
parser.add_argument('--lr_decay_gamma',
dest='lr_decay_gamma',
help='learning rate decay ratio',
default=0.1,
type=float)
args = parser.parse_args()
# Create model
model = Model().cuda()
macs, params = get_model_complexity_info(model, (3, 480, 640),
as_strings=True,
print_per_layer_stat=True,
verbose=True)
#print(summary(model, (3, 480, 640)))
print('{:<30} {:<8}'.format('Computational complexity: ', macs))
print('{:<30} {:<8}'.format('Number of parameters: ', params))
# hyperparams
lr = args.lr
bs = args.bs
lr_decay_step = args.lr_decay_step
lr_decay_gamma = args.lr_decay_gamma
DOUBLE_BIAS = True
WEIGHT_DECAY = 0.0001
# params
params = []
for key, value in dict(model.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
params += [{'params': [value], 'lr': lr * (DOUBLE_BIAS + 1), \
'weight_decay': 4e-5 and WEIGHT_DECAY or 0}]
else:
params += [{'params': [value], 'lr': lr, 'weight_decay': 4e-5}]
# optimizer
if args.optimizer == "adam":
optimizer = torch.optim.Adam(params,
lr=lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=4e-5)
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'mobilenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size,
trainCount=159569)
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
prefix, args.lr, args.epochs, args.bs),
flush_secs=30)
# Loss
# l1_criterion = nn.L1Loss()
# rmse = RMSE()
l1_criterion = nn.L1Loss()
grad_criterion = GradLoss()
normal_criterion = NormalLoss()
# eval_metric = RMSE_log()
now = datetime.datetime.now() # current date and time
runID = now.strftime("%m-%d-%Y_%H-%M-%S") + '-n' + str(
len(train_loader)) + '-e' + str(
args.epochs) + '-bs' + str(batch_size) + '-lr' + str(args.lr)
outputPath = './models_fk_mobile_sum_weights/'
runPath = outputPath + runID
pathlib.Path(runPath).mkdir(parents=True, exist_ok=True)
# constants
grad_factor = 10.
normal_factor = 10.
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=5.0)) * 0.5, 0,
1) # sbasak01
# FK01 loss modification
grad_real, grad_fake = imgrad_yx(depth_n), imgrad_yx(output)
grad_loss = grad_criterion(
grad_fake, grad_real) * grad_factor # * (epoch > 3)
normal_loss = normal_criterion(
grad_fake, grad_real) * normal_factor # * (epoch > 7)
loss = (0.28 * l_ssim) + (0.22 * l_depth) + (0.30 * grad_loss) + (
0.20 * normal_loss) # Fk01 loss modification
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print(
'Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'L1_Loss: {l1_loss:.4f} SSIM_Loss: {ssim_loss:.4f} grad_loss: {gradloss:.4f} normal_loss: {'
'normalloss:.4f} '.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta,
l1_loss=l_depth,
ssim_loss=l_ssim,
gradloss=grad_loss, # Fk1 loss modification
normalloss=normal_loss # Fk01 loss modification
))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# save Model intermediate
path = runPath + '/weights.epoch{0}_model.pth'.format(epoch)
torch.save(model.cpu().state_dict(), path) # saving model
model.cuda()
def trainAndVal(loader, model, l1_criterion, optimizer=None):
batch_time = AverageMeter()
losses = AverageMeter()
rsmes = AverageMeter()
if (optimizer):
# switch to train mode
model.train()
print('Train', flush=True)
else:
# switch to evaluate mode
model.eval()
print('Val', flush=True)
N = len(loader)
end = time.time()
start = end
if (optimizer is None):
predictions = []
testSetDepths = []
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# process epoch
for i, sample_batched in enumerate(loader):
# Prepare sample and target
image = sample_batched['image'].to(device)
depth = sample_batched['depth'].to(device)
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# measure accuracy and record loss
losses.update(loss.data, image.size(0))
rmse = (depth_n.data.cpu() - output.data.cpu())**2
rmse = np.sqrt(rmse.mean())
rsmes.update(rmse, image.size(0))
if (optimizer):
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.avg * (N - i))))
total = str(
datetime.timedelta(seconds=int((time.time() - start) +
batch_time.avg * (N - i))))
minDepth = 10
maxDepth = 1000
if (optimizer is None):
predictions.append(output.squeeze().data.cpu())
testSetDepths.append(depth_n.squeeze().data.cpu())
if i % 5 == 0:
p = 100 * i / N
bar = "[%-10s] %d%%" % ('=' * int(p * 10 / 100) + '.' *
(10 - int(p * 10 / 100)), p)
print('[{0}/{1}] {2} - '
'Batch Time: {batch_time.val:.2f} ({batch_time.avg:.2f}) '
'ETA: {eta}/{total} '
'Loss: {loss.val:.3f} ({loss.avg:.3f}) '
'RSME: {rsme.val:.3f} ({rsme.avg:.3f})'.format(
i,
N,
bar,
batch_time=batch_time,
eta=eta,
total=total,
loss=losses,
rsme=rsmes),
flush=True)
break
if (optimizer is None):
predictions = np.vstack(predictions)
testSetDepths = np.vstack(testSetDepths)
e = compute_errors(predictions, testSetDepths)
print("{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}".format(
'a1', 'a2', 'a3', 'rel', 'rms', 'log_10'))
print("{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}".
format(e[0], e[1], e[2], e[3], e[4], e[5]))
return losses.avg
def main():
# Arguments
parser = argparse.ArgumentParser(
description='Finetuning for depth estimation')
parser.add_argument('--epochs',
default=60,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--model_type',
default='densedepth',
help='type of the depth estimation network')
parser.add_argument('--layers',
default=161,
type=int,
help='number of layers of encoder')
parser.add_argument('--bs', default=8, type=int, help='batch size')
args = parser.parse_args()
# Create model
if args.model_type == 'densedepth':
model = Model(layers=args.layers)
else:
model = ResNet(layers=args.layers)
model = model.cuda()
model = nn.DataParallel(model)
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Loss
l1_criterion = nn.L1Loss()
masked_criterion = MaskedL1Loss()
for epoch in range(args.epochs):
model.train()
for i, sample_batched in enumerate(tqdm.tqdm(train_loader)):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
if args.model_type == 'densedepth':
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5,
0, 1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
else:
loss = masked_criterion(output, depth_n)
# Update step
loss.backward()
optimizer.step()
torch.save(model.module.state_dict(),
'checkpoints/%s_%d.pth' % (args.model_type, args.layers))
print('Epoch:%d Model Saved!' % (epoch + 1))
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=config.TRAIN_EPOCH,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=1e-4,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
# Create model
model = Model().cuda()
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
nyu_list = open('nyu_list.txt').readlines()
scannet_list = open('scannet_list.txt').readlines()
real_list = nyu_list + scannet_list
irs_list = open('irs_list.txt').readlines()
mask_list = open('mask_list.txt').readlines()
test_list = open('test_list.txt').readlines()
train_dataset = TripleDataset(real_list, irs_list, mask_list, None)
test_dataset = NyuV2Dataset(test_list)
train_loader = DataLoader(dataset=train_dataset,
batch_size=config.bs,
num_workers=8,
shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
batch_size=4,
num_workers=8,
shuffle=True)
# Load data
# train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
prefix, args.lr, args.epochs, args.bs),
flush_secs=30)
# Loss
l1_criterion = nn.L1Loss().cuda()
niter = 0
if config.load_model and config.model_name != '':
checkpoint = torch.load('Checkpoints\\%s' % config.model_name)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
niter = int(config.model_name.split('_')[-3])
print('load success -> %s' % config.model_name)
# Start training...
for epoch in range(0, args.epochs):
# batch_time = AverageMeter()
loss_l1 = AverageMeter()
loss_ssim = AverageMeter()
loss_edge = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
# end = time.time()
idx = 0
for syn_input, real_input, syn_label, real_label in (train_loader):
if niter % config.save_interval == 0:
if not os.path.exists('Checkpoints\\%s' % config.save_name):
os.makedirs('Checkpoints\\%s' % config.save_name)
save_name = '%s\\net_params_%s_%s.pkl' % (
config.save_name, niter,
datetime.datetime.now().strftime("%Y-%m-%d_%H-%M"))
state = {
'net': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(state, 'Checkpoints\\%s' % save_name)
print('save success -> %s' % save_name)
optimizer.zero_grad()
_input = torch.cat([syn_input, real_input], dim=0)
_label = torch.cat([syn_label, real_label], dim=0)
# Prepare sample and target
image = torch.autograd.Variable(_input.float().cuda())
depth = torch.autograd.Variable(
_label.float().cuda(non_blocking=True))
# Normalize depth
# depth_n = DepthNorm(depth)
# Predict
output = model(image)
# pred_edge = sobel_(output)
# edge = sobel_(depth)
# pred_edge = torch.where(pred_edge>0.2, pred_edge, torch.full_like(pred_edge, 0))
# edge = torch.where(edge > 0.2, edge, torch.full_like(pred_edge, 0))
# Compute the loss
# l_sobel = nn.L1Loss()(edge, pred_edge)
l_depth = l1_criterion(output, depth)
l_ssim = torch.clamp(
(1 - ssim(output, depth, val_range=10.0)) * 0.5, 0, 1)
# print('nan', torch.sum(torch.isnan(image)).item())
# print('inf', torch.sum(torch.isinf(image)).item())
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# if torch.isnan(l_sobel) or torch.isnan(l_depth) or torch.isnan(l_ssim) or torch.isinf(l_sobel) or torch.isinf(l_depth) or torch.isinf(l_ssim):
# print(0)
# print('nan', torch.sum(torch.isnan(image)).item())
# print('inf', torch.sum(torch.isinf(image)).item())
# print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim))
# print(torch.isinf(l_sobel), torch.isinf(l_depth), torch.isinf(l_ssim))
# return
# Update step
loss_l1.update(l_depth.data.item(), image.size(0))
loss_ssim.update(l_ssim.data.item(), image.size(0))
# loss_edge.update(l_sobel.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
# batch_time.update(time.time() - end)
# end = time.time()
# eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i))))
# Log progress
niter += 1
if idx % 50 == 0:
# Print to console
# print('Epoch: [{0}][{1}/{2}]\t'
# 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
# 'ETA {eta}\t'
# 'Loss {loss.val:.4f} ({loss.avg:.4f})'
# .format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta))
# Log to tensorboard
writer.add_scalar('Train/L1', loss_l1.val, niter)
writer.add_scalar('Train/SSIM', loss_ssim.val, niter)
# writer.add_scalar('Train/EDGE', loss_edge.val, niter)
if idx % 1000 == 0:
LogProgress(model, writer, test_loader, niter)
idx += 1
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=20,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--data-folder',
type=str,
help='location of dataset folder')
parser.add_argument('--pretrained-weights',
type=str,
default=None,
help='location of pretrained weights')
parser.add_argument('--label-file',
type=str,
help='name of label txt file')
parser.add_argument('--bs', default=4, type=int, help='batch size')
parser.add_argument('--num-instances',
default=None,
type=int,
help='number of instances of same image in the batch')
parser.add_argument('--lamb',
default=0.,
type=float,
help='lambda: multiplier for added additional loss')
parser.add_argument('--evaluate',
type=str,
default=None,
help='path of model file for testing')
args = parser.parse_args()
print(args)
# Create model
model = PTResModel(pretrained_weights=args.pretrained_weights)
evaluating = (args.evaluate != None)
if (evaluating):
print('Evaluating ', args.evaluate)
checkpoint = torch.load(args.evaluate)
model.load_state_dict(checkpoint)
model = model.cuda()
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
# Load data
data_test, nyu2_test = loadToMem(args.data_folder,
txtfile='/nyu2_test_updated.csv')
transformed_testing = depthDatasetMemory(
data_test, nyu2_test, transform=getNoTransform(is_test=True))
test_loader = DataLoader(transformed_testing, 1, shuffle=False)
if (evaluating):
evaluate_model(model, test_loader, args)
return
data, nyu2_train = loadToMem(args.data_folder, txtfile=args.label_file)
transformed_training = depthDatasetMemory(
data, nyu2_train, transform=getDefaultTrainTransform())
if (args.num_instances):
train_loader = DataLoader(transformed_training,
batch_size,
sampler=RandomIdentitySampler(
transformed_training,
num_instances=args.num_instances),
num_workers=4,
drop_last=True)
else:
train_loader = DataLoader(transformed_training,
batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
# Logging
writer = SummaryWriter('logs', flush_secs=30)
# Loss
l1_criterion = nn.L1Loss()
l1_criterion = l1_criterion.cuda()
ps = AllPositivePairSelector(balance=False)
criterion2 = OnlineContrastiveLoss(1., ps)
criterion2 = criterion2.cuda()
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
transformed_training.reset_seeds()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
image_id = torch.autograd.Variable(
sample_batched['image_id'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output, feats = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
if (args.lamb > 0):
l_mse = criterion2(feats, image_id, neg_loss=False)
loss += args.lamb * l_mse
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 300 == 0:
LogProgress(model, writer, test_loader, niter)
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
torch.save(model.state_dict(), "models_" + str(epoch) + '.pth.tar')
evaluate_model(model, test_loader, args)
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=20,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
# Create model
model = Model().cuda()
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
prefix, args.lr, args.epochs, args.bs),
flush_secs=30)
# Loss
l1_criterion = nn.L1Loss()
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 300 == 0:
LogProgress(model, writer, test_loader, niter)
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# Forward pass of the mini-batch
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
outputs = model(image)
# Compute the loss
l_depth = l1_criterion(outputs, depth_n)
l_ssim = torch.clamp(
(1 - ssim(outputs, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
running_loss += loss.item()
# Save output images, one at a time, to results
inputs_tensor = image.detach().cpu()
output_tensor = outputs.detach().cpu()
label_tensor = depth_n.detach().cpu()
depth_metric = depth * (config.train.max_depth / 1000.0)
outputs_tmp = DepthNorm(outputs)
outputs_metric = outputs_tmp * (config.train.max_depth / 1000.0)
metrics = compute_errors(depth_metric, outputs_metric)
# print(metrics)
def main():
# Arguments
parser = argparse.ArgumentParser(description='High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('-c', '--configFile', required=True, help='Path to config yaml file', metavar='path/to/config')
args = parser.parse_args()
CONFIG_FILE_PATH = args.configFile
with open(CONFIG_FILE_PATH) as fd:
config_yaml = oyaml.load(fd) # Returns an ordered dict. Used for printing
config = AttrDict(config_yaml)
print(colored('Config being used for training:\n{}\n\n'.format(oyaml.dump(config_yaml)), 'green'))
# Create a new directory to save logs
runs = sorted(glob.glob(os.path.join(config.train.logsDir, 'exp-*')))
prev_run_id = int(runs[-1].split('-')[-1]) if runs else 0
MODEL_LOG_DIR = os.path.join(config.train.logsDir, 'exp-{:03d}'.format(prev_run_id + 1))
CHECKPOINT_DIR = os.path.join(MODEL_LOG_DIR, 'checkpoints')
os.makedirs(CHECKPOINT_DIR)
print('Saving logs to folder: ' + colored('"{}"'.format(MODEL_LOG_DIR), 'blue'))
# Save a copy of config file in the logs
shutil.copy(CONFIG_FILE_PATH, os.path.join(MODEL_LOG_DIR, 'config.yaml'))
# Create a tensorboard object and Write config to tensorboard
writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph')
string_out = io.StringIO()
oyaml.dump(config_yaml, string_out, default_flow_style=False)
config_str = string_out.getvalue().split('\n')
string = ''
for line in config_str:
string = string + ' ' + line + '\n\r'
writer.add_text('Config', string, global_step=None)
# Create model
model = Model()
print('Model created.')
# to continue training from a checkpoint
if config.train.continueTraining:
print('Transfer Learning enabled. Model State to be loaded from a prev checkpoint...')
if not os.path.isfile(config.train.pathPrevCheckpoint):
raise ValueError('Invalid path to the given weights file for transfer learning.\
The file {} does not exist'.format(config.train.pathPrevCheckpoint))
CHECKPOINT = torch.load(config.train.pathPrevCheckpoint, map_location='cpu')
if 'model_state_dict' in CHECKPOINT:
# Newer weights file with various dicts
print(colored('Continuing training from checkpoint...Loaded data from checkpoint:', 'green'))
print('Config Used to train Checkpoint:\n', oyaml.dump(CHECKPOINT['config']), '\n')
print('From Checkpoint: Last Epoch Loss:', CHECKPOINT['epoch_loss'], '\n\n')
model.load_state_dict(CHECKPOINT['model_state_dict'])
elif 'state_dict' in CHECKPOINT:
# reading original authors checkpoints
if config.train.model != 'rednet':
# original author deeplab checkpoint
CHECKPOINT['state_dict'].pop('decoder.last_conv.8.weight')
CHECKPOINT['state_dict'].pop('decoder.last_conv.8.bias')
else:
# rednet checkpoint
# print(CHECKPOINT['state_dict'].keys())
CHECKPOINT['state_dict'].pop('final_deconv.weight')
CHECKPOINT['state_dict'].pop('final_deconv.bias')
CHECKPOINT['state_dict'].pop('out5_conv.weight')
CHECKPOINT['state_dict'].pop('out5_conv.bias')
CHECKPOINT['state_dict'].pop('out4_conv.weight')
CHECKPOINT['state_dict'].pop('out4_conv.bias')
CHECKPOINT['state_dict'].pop('out3_conv.weight')
CHECKPOINT['state_dict'].pop('out3_conv.bias')
CHECKPOINT['state_dict'].pop('out2_conv.weight')
CHECKPOINT['state_dict'].pop('out2_conv.bias')
model.load_state_dict(CHECKPOINT['state_dict'], strict=False)
else:
# Old checkpoint containing only model's state_dict()
model.load_state_dict(CHECKPOINT)
# Enable Multi-GPU training
print("Let's use", torch.cuda.device_count(), "GPUs!")
if torch.cuda.device_count() > 1:
print('Multiple GPUs being used, can\'t save model graph to Tensorboard')
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Training parameters
optimizer = torch.optim.Adam( model.parameters(), config.train.optimAdam.learningRate )
batch_size = config.train.batchSize
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader_list = []
test_loader_list = []
for dataset in config.train.datasetsTrain:
train_data = getTrainingTestingData('rgb', 'train', dataset.images, dataset.labels)
train_loader_list.append(train_data)
for dataset in config.train.datasetsVal:
print(dataset.images)
test_data = getTrainingTestingData('rgb', 'eval', dataset.images, dataset.labels)
test_loader_list.append(test_data)
train_loader = DataLoader(torch.utils.data.ConcatDataset(train_loader_list), batch_size, num_workers=config.train.numWorkers, shuffle=True, drop_last=True, pin_memory=True)
test_loader = DataLoader(torch.utils.data.ConcatDataset(test_loader_list), batch_size, num_workers=config.train.numWorkers, shuffle=False, drop_last=True, pin_memory=True)
print(len(torch.utils.data.ConcatDataset(train_loader_list)))
print(len(train_loader))
print(len(test_loader))
# Create a tensorboard object and Write config to tensorboard
writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph')
# Loss
l1_criterion = nn.L1Loss()
total_iter_num = 0
# Start training...
for epoch in range(config.train.numEpochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Log the current Epoch Number
writer.add_scalar('data/Epoch Number', epoch, total_iter_num)
# Switch to train mode
model.train()
end = time.time()
running_loss = 0.0
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
total_iter_num += 1
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm( depth )
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp((1 - ssim(output, depth_n, val_range = 1000.0 / 10.0)) * 0.5, 0, 1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i))))
# Log progress
niter = epoch*N+i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'
.format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 50 == 0:
LogProgress(model, writer, test_loader, niter)
# Log Epoch Loss
epoch_loss = running_loss / (len(train_loader))
writer.add_scalar('data/Train Epoch Loss', epoch_loss, total_iter_num)
print('\nTrain Epoch Loss: {:.4f}'.format(epoch_loss))
metrics = compute_errors(depth_n, output)
print(metrics)
for keys, values in metrics.items():
print(str(keys) + ':' + str(values))
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# Save the model checkpoint every N epochs
if (epoch % config.train.saveModelInterval) == 0:
filename = os.path.join(CHECKPOINT_DIR, 'checkpoint-epoch-{:04d}.pth'.format(epoch))
if torch.cuda.device_count() > 1:
model_params = model.module.state_dict() # Saving nn.DataParallel model
else:
model_params = model.state_dict()
torch.save(
{
'model_state_dict': model_params,
'optimizer_state_dict': optimizer.state_dict(),
'epoch': epoch,
'total_iter_num': total_iter_num,
'epoch_loss': epoch_loss,
'config': config_yaml
}, filename)
def LogProgress(model, writer, test_loader, global_step):
model.eval()
sequential = test_loader
sample_batched = next(iter(sequential))
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# print()
# if epoch == 0:
# writer.add_image('Train.1.Image', vutils.make_grid(image.data, nrow=6, normalize=True), epoch)
# if epoch == 0:
writer.add_image(
'Train.2.Depth',
colorize(vutils.make_grid(depth.data, nrow=6, normalize=False)),
global_step)
output = model(image)
plt.imshow(output.detach().cpu().numpy().squeeze(), cmap='plasma')
plt.show()
# Predict
# output = model(image)
edge = sobel_(output)
pred_edge = sobel_(depth_n)
# Compute the loss
l_sobel = nn.L1Loss()(edge, pred_edge)
l_depth = nn.L1Loss()(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
if torch.isnan(l_sobel):
print(1)
print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim))
writer.add_scalar('Test/L1', l_depth.item(), global_step)
writer.add_scalar('Test/SSIM', l_ssim.item(), global_step)
writer.add_scalar('Test/EDGE', l_sobel.item(), global_step)
# normal = Norm(output)
# normal = (normal + 1.0) / 2.
# normal = normal[3, :, :, :]
# normal = normal.detach().cpu().numpy().astype(np.uint8)
# normal = np.transpose(normal, (1, 2, 0))
# print()
# plt.imshow(normal)
# plt.show()
output = DepthNorm(output)
# writer.add_image('Train.3.Normal', vutils.make_grid(normal.data, nrow=6, normalize=False), epoch)
writer.add_image(
'Test.2.Ours',
colorize(vutils.make_grid(output.data, nrow=6, normalize=False)),
global_step)
writer.add_image(
'Test.3.Diff',
colorize(
vutils.make_grid(torch.abs(output - depth).data,
nrow=6,
normalize=False)), global_step)
del image
del depth
del output
del edge
del pred_edge
