def train(train_loader, model, criterion, optimizer, epoch, logger):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
batch_num = len(train_loader)
for i, (input, target) in enumerate(train_loader):
# itr_count += 1
input, target = input.cuda(), target.cuda()
# print('input size = ', input.size())
# print('target size = ', target.size())
torch.cuda.synchronize()
data_time = time.time() - end
# compute pred
end = time.time()
pred = model(input) # @wx 注意输出
# print('pred size = ', pred.size())
# print('target size = ', target.size())
loss = criterion(pred, target)
optimizer.zero_grad()
loss.backward() # compute gradient and do SGD step
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
if (i + 1) % args.print_freq == 0:
print('=> output: {}'.format(output_directory))
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'Loss={Loss:.5f} '
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'RML={result.absrel:.2f}({average.absrel:.2f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
epoch, i + 1, len(train_loader), data_time=data_time,
gpu_time=gpu_time, Loss=loss.item(), result=result, average=average_meter.average()))
current_step = epoch * batch_num + i
logger.add_scalar('Train/RMSE', result.rmse, current_step)
logger.add_scalar('Train/rml', result.absrel, current_step)
logger.add_scalar('Train/Log10', result.lg10, current_step)
logger.add_scalar('Train/Delta1', result.delta1, current_step)
logger.add_scalar('Train/Delta2', result.delta2, current_step)
logger.add_scalar('Train/Delta3', result.delta3, current_step)
avg = average_meter.average()
def train(train_loader, model, criterion, optimizer, epoch):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
for i, (input, target) in enumerate(train_loader):
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute pred
end = time.time()
pred = model(input)
loss = criterion(pred, target)
optimizer.zero_grad()
loss.backward() # compute gradient and do SGD step
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
if (i + 1) % args.print_freq == 0:
print('=> output: {}'.format(output_directory))
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
epoch,
i + 1,
len(train_loader),
data_time=data_time,
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
with open(train_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'gpu_time': avg.gpu_time,
'data_time': avg.data_time
})
def train(train_loader, model, criterion, optimizer, epoch):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
for i, (input, target) in enumerate(train_loader):
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute pred
end = time.time()
pred = model(input)
loss = criterion(pred, target)
optimizer.zero_grad()
loss.backward() # compute gradient and do SGD step
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
depth_in = np.hstack(input.data.cpu().numpy()[:4, 3] / 10.)
depth_in = cv2.applyColorMap((depth_in * 255).astype(np.uint8), cv2.COLORMAP_HOT)
tgt_out = np.hstack(np.squeeze(target[:4].data.cpu().numpy())) / 10.
tgt_out = cv2.applyColorMap((tgt_out * 255).astype(np.uint8), cv2.COLORMAP_HOT)
out = np.hstack(np.squeeze(pred[:4].data.cpu().numpy()))
out = np.clip(out / 10., 0., 1.)
out = cv2.applyColorMap((out * 255).astype(np.uint8), cv2.COLORMAP_HOT)
if i % 20 == 0:
cv2.imshow("Training Results", np.vstack([depth_in, tgt_out, out]))
cv2.waitKey(1)
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
if (i + 1) % args.print_freq == 0:
print('=> output: {}'.format(output_directory))
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
epoch, i+1, len(train_loader), data_time=data_time,
gpu_time=gpu_time, result=result, average=average_meter.average()))
avg = average_meter.average()
with open(train_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({'mse': avg.mse, 'rmse': avg.rmse, 'absrel': avg.absrel, 'lg10': avg.lg10,
'mae': avg.mae, 'delta1': avg.delta1, 'delta2': avg.delta2, 'delta3': avg.delta3,
'gpu_time': avg.gpu_time, 'data_time': avg.data_time})
def validate(epoch, valData, model, logger):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
# skip = len(valData) // 9 # save images every skip iters
for i, (image, depth) in enumerate(valData):
image = image.cuda()
depth = depth.cuda()
# normal = normal.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
end = time.time()
with torch.no_grad():
pred = model(image)
torch.cuda.synchronize()
gpu_time = time.time() - end
result = Result()
result.evaluate(pred.data, depth.data)
average_meter.update(result, gpu_time, data_time, image.size(0))
end = time.time()
if (i + 1) % 10 == 0:
print('Test Epoch: [{0}/{1}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'RML={result.absrel:.2f}({average.absrel:.2f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
i + 1, len(valData), data_time=data_time,
gpu_time=gpu_time, result=result, average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'Rel={average.absrel:.3f}\n'
'Log10={average.lg10:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'Delta2={average.delta2:.3f}\n'
'Delta3={average.delta3:.3f}\n'
't_GPU={time:.3f}\n'.format(
average=avg, time=avg.gpu_time))
logger.add_scalar('Test/rmse', avg.rmse, epoch)
logger.add_scalar('Test/Rel', avg.absrel, epoch)
logger.add_scalar('Test/log10', avg.lg10, epoch)
logger.add_scalar('Test/Delta1', avg.delta1, epoch)
logger.add_scalar('Test/Delta2', avg.delta2, epoch)
logger.add_scalar('Test/Delta3', avg.delta3, epoch)
return avg
def train(epoch, trainData, model, crite, optimizer, logger):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
for i, (image, depth) in enumerate(trainData):
image = image.cuda()
depth = depth.cuda()
# normal = normal.cuda()
# image = torch.autograd.Variable(image)
# depth = torch.autograd.Variable(depth)
torch.cuda.synchronize()
data_time = time.time() - end
end = time.time()
optimizer.zero_grad()
pred = model(image)
loss = crite(pred, depth)
loss.backward()
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
result = Result()
result.evaluate(pred.data, depth.data)
average_meter.update(result, gpu_time, data_time, image.size(0))
end = time.time()
if (i + 1) % 10 == 0:
print('=> output: {}'.format(opt.output_dir))
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'Loss={Loss:.5f} '
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'RML={result.absrel:.2f}({average.absrel:.2f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
epoch, i + 1, len(trainData), data_time=data_time,
gpu_time=gpu_time, Loss=loss.item(), result=result, average=average_meter.average()))
current_step = epoch * len(trainData) + i
logger.add_scalar('Train/loss', loss, current_step)
logger.add_scalar('Train/RMSE', result.rmse, current_step)
logger.add_scalar('Train/rml', result.absrel, current_step)
logger.add_scalar('Train/Log10', result.lg10, current_step)
logger.add_scalar('Train/Delta1', result.delta1, current_step)
logger.add_scalar('Train/Delta2', result.delta2, current_step)
logger.add_scalar('Train/Delta3', result.delta3, current_step)
def compute_depth_metrics(self, verbose=True) -> Result:
"""Computes metrics on the difference between raw and fixed depth values"""
avg = AverageMeter()
for i, path in enumerate(self.paths):
_, depth_raw, depth_fix = self.load_images(path)
depth_raw = torch.tensor(depth_raw)
depth_fix = torch.tensor(depth_fix)
res = Result()
res.evaluate(depth_raw, depth_fix)
avg.update(res, 0, 0, 1)
if verbose:
stdout.write(f"=> computing img {i}/{len(self)}\r")
if verbose:
stdout.write("\n")
return avg.average()
def train_coarse(train_loader, model, criterion, optimizer, epoch):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
for i, (input, target) in enumerate(train_loader):
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute pred
end = time.time()
pred = model(input)
loss = criterion(pred, target)
optimizer.zero_grad()
loss.backward() # compute gradient and do SGD step
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
end = time.time()
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
eval_time = time.time() - end
if (i + 1) % args.print_freq == 0:
history_loss.append(loss.item())
print('=> output: {}'.format(output_directory))
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
epoch,
i + 1,
len(train_loader),
data_time=data_time,
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
def train(train_loader, val_loader, model, criterion, optimizer, epoch, lr):
average_meter = AverageMeter()
model.train() # switch to train mode
global batch_num, best_result
end = time.time()
#every batch
for i, (Y, Y_1_2, Y_1_4, Y_1_8, LR, LR_8, HR,
name) in enumerate(train_loader): #处理被train_loader进去的每一个数据
batch_num = batch_num + 1
Y = Y.cuda()
Y_1_2 = Y_1_2.cuda()
Y_1_4 = Y_1_4.cuda()
Y_1_8 = Y_1_8.cuda()
LR = LR.cuda()
LR_8 = LR_8.cuda()
HR = HR.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
end = time.time()
if args.arch == 'VDSR_16':
pred_HR = model(LR)
loss = criterion(pred_HR, HR, Y)
elif args.arch == 'VDSR_16_2':
pred_HR = model(Y, LR)
loss = criterion(pred_HR, HR, Y)
elif args.arch == 'VDSR':
pred_HR, residule = model(LR_8, Y)
loss = criterion(pred_HR, HR, Y)
elif args.arch == 'ResNet_bicubic':
pred_HR, residule = model(LR_8, Y)
loss = criterion(pred_HR, HR, Y)
elif args.arch == 'resnet50_15_6' or 'resnet50_15_11' or 'resnet50_15_12':
pred_HR = model(Y_1_2, LR)
loss = criterion(pred_HR, HR, Y)
elif args.arch == 'resnet50_15_2' or 'resnet50_15_3' or 'resnet50_15_5' or 'resnet50_15_8' or 'resnet50_15_9':
pred_HR, residule = model(Y, LR, LR_8)
loss = criterion(pred_HR, HR, Y)
else:
if config.loss_num == 2:
if config.LOSS_1 == 'l2':
# 均方差
criterion1 = criteria.MaskedMSELoss().cuda()
elif config.LOSS_1 == 'l1':
criterion1 = criteria.MaskedL1Loss().cuda()
elif config.LOSS_1 == 'l1_canny':
# 均方差
criterion1 = criteria.MaskedL1_cannyLoss().cuda()
elif config.LOSS_1 == 'l1_from_rgb_sobel':
# 均方差
criterion1 = criteria.MaskedL1_from_rgb_sobel_Loss().cuda()
elif aconfig.LOSS_1 == 'l1_canny_from_GT_canny':
criterion1 = criteria.MaskedL1_canny_from_GT_Loss().cuda()
elif aconfig.LOSS_1 == 'l1_from_GT_sobel':
criterion1 = criteria.MaskedL1_from_GT_sobel_Loss().cuda()
elif config.LOSS_1 == 'l2_from_GT_sobel_Loss':
criterion1 = criteria.MaskedL2_from_GT_sobel_Loss().cuda()
if config.use_different_size_Y == 1:
pred_HR, pred_thermal0 = model(Y, Y_1_2, Y_1_4, Y_1_8, LR)
else:
pred_HR, pred_thermal0 = model(Y, LR)
#final loss
loss0 = criterion(pred_HR, HR, Y)
#therma upsample loss
loss1 = criterion1(pred_thermal0, HR, Y)
loss = config.LOSS_0_weight * loss0 + config.LOSS_1_weight * loss1
else:
if config.use_different_size_Y == 1:
pred_HR, pred_thermal0 = model(Y, Y_1_2, Y_1_4, Y_1_8, LR)
#writer = SummaryWriter(log_dir='logs')
#writer.add_graph(model, input_to_model=(Y,Y_1_2,Y_1_4,Y_1_8,LR,))
else:
pred_HR, pred_thermal0 = model(Y, LR)
loss = criterion(pred_HR, HR, Y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred_HR, HR, loss.cpu().detach().numpy())
average_meter.update(result, gpu_time, data_time, Y.size(0))
end = time.time()
if (i + 1) % args.print_freq == 0:
print('=> output: {}'.format(output_directory))
print('Dataset Epoch: {0} [{1}/{2}]\t'
'Batch Epoch: {3} \t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n'
'PSNR={result.psnr:.5f}({average.psnr:.5f}) '
'MSE={result.mse:.3f}({average.mse:.3f}) '
'RMSE={result.rmse:.3f}({average.rmse:.3f}) '
'MAE={result.mae:.3f}({average.mae:.3f}) '
'Delta1={result.delta1:.4f}({average.delta1:.4f}) '
'REL={result.absrel:.4f}({average.absrel:.4f}) '
'Lg10={result.lg10:.4f}({average.lg10:.4f}) '
'Loss={result.loss:}({average.loss:}) '.format(
epoch,
i + 1,
len(train_loader),
batch_num,
data_time=data_time,
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
else:
pass
if (batch_num + 1) % config.save_fc == 0:
print("==============Time to evaluate=================")
utils.adjust_learning_rate(optimizer, batch_num, lr)
print("==============SAVE_MODEL=================")
avg = average_meter.average()
average_meter = AverageMeter()
with open(train_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'dataset epoch': epoch,
'batch epoch': batch_num + 1,
'psnr': 10 * math.log(1 / (avg.mse), 10),
'mse': result.mse,
'rmse': result.rmse,
'absrel': result.absrel,
'lg10': result.lg10,
'mae': result.mae,
'delta1': result.delta1,
'delta2': result.delta2,
'delta3': result.delta3,
'gpu_time': result.gpu_time,
'data_time': result.data_time,
'loss': result.loss
})
#------------------#
# VALIDATION #
#------------------#
result_val, img_merge = validate(
val_loader, model, epoch,
batch_num) # evaluate on validation set，每次训练完以后都要测试一下
#------------------#
# SAVE BEST MODEL #
#------------------#
is_best = result_val.rmse < best_result.rmse
if is_best:
best_result = result_val
with open(best_txt, 'w') as txtfile:
txtfile.write(
"dataset epoch={}\nbatch epoch={}\npsnr={:.5f}\nmse={:.3f}\nrmse={:.3f}\nabsrel={:.3f}\nlg10={:.3f}\nmae={:.3f}\ndelta1={:.3f}\nt_gpu={:.4f}\n"
.format(epoch, batch_num + 1,
10 * math.log(1 / (best_result.mse), 10),
best_result.mse, best_result.rmse,
best_result.absrel, best_result.lg10,
best_result.mae, best_result.delta1,
best_result.gpu_time))
if img_merge is not None:
img_filename = output_directory + '/comparison_best.png'
utils.save_image(img_merge, img_filename)
utils.save_checkpoint(
{
'args': args,
'epoch': epoch,
'batch_epoch': batch_num,
'arch': args.arch,
'model': model,
'best_result': best_result,
'optimizer': optimizer,
}, is_best, epoch, batch_num, output_directory)
def train(train_loader, model, criterion,smoothloss,photometric_loss,optimizer, epoch):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
iheight, iwidth = 480, 640 # raw image size
fx_rgb = 5.1885790117450188e+02;
fy_rgb = 5.1946961112127485e+02;
cx_rgb = 3.2558244941119034e+02;
cy_rgb = 2.5373616633400465e+02;
new_intrinsics = Intrinsics(304,228,fx_rgb*(250.0/iheight),fy_rgb*(250.0/iheight),155.1, 121.0).cuda()
for i, (batch_data, intrinsics) in enumerate(train_loader):
#input, target = input.cuda(), target.cuda()
batch_data = {key:val.cuda() for key,val in batch_data.items() if val is not None}
oheight, owidth = intrinsics["output_size"]
#new_intrinsics = Intrinsics(owidth,oheight,intrinsics["fx"],intrinsics["fy"],intrinsics["cx"],intrinsics["cy"]).cuda()
target = batch_data['gt']
torch.cuda.synchronize()
data_time = time.time() - end
# compute pred
end = time.time()
#pred = model(input[:,3:,:,:],input[:,:3,:,:]) # (depth,image)
#candidates = {"rgb":input[:,:3,:,:], "d":input[:,3:,:,:], "gt":input[:,3:,:,:]}
#batch_data = {key:val.cuda() for key,val in candidates.items() if val is not None}
pred = model(batch_data) # (depth,image)
#loss = criterion(pred, input[:,3:,:,:]) + 0.01*smoothloss(pred)
photoloss = 0.0
if args.use_pose:
# warp near frame to current frame
#hh, ww = pred.size(2), pred.size(3)
#new_intrinsics = new_intrinsics.scale(hh,ww)
mask = (batch_data['d'] < 1e-3).float()
pred_array = multiscale(pred)
rgb_curr_array = multiscale(batch_data['rgb'])
rgb_near_array = multiscale(batch_data['rgb_near'])
mask_array = multiscale(mask)
num_scales = len(pred_array)
# compute photometric loss at multiple scales
for scale in range(len(pred_array)):
pred_ = pred_array[scale]
rgb_curr_ = rgb_curr_array[scale]
rgb_near_ = rgb_near_array[scale]
mask_ = None
if mask is not None:
mask_ = mask_array[scale]
# compute the corresponding intrinsic parameters
height_, width_ = pred_.size(2), pred_.size(3)
intrinsics_ = new_intrinsics.scale(height_, width_)
warped_ = homography_from(rgb_near_,pred_,batch_data["r_mat"],batch_data["t_vec"],intrinsics_)
#warped = homography_from(batch_data["rgb_near"],pred,batch_data["r_mat"],batch_data["t_vec"],new_intrinsics)
#photoloss = photometric_loss(batch_data["rgb"],warped,mask)
photoloss += photometric_loss(rgb_curr_,warped_,mask_)*(2**(scale-num_scales))
loss = criterion(pred, target) + 0.01*smoothloss(pred) + 0.1*photoloss
optimizer.zero_grad()
loss.backward() # compute gradient and do SGD step
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, batch_data['rgb'].size(0))
end = time.time()
if (i + 1) % args.print_freq == 0:
print('=> output: {}'.format(output_directory))
print('Train Epoch: {0} [{1}/{2}]\t'
'LOSS={loss:.3f} '
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
epoch, i+1, len(train_loader),loss=loss.item(), data_time=data_time,
gpu_time=gpu_time, result=result, average=average_meter.average()))
avg = average_meter.average()
with open(train_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({'mse': avg.mse, 'rmse': avg.rmse, 'absrel': avg.absrel, 'lg10': avg.lg10,
'mae': avg.mae, 'delta1': avg.delta1, 'delta2': avg.delta2, 'delta3': avg.delta3,
'gpu_time': avg.gpu_time, 'data_time': avg.data_time})
def evaluate(params, loader, model, experiment):
print("Testing...")
with experiment.test():
with torch.no_grad():
average = AverageMeter()
img_idxs = np.random.randint(0,
len(loader),
size=min(len(loader), 50))
end = time.time()
for i, (inputs, targets) in enumerate(loader):
inputs, targets = inputs.to(params["device"]), targets.to(
params["device"])
data_time = time.time() - end
# Predict
end = time.time()
outputs = model(inputs)
gpu_time = time.time() - end
result = Result()
result.evaluate(outputs.data, targets.data)
average.update(result, gpu_time, data_time, inputs.size(0))
end = time.time()
# Log images to comet
if i in img_idxs:
img_merged = utils.log_image_to_comet(
inputs[0],
targets[0],
outputs[0],
epoch=0,
id=i,
experiment=experiment,
result=result,
prefix="test")
if params["save_test_images"]:
filename = os.path.join(
params["experiment_dir"],
"comparison_epoch_{}_{}.png".format(
str(params["start_epoch"]),
np.where(img_idxs == i)[0][0]))
utils.save_image(img_merged, filename)
if (i + 1) % params["stats_frequency"] == 0 and i != 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(loader),
gpu_time=gpu_time,
result=result,
average=average.average()))
# Mean validation loss
avg = average.average()
utils.log_comet_metrics(experiment, avg, None)
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if params["save_test_metrics"]:
filename = os.path.join(params["experiment_dir"],
"results.csv")
with open(filename, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'irmse': avg.irmse,
'imae': avg.imae,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
def validate(val_loader, model, epoch, write_to_file=True):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
for i, (input, target) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
##############################################################
## Start of PnP-Depth modification ##
##############################################################
# Original inference
with torch.no_grad():
ori_pred = model.pnp_forward_front(model.pnp_forward_rear(
input)) # equivalent to `ori_pred = model(input)`
# Inference with PnP
sparse_target = input[:, -1:] # NOTE: written for rgbd input
criterion = criteria.MaskedL1Loss().cuda(
) # NOTE: criterion function defined here only for clarity
pnp_iters = 5 # number of iterations
pnp_alpha = 0.01 # update/learning rate
pnp_z = model.pnp_forward_front(input)
for pnp_i in range(pnp_iters):
if pnp_i != 0:
pnp_z = pnp_z - pnp_alpha * torch.sign(pnp_z_grad) # iFGM
pnp_z = Variable(pnp_z, requires_grad=True)
pred = model.pnp_forward_rear(pnp_z)
if pnp_i < pnp_iters - 1:
pnp_loss = criterion(pred, sparse_target)
pnp_z_grad = Grad([pnp_loss], [pnp_z], create_graph=True)[0]
##############################################################
## End of PnP-Depth modification ##
##############################################################
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
# save 8 images for visualization
skip = 50
if args.modality == 'd':
img_merge = None
else:
if args.modality == 'rgb':
rgb = input
elif args.modality == 'rgbd':
rgb = input[:, :3, :, :]
depth = input[:, 3:, :, :]
if i == 0:
if args.modality == 'rgbd':
img_merge = utils.merge_into_row_with_gt(
rgb, depth, target, pred)
else:
img_merge = utils.merge_into_row(rgb, target, pred)
elif (i < 8 * skip) and (i % skip == 0):
if args.modality == 'rgbd':
row = utils.merge_into_row_with_gt(rgb, depth, target,
pred)
else:
row = utils.merge_into_row(rgb, target, pred)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = output_directory + '/comparison_' + str(
epoch) + '.png'
utils.save_image(img_merge, filename)
if (i + 1) % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
def train(epoch):
epoch_loss = 0
epoch_loss_depth = 0
epoch_loss_pos = 0
epoch_loss_temperature = 0
epoch_psnr = 0
epoch_sparsity = 0
epoch_loss_slic = 0.0
epoch_loss_color = 0.0
epoch_loss_pos = 0.0
epoch_start = time.time()
end = time.time()
average_meter = AverageMeter()
# train/eval modes make difference on batch normalization layer
modelME.module.netM.train()
modelME.module.netE.eval()
# setup learning rate decay
lr = opt.lr * (0.5**(epoch // (opt.nEpochs // 5)))
# setup temperature for SSA
temperature = opt.temperature
modelME.module.netM.temperature.fill_(temperature)
# use ADAM for the first 2 epoch, then SGD, to speedup training
if epoch <= 2:
optimizer = optim.Adam(
[{
'params': modelME.module.netM.bias_parameters(),
'lr': opt.lr,
'weight_decay': opt.bias_decay
}, {
'params': modelME.module.netM.weight_parameters(),
'lr': opt.lr,
'weight_decay': opt.weight_decay
}, {
'params': modelME.module.netE.bias_parameters(),
'lr': 0.0,
'weight_decay': opt.bias_decay
}, {
'params': modelME.module.netE.weight_parameters(),
'lr': 0.0,
'weight_decay': opt.weight_decay
}],
lr=lr,
betas=(opt.momentum, opt.beta))
else:
optimizer = optim.SGD(
[{
'params': modelME.module.netM.bias_parameters(),
'lr': opt.lr,
'weight_decay': opt.bias_decay
}, {
'params': modelME.module.netM.weight_parameters(),
'lr': opt.lr,
'weight_decay': opt.weight_decay
}, {
'params': modelME.module.netE.bias_parameters(),
'lr': 0.0,
'weight_decay': opt.bias_decay
}, {
'params': modelME.module.netE.weight_parameters(),
'lr': 0.0,
'weight_decay': opt.weight_decay
}],
lr=lr,
momentum=opt.momentum)
for iteration, batch in enumerate(train_loader, 1):
image_target, depth_target, depth_mask = Variable(batch[0]), Variable(
batch[1]), Variable(batch[2])
image_target_lab = rgb2Lab_torch(
image_target.cuda()) # image in lab color space
image_target_labxy_feat_tensor = build_LABXY_feat(
image_target_lab, train_XY_feat_stack) # B* (3+2 )* H * W
depth_input = depth_target.clone()
if torch.cuda.is_available():
image_target = image_target.cuda()
depth_target = depth_target.cuda()
depth_input = depth_input.cuda()
depth_mask = depth_mask.cuda()
image_target_labxy_feat_tensor = image_target_labxy_feat_tensor.cuda(
)
torch.cuda.synchronize()
data_time = time.time() - end
# Compute prediction
end = time.time()
optimizer.zero_grad()
depth_recon, corrupt_mask_soft, corrupt_mask_binary, pooled_xy_tensor, reconstr_xy_tensor, curr_spixl_map, prob = modelME(
image_target, depth_input, True)
mask_sparsity = corrupt_mask_binary.sum() / (
corrupt_mask_binary.shape[0] * corrupt_mask_binary.shape[1] *
corrupt_mask_binary.shape[2] * corrupt_mask_binary.shape[3])
batch_size_cur = image_target.shape[0]
loss_depth = criterion_depth(depth_recon, depth_target, depth_mask)
loss_mse = criterion_mse(depth_recon, depth_target,
depth_mask) # in [0, 1 range]
loss_slic, loss_color, loss_pos = compute_color_pos_loss(
prob,
image_target_labxy_feat_tensor[:batch_size_cur, :, :, :],
pos_weight=opt.pos_weight,
kernel_size=opt.downsize)
loss_temperature = modelME.module.netM.temperature**2
loss = loss_depth + opt.slic_weight * loss_slic + opt.proj_weight * loss_temperature
loss.backward()
optimizer.step()
torch.cuda.synchronize()
psnr = 10 * log10(1 / loss_mse.data.item())
epoch_loss += loss.data.item()
epoch_loss_depth += loss_depth.data.item()
epoch_loss_pos += loss_pos.data.item()
epoch_loss_temperature += loss_temperature.data.item()
epoch_psnr += psnr
epoch_sparsity += mask_sparsity
epoch_loss_slic += loss_slic.data.item()
epoch_loss_color += loss_color.data.item()
epoch_loss_pos += loss_pos.data.item()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(depth_recon.data, depth_target.data)
average_meter.update(result, gpu_time, data_time, image_target.size(0))
epoch_end = time.time()
print(
"===> Epoch {} Complete: lr: {}, Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Time: {:.4f}"
.format(epoch, lr, epoch_loss / len(train_loader),
epoch_psnr / len(train_loader), (epoch_end - epoch_start)))
log_value('train_loss', epoch_loss / len(train_loader), epoch)
log_value('train_loss_depth', epoch_loss_depth / len(train_loader), epoch)
log_value('train_loss_pos', epoch_loss_pos / len(train_loader), epoch)
log_value('train_loss_temperature',
epoch_loss_temperature / len(train_loader), epoch)
log_value('train_psnr', epoch_psnr / len(train_loader), epoch)
log_value('train_sparsity', epoch_sparsity / len(train_loader), epoch)
log_value('train_loss_slic', epoch_loss_slic / len(train_loader), epoch)
log_value('train_loss_color', epoch_loss_color / len(train_loader), epoch)
log_value('train_loss_pos', epoch_loss_pos / len(train_loader), epoch)
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
epoch,
1,
len(train_loader),
data_time=data_time,
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
with open(train_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'gpu_time': avg.gpu_time,
'data_time': avg.data_time
})
def validate(val_loader, model, epoch):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
skip = len(val_loader) // 8 # save images every skip iters
count_b = 0
count = 0
x = random.randint(0, len(val_loader))
for i, (input, target, label, mask) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
input = torch.squeeze(input, 0)
mask = mask.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred, pred_mask, c1, c2, c3 = model(input)
# pred,c1,c2,c3 = model(input)
b, c, h, w = pred.size()
# temp0 = torch.zeros_like(pred_mask)
# temp1 = torch.ones_like(pred_mask)
# pred_mask2 = torch.where(pred_mask>0.5,temp1,temp0)
#pred[:,1,:,:] = pred[:,1,:,:] * pred_mask2
temp0_2 = torch.zeros_like(c1)
temp1_2 = torch.ones_like(c1)
c1_2 = torch.where(c1 > 0.5, temp1_2, temp0_2)
c2_2 = torch.where(c2 > 0.5, temp1_2, temp0_2)
c3_2 = torch.where(c3 > 0.5, temp1_2, temp0_2)
torch.cuda.synchronize()
gpu_time = time.time() - end
target = torch.squeeze(target, 1)
# measure accuracy and record loss
c1_2 = c1_2.cpu().numpy()
c2_2 = c2_2.cpu().numpy()
c3_2 = c3_2.cpu().numpy()
l = label.numpy()
for k in range(l.shape[0]):
if c1_2[k] == 0 and c2_2[k] == 0 and c3_2[k] == 0 and l[k] == -1:
count_b += 1
if c1_2[k] == 1 and c2_2[k] == 0 and c3_2[k] == 0 and l[k] == 0:
count_b += 1
if c1_2[k] == 0 and c2_2[k] == 1 and c3_2[k] == 0 and l[k] == 1:
count_b += 1
if c1_2[k] == 0 and c2_2[k] == 0 and c3_2[k] == 1 and l[k] == 2:
count_b += 1
count += l.shape[0]
result = Result()
result.evaluate(pred, target, label)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
# save 8 images for visualization
rgb = input
# print(rgb.size(),target.size(),pred.size())
# exit(0)
# if i == x:
# img_merge = utils.merge_into_row(rgb, target, pred, pred_mask2,label)
# filename = output_directory + '/comparison_' + str(epoch) + '.png'
# utils.save_image(img_merge, filename)
if i == 0:
img_merge = utils.merge_into_row(
rgb, target, pred, target,
label) # (rgb, target, pred, pred_mask2,label)
elif (i < 8 * skip) and (i % skip == 0):
row = utils.merge_into_row(rgb, target, pred, target, label)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = output_directory + '/comparison_' + str(epoch) + '.png'
utils.save_image(img_merge, filename)
if (i + 1) % args.print_freq == 0:
print("acc: %f" % (count_b * 1.0 / count))
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'MAE={result.mae:.2f}({average.mae:.2f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print("epoch: %d, acc: %f" % (epoch, count_b * 1.0 / count))
print('\n*\n'
'MAE={average.mae:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
return avg, img_merge
def validate(val_loader, model, epoch, logger):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
skip = len(val_loader) // 9 # save images every skip iters
for i, (input, target) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred = model(input)
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
# save 8 images for visualization
if args.dataset == 'kitti':
rgb = input[0]
pred = pred[0]
target = target[0]
else:
rgb = input
if i == 0:
img_merge = utils.merge_into_row(rgb, target, pred)
elif (i < 8 * skip) and (i % skip == 0):
row = utils.merge_into_row(rgb, target, pred)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = output_directory + '/comparison_' + str(epoch) + '.png'
utils.save_image(img_merge, filename)
if (i + 1) % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'RML={result.absrel:.2f}({average.absrel:.2f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'Rel={average.absrel:.3f}\n'
'Log10={average.lg10:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'Delta2={average.delta2:.3f}\n'
'Delta3={average.delta3:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
logger.add_scalar('Test/rmse', avg.rmse, epoch)
logger.add_scalar('Test/Rel', avg.absrel, epoch)
logger.add_scalar('Test/log10', avg.lg10, epoch)
logger.add_scalar('Test/Delta1', avg.delta1, epoch)
logger.add_scalar('Test/Delta2', avg.delta2, epoch)
logger.add_scalar('Test/Delta3', avg.delta3, epoch)
return avg, img_merge
def validate(val_loader, model, epoch, write_to_file=True):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
for i, (input, target) in enumerate(val_loader):
# io.imshow(np.squeeze(input[:,3:,:,:].cpu().numpy()), interpolation='nearest')
# io.show()
#print(input.size())
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred = model(input)
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
if visualize:
if args.modality == 'd':
fig = plt.figure()
fig.suptitle(
'Error Percentage ' + str(round(result.absrel * 100, 2)) +
' GPU TIME ' + str(round(gpu_time, 2)) + ' FPS ' +
str(round(60.0 / (gpu_time + data_time), 2)),
fontsize=16)
plt.subplot(131)
plt.title("SPARSE (Input)")
plt.axis('off')
plt.imshow(np.squeeze(input.cpu().numpy()),
interpolation='nearest')
plt.subplot(132)
plt.title("TARGET (Ground Truth)")
plt.axis('off')
plt.imshow(np.squeeze(target.cpu().numpy()),
interpolation='nearest')
plt.colorbar(fraction=0.1, pad=0.04)
plt.subplot(133)
plt.title("PREDICTED")
plt.axis('off')
plt.imshow(np.squeeze(pred.cpu().numpy()),
interpolation='nearest')
plt.colorbar(fraction=0.1, pad=0.04)
# plt.waitforbuttonpress(timeout=2)
# plt.close()
plt.waitforbuttonpress()
plt.close()
if args.modality == 'rgbd':
# sparse = np.squeeze(input[:, 3:, :, :].cpu().numpy())
# print(sparse.shape)
# sleep(3)
fig = plt.figure()
fig.suptitle(
'Error Percentage ' + str(round(result.absrel * 100, 2)) +
' GPU TIME ' + str(round(gpu_time, 2)) + ' FPS ' +
str(round(60.0 / (gpu_time + data_time), 2)),
fontsize=16)
rgb1 = 255 * np.transpose(
np.squeeze(input[:, :3, :, :].cpu().numpy()),
(1, 2, 0)) # H, W, C
rgb1 = Image.fromarray(rgb1.astype('uint8'))
plt.subplot(221)
plt.title("RGB (Input)")
plt.axis('off')
plt.imshow(rgb1)
plt.subplot(222)
plt.title("SPARSE (Input)")
plt.axis('off')
plt.imshow(np.squeeze(input[:, 3:, :, :].cpu().numpy()),
interpolation='nearest')
plt.subplot(223)
plt.title("TARGET (Ground Truth)")
plt.axis('off')
plt.imshow(np.squeeze(target.cpu().numpy()),
interpolation='nearest')
plt.colorbar(fraction=0.1, pad=0.04)
plt.subplot(224)
plt.title("PREDICTED")
plt.axis('off')
plt.imshow(np.squeeze(pred.cpu().numpy()),
interpolation='nearest')
plt.colorbar(fraction=0.1, pad=0.04)
# plt.waitforbuttonpress(timeout=2)
# plt.close()
plt.waitforbuttonpress()
plt.close()
#
# save 8 images for visualization
skip = 50
if args.modality == 'd':
img_merge = None
else:
if args.modality == 'rgb':
rgb = input
elif args.modality == 'rgbd':
rgb = input[:, :3, :, :]
depth = input[:, 3:, :, :]
if i == 0:
if args.modality == 'rgbd':
img_merge = utils.merge_into_row_with_gt(
rgb, depth, target, pred)
else:
img_merge = utils.merge_into_row(rgb, target, pred)
elif (i < 8 * skip) and (i % skip == 0):
if args.modality == 'rgbd':
row = utils.merge_into_row_with_gt(rgb, depth, target,
pred)
else:
row = utils.merge_into_row(rgb, target, pred)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = output_directory + '/comparison_' + str(
epoch) + '.png'
utils.save_image(img_merge, filename)
if (i + 1) % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
def evaluate(val_loader, model, epoch, write_to_file=True):
#print(model)
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
for i, (input, target, mask, h5name) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
input_var = torch.autograd.Variable(input)
target_var = torch.autograd.Variable(target)
depth_pred = model(input_var)
result = Result()
output1 = torch.index_select(depth_pred.data, 1,
torch.cuda.LongTensor([0]))
result.evaluate(output1, target)
average_meter.update(result, input.size(0))
if args.modality == 'd':
img_merge = None
else:
if args.modality == 'rgb':
rgb = input
elif args.modality == 'rgbd':
rgb = input[:, :3, :, :]
sparse_d = input[:, -1, :, :]
sparse_depth_np = utils.depth_to_np(sparse_d)
sd_filename = output_directory + '/' + str(
epoch) + '_' + '_' + h5name[0] + '_s_depth.png'
utils.save_image(sparse_depth_np, sd_filename)
rgb_np, depth_np, pred_np = utils.image_3_row(
rgb, target, depth_pred)
rgb_filename = output_directory + '/' + str(
epoch) + '_' + '_' + h5name[0] + '_rgb.png'
utils.save_image(rgb_np, rgb_filename)
depth_filename = output_directory + '/' + str(
epoch) + '_' + '_' + h5name[0] + '_depth.png'
utils.save_image(depth_np, depth_filename)
pred_filename = output_directory + '/' + str(
epoch) + '_' + '_' + h5name[0] + '_pred.png'
utils.save_image(pred_np, pred_filename)
avg = average_meter.average()
print('\n\nnew prediction with data:\n*\n'
'RMSE={average.rmse:.5f}\n'
'MAE={average.mae:.5f}\n'
'Delta1={average.delta1:.5f}\n'
'Delta2={average.delta2:.5f}\n'
'Delta3={average.delta3:.5f}\n'
'REL={average.absrel:.5f}\n'
'Lg10={average.lg10:.5f}\n'.format(average=avg))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3
})
def train(epoch):
epoch_loss = 0
epoch_psnr = 0
epoch_start = time.time()
end = time.time()
average_meter = AverageMeter()
model.module.train()
# Step up learning rate decay
lr = opt.lr * (0.2**(epoch // (opt.nEpochs // 4)))
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
for iteration, batch in enumerate(train_loader, 1):
image_target, depth_target, depth_mask = Variable(batch[0]), Variable(
batch[1]), Variable(batch[2])
batch_size = depth_target.size(0)
depth_input = depth_target.clone()
image_height = image_target.size(2)
image_width = image_target.size(3)
# Corrupt the depth_input image
for i in range(0, batch_size):
n_depth_mask = depth_mask[i, 0, :, :].sum().item()
# Adjust the sampling rate based on the depth_mask
if n_depth_mask != 0:
sample_rate = opt.sample_rate / n_depth_mask * (image_height *
image_width)
sample_rate = np.clip(
sample_rate, 0.0, 1.0
) # some randome augementation can cause some crazy mask
else:
sample_rate = opt.sample_rate
corrupt_mask = np.random.binomial(1, (1 - sample_rate),
(image_height, image_width))
corrupt_mask.astype(np.bool)
corrupt_mask = torch.BoolTensor(corrupt_mask)
depth_input[i, 0, :, :].masked_fill_(corrupt_mask, (0.0))
if torch.cuda.is_available():
depth_input = depth_input.cuda()
depth_target = depth_target.cuda()
image_target = image_target.cuda()
depth_mask = depth_mask.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# Compute prediction
end = time.time()
optimizer.zero_grad()
rgb_sparse_d_input = torch.cat((image_target, depth_input),
1) # white input
depth_prediction = model(rgb_sparse_d_input) # white output
loss_depth = criterion_depth(depth_prediction, depth_target,
depth_mask)
loss_mse = criterion_mse(depth_prediction, depth_target, depth_mask)
psnr = 10 * log10(1 / loss_mse.data.item())
epoch_loss += loss_depth.data.item()
epoch_psnr += psnr
loss_depth.backward()
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(depth_prediction.data, depth_target.data)
average_meter.update(result, gpu_time, data_time, image_target.size(0))
end = time.time()
epoch_end = time.time()
print(
"===> Epoch {} Complete: lr: {}, Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Time: {:.4f}"
.format(epoch, lr, epoch_loss / len(train_loader),
epoch_psnr / len(train_loader), (epoch_end - epoch_start)))
log_value('train_loss', epoch_loss / len(train_loader), epoch)
log_value('train_psnr', epoch_psnr / len(train_loader), epoch)
print('Train Epoch: {0} [{1}/{2}]\t'
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
epoch,
i + 1,
len(train_loader),
data_time=data_time,
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
with open(train_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'gpu_time': avg.gpu_time,
'data_time': avg.data_time
})
def demo(val_loader, model, epoch, write_to_file=True):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
#viz.line([[0.,0.]],[0],win='demo2',opts=dict(title='resluts2',legend=['RMSE','MAE']))
viz.line([[0., 0., 0., 0., 0., 0.]], [0],
win='demo1',
opts=dict(
title='resluts1',
legend=['t_GPU', 'Delta1', 'REL', 'lG10', 'RMSE', 'MAE']))
for i, (input, target) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
# torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred = model(input)
# torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
# step 50 images for visualization
skip = 10
if args.modality == 'rgb':
rgb = input
img_merge = utils.merge_into_row(rgb, target, pred)
#row = utils.merge_into_row(rgb, target, pred)
#img_merge = utils.add_row(img_merge, row)
#filename = output_directory + '/comparison_' + str(epoch) + '.png'
#utils.save_image(img_merge, filename)
#print(img_merge)
if i % skip == 0:
img = np.transpose(img_merge, (2, 0, 1))
img = torch.from_numpy(img)
viz.images(img, win='depth_estimation')
t_GPU = float('{gpu_time:.3f}'.format(gpu_time=gpu_time))
RMSE = float('{result.rmse:.2f}'.format(result=result))
MAE = float('{result.mae:.2f}'.format(result=result))
Delta1 = float('{result.delta1:.3f}'.format(result=result))
REL = float('{result.absrel:.3f}'.format(result=result))
Lg10 = float('{result.lg10:.3f}'.format(result=result))
#print(t_GPU)
#viz.line([[RMSE,MAE]],[i],win='demo2',update='append')
viz.line([[t_GPU, Delta1, REL, Lg10, RMSE, MAE]], [i],
win='demo1',
update='append')
time.sleep(0.2)
avg = average_meter.average()
viz.text('\n*\n'
'RMSE={result.rmse:.2f}({average.rmse:.3f})\n\t'
'MAE={result.mae:.2f}({average.mae:.3f})\n\t'
'Delta1={result.delta1:.3f}({average.delta1:.3f})\n\t'
'REL={result.absrel:.3f}({average.absrel:.3f})\n\t'
'Lg10={result.lg10:.3f}({average.lg10:.3f})\n\t'
't_GPU={gpu_time:.3f}{time:.3f}\n'.format(gpu_time=gpu_time,
average=avg,
time=avg.gpu_time,
result=result))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
errors = dict()
for i, data in enumerate(dataset):
if i >= opt.num_test:
break
model.set_input(data)
model.test()
visuals = model.get_current_visuals()
img_path = model.get_image_paths()
if i % 5 == 0:
print('processing (%04d)-th image... %s' % (i, img_path))
save_images(webpage,
visuals,
img_path,
aspect_ratio=opt.aspect_ratio,
width=opt.display_winsize)
result = model.get_depth_errors()
average_meter.update(result, 0, 0, dataset_size)
errors[i] = result.to_dict()
# save the website
webpage.save()
avg = average_meter.average()
errors['average'] = avg.to_dict()
errorDict = {'errors': errors}
error_dir = os.path.join(opt.results_dir, 'loss_%s.txt' % opt.name)
with open(error_dir, 'w') as file:
file.write(json.dumps(errorDict)) # use `json.loads` to do the revers
print("Finish writing at ", error_dir)
def val_rand(epoch):
avg_loss = 0
avg_psnr = 0
frame_count = 0
epoch_start = time.time()
end = time.time()
average_meter = AverageMeter()
modelME.module.eval()
modelME.module.netM.eval()
modelME.module.netE.eval()
for batch in val_loader:
image_target, depth_target, depth_mask = Variable(batch[0]), Variable(
batch[1]), Variable(batch[2])
batch_size = depth_target.size(0)
depth_input = depth_target.clone()
image_height = image_target.size(2)
image_width = image_target.size(3)
# Corrupt the target image
for i in range(0, batch_size):
n_depth_mask = depth_mask[i, 0, :, :].sum().item()
# Adjust the sampling rate based on the depth_mask
sample_rate = opt.sample_rate / n_depth_mask * (image_height *
image_width)
corrupt_mask = np.random.binomial(1, (1 - sample_rate),
(image_height, image_width))
corrupt_mask.astype(np.bool)
corrupt_mask = torch.BoolTensor(corrupt_mask)
depth_input[i, 0, :, :].masked_fill_(corrupt_mask, (0.0))
if torch.cuda.is_available():
image_target = image_target.cuda()
depth_input = depth_input.cuda()
depth_target = depth_target.cuda()
depth_mask = depth_mask.cuda()
rgb_sparse_d_input = torch.cat((image_target, depth_input), 1)
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
depth_prediction = modelME.module.netE(rgb_sparse_d_input)
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(depth_prediction.data, depth_target.data)
average_meter.update(result, gpu_time, data_time, depth_input.size(0))
end = time.time()
for i in range(0, depth_input.shape[0]):
loss_depth = criterion_depth(depth_prediction[[i]],
depth_target[[i]], depth_mask[[i]])
loss_mse = criterion_mse(depth_prediction[[i]], depth_target[[i]],
depth_mask[[i]])
psnr = 10 * log10(1 / loss_mse.data.item())
avg_loss += loss_depth.data.item()
avg_psnr += psnr
frame_count += 1
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
with open(val_rand_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
epoch_end = time.time()
print(
"===> Epoch {} Random Validation: Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Time: {:.4f}"
.format(epoch, avg_loss / frame_count, avg_psnr / frame_count,
(epoch_end - epoch_start)))
log_value('val_loss_depth_rand', avg_loss / frame_count, epoch)
log_value('val_psnr_rand', avg_psnr / frame_count, epoch)
log_images('depth_input_rand',
reshape_4D_array(depth_input[[0], :, :, :].data.cpu().numpy(),
1),
step=1)
log_images('depth_prediction_rand',
reshape_4D_array(
depth_prediction[[0], :, :, :].data.cpu().numpy(), 1),
step=1)
log_images('depth_target_rand',
reshape_4D_array(depth_target[[0], :, :, :].data.cpu().numpy(),
1),
step=1)
log_images('image_target_rand',
reshape_4D_array(image_target[[0], :, :, :].data.cpu().numpy(),
1),
step=1)
def val(epoch):
avg_loss = 0
avg_loss_depth = 0
avg_loss_pos = 0
avg_loss_temperature = 0
avg_psnr = 0
avg_sparsity = 0
avg_loss_slic = 0.0
avg_loss_color = 0.0
avg_loss_pos = 0.0
frame_count = 0
epoch_start = time.time()
end = time.time()
average_meter = AverageMeter()
modelME.module.eval()
modelME.module.netM.eval()
modelME.module.netE.eval()
for batch in val_loader:
image_target, depth_target, depth_mask = Variable(batch[0]), Variable(
batch[1]), Variable(batch[2])
image_target_lab = rgb2Lab_torch(
image_target.cuda()) # image in lab color space
image_target_labxy_feat_tensor = build_LABXY_feat(
image_target_lab, train_XY_feat_stack) # B* (3+2 )* H * W
depth_input = depth_target.clone()
if torch.cuda.is_available():
image_target = image_target.cuda()
depth_input = depth_input.cuda()
depth_target = depth_target.cuda()
depth_mask = depth_mask.cuda()
image_target_labxy_feat_tensor = image_target_labxy_feat_tensor.cuda(
)
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
depth_recon, corrupt_mask_soft, corrupt_mask_binary, pooled_xy_tensor, reconstr_xy_tensor, curr_spixl_map, prob = modelME(
image_target, depth_input, False)
torch.cuda.synchronize()
mask_sparsity = corrupt_mask_binary.sum() / (
corrupt_mask_binary.shape[0] * corrupt_mask_binary.shape[1] *
corrupt_mask_binary.shape[2] * corrupt_mask_binary.shape[3])
# Generate the corrupted depth image
depth_input = corrupt_mask_binary * depth_input
rgb_sparse_d_input = torch.cat((image_target, depth_input),
1) # white input
with torch.no_grad():
restored_depth = modelME.module.netE(rgb_sparse_d_input)
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(restored_depth.data, depth_target.data)
average_meter.update(result, gpu_time, data_time, depth_input.size(0))
end = time.time()
for i in range(0, depth_input.shape[0]):
loss_depth = criterion_depth(restored_depth[[i]],
depth_target[[i]], depth_mask[[i]])
loss_mse = criterion_mse(restored_depth[[i]], depth_target[[i]],
depth_mask[[i]])
psnr = 10 * log10(1 / loss_mse.data.item())
loss_slic, loss_color, loss_pos = compute_color_pos_loss(
prob[[i]],
image_target_labxy_feat_tensor[[i]],
pos_weight=opt.pos_weight,
kernel_size=opt.downsize)
loss_temperature = modelME.module.netM.temperature**2
loss = loss_depth + opt.slic_weight * loss_slic + opt.proj_weight * loss_temperature
avg_loss += loss.data.item()
avg_loss_depth += loss_depth.data.item()
avg_loss_pos += loss_pos.data.item()
avg_loss_temperature += loss_temperature.data.item()
avg_psnr += psnr
avg_sparsity += mask_sparsity
avg_loss_slic += loss_slic.data.item()
avg_loss_color += loss_color.data.item()
avg_loss_pos += loss_pos.data.item()
frame_count += 1
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
with open(val_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
epoch_end = time.time()
print(
"===> Epoch {} Validation: Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Mask Sparsity: {:.4f}, Time: {:.4f}"
.format(epoch, avg_loss / frame_count, avg_psnr / frame_count,
avg_sparsity / frame_count, (epoch_end - epoch_start)))
log_value('val_loss', avg_loss / frame_count, epoch)
log_value('val_loss_depth', avg_loss_depth / frame_count, epoch)
log_value('val_loss_pos', avg_loss_pos / frame_count, epoch)
log_value('val_loss_temperature', avg_loss_temperature / frame_count,
epoch)
log_value('val_psnr', avg_psnr / frame_count, epoch)
log_value('val_sparsity', avg_sparsity / frame_count, epoch)
log_value('val_loss_slic', avg_loss_slic / frame_count, epoch)
log_value('val_loss_color', avg_loss_color / frame_count, epoch)
log_value('val_loss_pos', avg_loss_pos / frame_count, epoch)
# Draw the last image result
spixl_map = update_spixl_map(val_spixelID[[-1], :, :, :],
prob[[-1], :, :, :]) # 1x1x240x960
ori_sz_spixel_map = F.interpolate(
spixl_map.type(torch.float),
size=(opt.input_img_height, opt.input_img_width),
mode='nearest').type(torch.int) # 1x1x240x960
spix_index_np = ori_sz_spixel_map.squeeze().detach().cpu().numpy(
).transpose(0, 1) #240x960
spix_index_np = spix_index_np.astype(np.int64) #240x960, 1% here
image_rgb_np = image_target[[-1], :, :, :].squeeze().clamp(
0, 1).detach().cpu().numpy().transpose(1, 2, 0)
spixel_bd_image = mark_boundaries(image_rgb_np,
spix_index_np.astype(int),
color=(0, 1, 1))
spixel_viz = spixel_bd_image.astype(np.float32).transpose(2, 0, 1)
spixel_viz = np.expand_dims(spixel_viz, axis=0)
log_images('image_target_spixel', reshape_4D_array(spixel_viz, 1), step=1)
log_images('depth_input',
reshape_4D_array(depth_input[[-1], :, :, :].data.cpu().numpy(),
1),
step=1)
log_images('depth_prediction',
reshape_4D_array(
restored_depth[[-1], :, :, :].data.cpu().numpy(), 1),
step=1)
log_images('depth_target',
reshape_4D_array(depth_target[[-1], :, :, :].data.cpu().numpy(),
1),
step=1)
log_images('image_target',
reshape_4D_array(image_target[[-1], :, :, :].data.cpu().numpy(),
1),
step=1)
log_images('corrupt_mask_binary',
reshape_4D_array(
corrupt_mask_binary[[-1], :, :, :].data.cpu().numpy(), 1),
step=1)
log_images('corrupt_mask_soft',
reshape_4D_array(
corrupt_mask_soft[[-1], :, :, :].data.cpu().numpy(), 1),
step=1)
global LOSS_best
if avg_loss_depth < LOSS_best:
LOSS_best = avg_loss
model_out_path = opt.path_to_save + "/model_best.pth"
torch.save(modelME.module.state_dict(), model_out_path)
print("Checkpoint saved to {}".format(model_out_path))
def validate(val_loader, model, epoch, batch_epoch, write_to_file=True):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
for i, (Y, Y_1_2, Y_1_4, Y_1_8, LR, LR_8, HR,
name) in enumerate(val_loader):
Y = Y.cuda()
Y_1_2 = Y_1_2.cuda()
Y_1_4 = Y_1_4.cuda()
Y_1_8 = Y_1_8.cuda()
LR_8 = LR_8.cuda()
LR = LR.cuda()
HR = HR.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
# print("I am for validation in main 342")
if args.arch == 'VDSR_16':
pred_HR = model(LR)
elif args.arch == 'VDSR_16_2':
pred_HR = model(Y, LR)
elif args.arch == 'VDSR':
pred_HR, residule = model(LR_8, Y)
elif args.arch == 'ResNet_bicubic':
pred_HR, residule = model(LR_8, Y)
elif args.arch == 'resnet50_15_6' or 'resnet50_15_11' or 'resnet50_15_12':
pred_HR = model(Y_1_2, LR)
elif args.arch == 'resnet50_15_2' or 'resnet50_15_3' or 'resnet50_15_5' or 'resnet50_15_8' or 'resnet50_15_9':
pred_HR, residule = model(Y, LR, LR_8)
else:
if config.use_different_size_Y == 1:
pred_HR, pred_thermal0 = model(Y, Y_1_2, Y_1_4, Y_1_8, LR)
else:
pred_HR, pred_thermal = model(Y, LR)
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred_HR, HR)
average_meter.update(result, gpu_time, data_time,
Y.size(0)) #Y.size(0) batch_size
end = time.time()
# save 8 images for visualization,对验证集合生产图片
skip = config.skip
if i == 0:
img_merge = utils.merge_into_row_with_YT(Y, LR, pred_HR, HR)
elif (i < 8 * skip) and (i % skip == 0):
row = utils.merge_into_row_with_YT(Y, LR, pred_HR, HR)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip: #储存最终的图片
filename = output_directory + '/Compair_data_epoch_' + str(
epoch) + '_batch_eopch_' + str(batch_epoch + 1) + '.png'
utils.save_image(img_merge, filename)
print("生成第" + str(batch_epoch + 1) + "图片")
if (i + 1) % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'dataset epoch': epoch,
'batch epoch': batch_num + 1,
'psnr': 10 * math.log(1 / (avg.mse), 10),
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
def test(model):
global name
test_dir = config.test_dir
test_dataset = YT_dataset(test_dir, config, is_train_set=False)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
file_name = config.test_result_name + str('.csv')
test_csv_2 = os.path.join(config.test_output_dir, file_name)
#print(test_csv_2)
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
num = 0
for i, (Y, Y_1_2, Y_1_4, Y_1_8, LR, LR_8, HR,
name) in enumerate(test_loader):
Y = Y.cuda()
LR = LR.cuda()
LR_8 = LR_8.cuda()
HR = HR.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
end = time.time()
with torch.no_grad():
# print("I am for validation in main 342")
num = num + 1
if args.arch == 'VDSR_16':
pred_HR = model(LR)
elif args.arch == 'VDSR_16_2':
pred_HR = model(Y, LR)
elif args.arch == 'VDSR':
pred_HR, residule = model(LR_8, Y)
elif args.arch == 'ResNet_bicubic':
pred_HR, residule = model(LR_8, Y)
elif args.arch == 'resnet50_15_2' or 'resnet50_15_3':
pred_HR, residule = model(Y, LR, LR_8)
else:
pred_HR, thermal0 = model(Y, LR)
gpu_time = time.time() - end
image_pred_HR = trans_norm_tensor_into_T(pred_HR)
#image_pred_HR = trans_norm_tensor_into_T(thermal0)
image_LR = trans_norm_tensor_into_rgb(LR)
image_HR = trans_norm_tensor_into_rgb(HR)
#image_HR=trans_norm_tensor_into_T(HR)
name_HR = str(
config.test_output_dir_HR) + str(num) + str('_HR') + str(".png")
name_pred_HR = str(
config.test_output_dir_pred_HR) + str(num) + str(".png")
name_HR_bicubic_HR_res = str(
config.test_output_dir_res) + str(num) + str(".png")
#name_bicubic_HR = str(config.test_output_dir_bicubic_HR) + str(num) + str(".png")
#cv2.imwrite(name_HR,image_HR[...,::-1])
#print(name_pred_HR)
#cv2.imwrite(name_pred_HR,image_pred_HR[...,::-1])
#cv2.waitKey(0)
torch.cuda.synchronize()
# measure accuracy and record loss
result = Result()
result.evaluate(pred_HR, HR)
average_meter.update(result, gpu_time, data_time,
Y.size(0)) # Y.size(0) batch_size
end = time.time()
# save 8 images for visualization,对验证集合生产图片
dir_small = '/home1/sas/datasets/crop/smallthan7/'
dir_big = '/home1/sas/datasets/crop/bigthan/'
dir_out = '/home/zju231/sas/data/old_val_resnet50_15_1/'
#dir_out='/home1/sas/results_model/pytorch/YT-SD/resnet50_30_thermal0/'
#dir_out=test_dir
name1 = dir_out + str(name[0][:-4]) + str('.png')
cv2.imwrite(name1, image_pred_HR)
name2 = dir_out + str(name[0][:-4]) + str('_HR.png')
#cv2.imwrite(name2, image_HR[..., ::-1])
name3 = dir_out + str(name[0][:-4]) + str('_Y.png')
#cv2.imwrite(name3, np.squeeze(np.array(Y.cpu())) * 255)
name4 = dir_out + str(name[0][:-4]) + str('_LR.png')
#cv2.imwrite(name4, image_LR[..., ::-1])
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'PSNR={result.psnr:.5f}({average.psnr:.5f}) '
'RMSE={result.rmse:.5f}({average.rmse:.5f}) '
'MAE={result.mae:.5f}({average.mae:.5f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(test_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'PSNR={average.psnr:.5f}\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if 1:
with open(test_csv_2, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
writer.writerow({
'psnr': 10 * math.log(1 / (avg.mse), 10),
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
def validation(device,
data_loader,
model,
ord_loss,
output_dir,
epoch,
logger,
PRINT_FREQ,
BETA,
GAMMA,
ORD_NUM=80.0):
avg80 = AverageMeter()
avg50 = AverageMeter()
model.eval()
end = time.time()
skip = 1
img_list = []
evalbar = tqdm(total=len(data_loader))
for i, (_input, _sparse_depth, _dense_depth) in enumerate(data_loader):
_input, _sparse_depth, _dense_depth = _input.to(
device), _sparse_depth.to(device), _dense_depth.to(device)
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
_pred_prob, _pred_label = model(_input)
loss = ord_loss(_pred_prob, _dense_depth)
torch.cuda.synchronize()
gpu_time = time.time() - end
pred_depth = utils.label2depth_sid(_pred_label,
K=ORD_NUM,
alpha=1.0,
beta=BETA,
gamma=GAMMA)
abs_rel, sq_rel, rmse, rmse_log, a1, a2, a3 = compute_errors(
_sparse_depth, pred_depth.to(device))
# measure accuracy and record loss
result80 = Result()
result80.evaluate(pred_depth, _sparse_depth.data, cap=80)
result50 = Result()
result50.evaluate(pred_depth, _sparse_depth.data, cap=50)
avg80.update(result80, gpu_time, data_time, _input.size(0))
avg50.update(result50, gpu_time, data_time, _input.size(0))
end = time.time()
# save images for visualization
if i == 0:
img_merge = utils.merge_into_row(_input, _dense_depth, pred_depth)
elif (i < 8 * skip) and (i % skip == 0):
row = utils.merge_into_row(_input, _dense_depth, pred_depth)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = os.path.join(output_dir,
'eval_{}.png'.format(int(epoch)))
print('save validation figures at {}'.format(filename))
utils.save_image(img_merge, filename)
if (i + 1) % PRINT_FREQ == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'AbsRel={result.absrel:.2f}({average.absrel:.2f}) '
'SqRel={result.sqrel:.2f}({average.sqrel:.2f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'RMSE_log={result.rmse_log:.3f}({average.rmse_log:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result80,
average=avg80.average()))
# update progress bar and show loss
evalbar.set_postfix(
ORD_LOSS=
'{:.2f},RMSE/log= {:.2f}/{:.2f},delta={:.2f}/{:.2f}/{:.2f},AbsRel/SqRe;={:.2f}/{:.2f}'
.format(loss, rmse, rmse_log, a1, a2, a3, abs_rel, sq_rel))
evalbar.update(1)
i = i + 1
print('\n**** CAP=80 ****\n'
'RMSE={average.rmse:.3f}\n'
'RMSE_log={average.rmse_log:.3f}\n'
'AbsRel={average.absrel:.3f}\n'
'SqRel={average.sqrel:.3f}\n'
'Log10={average.lg10:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'Delta2={average.delta2:.3f}\n'
'Delta3={average.delta3:.3f}\n'
'iRMSE={average.irmse:.3f}\n'
'iMAE={average.imae:.3f}\n'
't_GPU={average.gpu_time:.3f}\n'.format(average=avg80.average()))
print('\n**** CAP=50 ****\n'
'RMSE={average.rmse:.3f}\n'
'RMSE_log={average.rmse_log:.3f}\n'
'AbsRel={average.absrel:.3f}\n'
'SqRel={average.sqrel:.3f}\n'
'Log10={average.lg10:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'Delta2={average.delta2:.3f}\n'
'Delta3={average.delta3:.3f}\n'
'iRMSE={average.irmse:.3f}\n'
'iMAE={average.imae:.3f}\n'
't_GPU={average.gpu_time:.3f}\n'.format(average=avg50.average()))
logger.add_scalar('VAL_CAP80/RMSE', avg80.average().rmse, epoch)
logger.add_scalar('VAL_CAP80/RMSE_log', avg80.average().rmse_log, epoch)
logger.add_scalar('VAL_CAP80/AbsRel', avg80.average().absrel, epoch)
logger.add_scalar('VAL_CAP80/SqRel', avg80.average().sqrel, epoch)
logger.add_scalar('VAL_CAP80/Delta1', avg80.average().delta1, epoch)
logger.add_scalar('VAL_CAP80/Delta2', avg80.average().delta2, epoch)
logger.add_scalar('VAL_CAP80/Delta3', avg80.average().delta3, epoch)
logger.add_scalar('VAL_CAP50/RMSE', avg50.average().rmse, epoch)
logger.add_scalar('VAL_CAP50/RMSE_log', avg50.average().rmse_log, epoch)
logger.add_scalar('VAL_CAP50/AbsRel', avg50.average().absrel, epoch)
logger.add_scalar('VAL_CAP50/SqRel', avg50.average().sqrel, epoch)
logger.add_scalar('VAL_CAP50/Delta1', avg50.average().delta1, epoch)
logger.add_scalar('VAL_CAP50/Delta2', avg50.average().delta2, epoch)
logger.add_scalar('VAL_CAP50/Delta3', avg50.average().delta3, epoch)
def validate(val_loader, model, epoch, write_to_file=True):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
eval_file = output_directory + '/evaluation.csv'
f = open(eval_file, "w+")
f.write("Max_Error,Depth,RMSE,GPU_TIME,Number_Of_Frame\r\n")
for i, (input, target) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
# torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred = model(input)
# torch.cuda.synchronize()
gpu_time = time.time() - end
abs_err = (target.data - pred.data).abs().cpu()
max_err_ind = np.unravel_index(np.argmax(abs_err, axis=None),
abs_err.shape)
max_err_depth = target.data[max_err_ind]
max_err = abs_err[max_err_ind]
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
f.write(
f'{max_err},{max_err_depth},{result.rmse:.2f},{gpu_time},{i+1}\r\n'
)
# save 8 images for visualization
skip = 50
if args.modality == 'rgb':
rgb = input
if i == 0:
img_merge = utils.merge_into_row_with_gt(rgb, target, pred,
(target - pred).abs())
elif (i < 8 * skip) and (i % skip == 0):
row = utils.merge_into_row_with_gt(rgb, target, pred,
(target - pred).abs())
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = output_directory + '/comparison_' + str(epoch) + '.png'
utils.save_image(img_merge, filename)
if (i + 1) % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
f.close()
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
def train(params, train_loader, val_loader, model, criterion, optimizer,
scheduler, experiment):
phase_list = ["phase_1"]
current_epoch = 1
mean_val_loss = -1
try:
train_step = int(
np.ceil(params["num_training_examples"] / params["batch_size"]) *
params["start_epoch"])
val_step = int(
np.ceil(params["num_validation_examples"] / params["batch_size"] *
params["start_epoch"]))
for epoch in range(params["num_epochs"] - params["start_epoch"]):
current_epoch = params["start_epoch"] + epoch + 1
epoch_loss = 0.0
running_loss = 0.0
average = AverageMeter()
img_idxs = np.random.randint(0, len(train_loader), size=5)
model.train()
with experiment.train():
for i, (inputs, target) in enumerate(train_loader):
# Send data to GPU
inputs, target = inputs.to(params["device"]), target.to(
params["device"])
# Zero the parameter gradients
optimizer.zero_grad()
prediction = model(inputs)
loss = criterion(prediction, target, phase_list=phase_list)
# Compute loss
loss.backward()
optimizer.step()
# Calculate metrics
result = Result()
result.evaluate(prediction.data, target.data)
average.update(result, 0, 0, inputs.size(0))
epoch_loss += loss.item()
# Log to Comet
utils.log_comet_metrics(experiment,
result,
loss.item(),
step=train_step,
epoch=current_epoch)
train_step += 1
# Log images to Comet
if i in img_idxs:
utils.log_image_to_comet(inputs[0], target[0],
prediction[0], current_epoch,
i, experiment, result,
"train", train_step)
# Print statistics
running_loss += loss.item()
if (i + 1) % params["stats_frequency"] == 0 and i != 0:
print('[%d, %5d] loss: %.3f' %
(current_epoch, i + 1,
running_loss / params["stats_frequency"]))
running_loss = 0.0
# Log epoch metrics to Comet
mean_train_loss = epoch_loss / len(train_loader)
utils.log_comet_metrics(experiment,
average.average(),
mean_train_loss,
prefix="epoch",
step=train_step,
epoch=current_epoch)
# Validation each epoch
epoch_loss = 0.0
average = AverageMeter()
with experiment.validate() and torch.no_grad():
img_idxs = np.random.randint(0, len(val_loader), size=5)
model.eval()
for i, (inputs, target) in enumerate(val_loader):
inputs, target = inputs.to(params["device"]), target.to(
params["device"])
prediction = model(inputs)
loss = criterion(prediction, target, phase_list=phase_list)
# Calculate metrics
result = Result()
result.evaluate(prediction.data, target.data)
average.update(result, 0, 0, inputs.size(0))
epoch_loss += loss.item()
# Log to Comet
utils.log_comet_metrics(experiment,
result,
loss.item(),
step=val_step,
epoch=current_epoch)
val_step += 1
# Log images to Comet
if i in img_idxs:
utils.log_image_to_comet(inputs[0], target[0],
prediction[0], current_epoch,
i, experiment, result, "val",
val_step)
# Log epoch metrics to Comet
mean_val_loss = epoch_loss / len(val_loader)
utils.log_comet_metrics(experiment,
average.average(),
mean_val_loss,
prefix="epoch",
step=val_step,
epoch=current_epoch)
print("Validation Loss [%d]: %.3f" %
(current_epoch, mean_val_loss))
# Update training phase
try:
if current_epoch + 1 == params["loss"]["phase_2"]["start"]:
phase_list.append("phase_2")
if current_epoch + 1 == params["loss"]["phase_3"]["start"]:
phase_list.append("phase_3")
except KeyError:
pass
# Save periodically
if (current_epoch + 1) % params["save_frequency"] == 0:
save_path = utils.get_save_path(current_epoch,
params["experiment_dir"])
utils.save_model(model, optimizer, save_path, current_epoch,
mean_val_loss, params["max_checkpoints"])
experiment.log_model(save_path.split("/")[-1], save_path)
print("Saving new checkpoint")
experiment.log_epoch_end(current_epoch)
scheduler.step()
# Save the final model
save_path = utils.get_save_path(params["num_epochs"],
params["experiment_dir"])
utils.save_model(model, optimizer, save_path, current_epoch,
mean_val_loss, params["max_checkpoints"])
experiment.log_model(save_path.split("/")[-1], save_path)
print("Model saved to ", os.path.abspath(save_path))
except KeyboardInterrupt:
print("Saving model and quitting...")
save_path = utils.get_save_path(current_epoch,
params["experiment_dir"])
utils.save_model(model, optimizer, save_path, current_epoch,
mean_val_loss, params["max_checkpoints"])
experiment.log_model(save_path.split("/")[-1], save_path)
print("Model saved to ", os.path.abspath(save_path))
def validate(val_loader, model, write_to_file=True):
average_meter = AverageMeter()
model.eval() # switch to evaluate mode
end = time.time()
print_freq = 10
for i, (input, target) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred = model(input)
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
# save 8 images for visualization
skip = 10
rgb = input
if i == 0:
import matplotlib.pyplot as plt
plt.imsave('pred.png', np.squeeze(pred.cpu().numpy()))
img_merge = utils.merge_into_row(rgb, target, pred)
elif (i < 8 * skip) and (i % skip == 0):
row = utils.merge_into_row(rgb, target, pred)
img_merge = utils.add_row(img_merge, row)
if (i + 1) % print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
def train(train_loader, model, criterion, optimizer, epoch, logger):
average_meter = AverageMeter()
model.train() # switch to train mode
end = time.time()
batch_num = len(train_loader)
current_step = batch_num * args.batch_size * epoch
for i, (input, target) in enumerate(train_loader):
lr = utils.update_ploy_lr(optimizer, args.lr, current_step,
args.max_iter)
input, target = input.cuda(), target.cuda()
data_time = time.time() - end
current_step += input.data.shape[0]
if current_step == args.max_iter:
logger.close()
print('Iteration finished!')
break
torch.cuda.synchronize()
# compute pred
end = time.time()
pred_d, pred_ord = model(input) # @wx 注意输出
loss = criterion(pred_ord, target)
optimizer.zero_grad()
loss.backward() # compute gradient and do SGD step
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
depth = nyu_dataloader.get_depth_sid(pred_d)
target_dp = nyu_dataloader.get_depth_sid(target)
result.evaluate(depth.data, target_dp.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
if (i + 1) % args.print_freq == 0:
print('=> output: {}'.format(output_directory))
print('Train Epoch: {0} [{1}/{2}]\t'
'learning_rate={lr:.8f} '
't_Data={data_time:.3f}({average.data_time:.3f}) '
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'Loss={loss:.3f} '
'RMSE={result.rmse:.3f}({average.rmse:.3f}) '
'RML={result.absrel:.3f}({average.absrel:.3f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
epoch,
i + 1,
batch_num,
lr=lr,
data_time=data_time,
loss=loss.item(),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
logger.add_scalar('Learning_rate', lr, current_step)
logger.add_scalar('Train/Loss', loss.item(), current_step)
logger.add_scalar('Train/RMSE', result.rmse, current_step)
logger.add_scalar('Train/rml', result.absrel, current_step)
logger.add_scalar('Train/Log10', result.lg10, current_step)
logger.add_scalar('Train/Delta1', result.delta1, current_step)
logger.add_scalar('Train/Delta2', result.delta2, current_step)
logger.add_scalar('Train/Delta3', result.delta3, current_step)
avg = average_meter.average()
def validate(val_loader, model, epoch, write_to_file=True):
average_meter = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for i, (input, target) in enumerate(val_loader):
input, target = input.cuda(), target.cuda()
# torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.no_grad():
pred = model(input)
# torch.cuda.synchronize()
gpu_time = time.time() - end
# measure accuracy and record loss
result = Result()
result.evaluate(pred.data, target.data)
average_meter.update(result, gpu_time, data_time, input.size(0))
end = time.time()
# save 8 images for visualization
skip = 50
if args.modality == 'd':
img_merge = None
else:
if args.modality == 'rgb':
rgb = input
elif args.modality == 'rgbd':
rgb = input[:, :3, :, :]
depth = input[:, 3:, :, :]
if i == 0:
if args.modality == 'rgbd':
img_merge = utils.merge_into_row_with_gt(
rgb, depth, target, pred)
else:
img_merge = utils.merge_into_row(rgb, target, pred)
elif (i < 8 * skip) and (i % skip == 0):
if args.modality == 'rgbd':
row = utils.merge_into_row_with_gt(rgb, depth, target,
pred)
else:
row = utils.merge_into_row(rgb, target, pred)
img_merge = utils.add_row(img_merge, row)
elif i == 8 * skip:
filename = output_directory + '/comparison_' + str(
epoch) + '.png'
utils.save_image(img_merge, filename)
if (i + 1) % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'MAE={result.mae:.2f}({average.mae:.2f})\n\t'
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'REL={result.absrel:.3f}({average.absrel:.3f}) '
'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
i + 1,
len(val_loader),
gpu_time=gpu_time,
result=result,
average=average_meter.average()))
avg = average_meter.average()
print('\n*\n'
'RMSE={average.rmse:.3f}\n'
'MAE={average.mae:.3f}\n'
'Delta1={average.delta1:.3f}\n'
'REL={average.absrel:.3f}\n'
'Lg10={average.lg10:.3f}\n'
't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time))
if write_to_file:
with open(test_csv, 'a') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow({
'mse': avg.mse,
'rmse': avg.rmse,
'absrel': avg.absrel,
'lg10': avg.lg10,
'mae': avg.mae,
'delta1': avg.delta1,
'delta2': avg.delta2,
'delta3': avg.delta3,
'data_time': avg.data_time,
'gpu_time': avg.gpu_time
})
return avg, img_merge
def train_one_epoch(device,
train_loader,
model,
output_dir,
ord_loss,
optimizer,
epoch,
logger,
PRINT_FREQ,
BETA,
GAMMA,
ORD_NUM=80.0):
avg80 = AverageMeter()
model.train() # switch to train mode
iter_per_epoch = len(train_loader)
trainbar = tqdm(total=iter_per_epoch)
end = time.time()
for i, (_input, _sparse_depth, _dense_depth) in enumerate(train_loader):
_input, _sparse_depth, _dense_depth = _input.to(
device), _sparse_depth.to(device), _dense_depth.to(device)
torch.cuda.synchronize()
data_time = time.time() - end
# compute output
end = time.time()
with torch.autograd.detect_anomaly():
_pred_prob, _pred_label = model(_input)
loss = ord_loss(
_pred_prob,
_dense_depth) # calculate ord loss with dense_depth
optimizer.zero_grad()
loss.backward()
optimizer.step()
torch.cuda.synchronize()
gpu_time = time.time() - end
pred_depth = utils.label2depth_sid(_pred_label,
K=ORD_NUM,
alpha=1.0,
beta=BETA,
gamma=GAMMA)
# calculate metrices with ground truth sparse depth
abs_rel, sq_rel, rmse, rmse_log, a1, a2, a3 = compute_errors(
_sparse_depth, pred_depth.to(device))
result80 = Result()
result80.evaluate(pred_depth, _sparse_depth.data, cap=80)
avg80.update(result80, gpu_time, data_time, _input.size(0))
end = time.time()
# update progress bar and show loss
trainbar.set_postfix(
ORD_LOSS=
'{:.2f},RMSE/log= {:.2f}/{:.2f},delta={:.2f}/{:.2f}/{:.2f},AbsRel/SqRe;={:.2f}/{:.2f}'
.format(loss, rmse, rmse_log, a1, a2, a3, abs_rel, sq_rel))
trainbar.update(1)
if (i + 1) % PRINT_FREQ == 0:
print('Test: [{0}/{1}]\t'
't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
'AbsRel={result.absrel:.2f}({average.absrel:.2f}) '
'SqRel={result.sqrel:.2f}({average.sqrel:.2f}) '
'Log10={result.lg10:.3f}({average.lg10:.3f}) '
'RMSE_log={result.rmse_log:.3f}({average.rmse_log:.3f}) '
'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
'Delta2={result.delta2:.3f}({average.delta2:.3f}) '
'Delta3={result.delta3:.3f}({average.delta3:.3f})'.format(
i + 1,
len(train_loader),
gpu_time=gpu_time,
result=result80,
average=avg80.average()))
current_step = int(epoch * iter_per_epoch + i + 1)
img_merge = utils.batch_merge_into_row(_input, _dense_depth.data,
pred_depth)
filename = os.path.join(output_dir,
'step_{}.png'.format(current_step))
utils.save_image(img_merge, filename)
logger.add_scalar('TRAIN/RMSE', avg80.average().rmse, current_step)
logger.add_scalar('TRAIN/RMSE_log',
avg80.average().rmse_log, current_step)
logger.add_scalar('TRAIN/AbsRel',
avg80.average().absrel, current_step)
logger.add_scalar('TRAIN/SqRel',
avg80.average().sqrel, current_step)
logger.add_scalar('TRAIN/Delta1',
avg80.average().delta1, current_step)
logger.add_scalar('TRAIN/Delta2',
avg80.average().delta2, current_step)
logger.add_scalar('TRAIN/Delta3',
avg80.average().delta3, current_step)
# reset average meter
result80 = Result()
avg80 = AverageMeter()
