# handling GPU/CPU
cuda = torch.cuda.is_available() and not args.cpu
if cuda:
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
print("\n" + "=> using '{}' for computation.".format(device))
# define loss functions
depth_criterion = criteria.MaskedMSELoss() if (
args.criterion == 'l2') else criteria.MaskedL1Loss()
photometric_criterion = criteria.PhotometricLoss()
smoothness_criterion = criteria.SmoothnessLoss()
if args.use_pose:
# hard-coded KITTI camera intrinsics
K = load_calib()
fu, fv = float(K[0, 0]), float(K[1, 1])
cu, cv = float(K[0, 2]), float(K[1, 2])
kitti_intrinsics = Intrinsics(owidth, oheight, fu, fv, cu, cv)
if cuda:
kitti_intrinsics = kitti_intrinsics.cuda()
def iterate(mode, args, loader, model, optimizer, logger, epoch):
block_average_meter = AverageMeter()
average_meter = AverageMeter()
meters = [block_average_meter, average_meter]
def main():
global args, best_result, output_directory, train_csv, test_csv
# evaluation mode
start_epoch = 0
if args.evaluate:
assert os.path.isfile(args.evaluate), \
"=> no best model found at '{}'".format(args.evaluate)
print("=> loading best model '{}'".format(args.evaluate))
checkpoint = torch.load(args.evaluate)
output_directory = os.path.dirname(args.evaluate)
args = checkpoint['args']
start_epoch = checkpoint['epoch'] + 1
best_result = checkpoint['best_result']
model = checkpoint['model']
print("=> loaded best model (epoch {})".format(checkpoint['epoch']))
_, val_loader = create_data_loaders(args)
args.evaluate = True
validate(val_loader, model, checkpoint['epoch'], write_to_file=False)
return
# optionally resume from a checkpoint
elif args.resume:
chkpt_path = args.resume
assert os.path.isfile(chkpt_path), \
"=> no checkpoint found at '{}'".format(chkpt_path)
print("=> loading checkpoint '{}'".format(chkpt_path))
checkpoint = torch.load(chkpt_path)
args = checkpoint['args']
start_epoch = checkpoint['epoch'] + 1
best_result = checkpoint['best_result']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
output_directory = os.path.dirname(os.path.abspath(chkpt_path))
print("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))
train_loader, val_loader = create_data_loaders(args)
args.resume = True
# create new model
else:
train_loader, val_loader = create_data_loaders(args)
print("=> creating Model ({}-{}) ...".format(args.arch, args.decoder))
in_channels = len(args.modality)
if args.arch == 'resnet50':
model = ResNet(layers=50, decoder=args.decoder, output_size=train_loader.dataset.output_size,
in_channels=in_channels, pretrained=args.pretrained)
elif args.arch == 'SmallNet':
model = SmallNet().cuda()
model_named_params = [p for _,p in model.named_parameters() if p.requires_grad]
elif args.arch == 'UNET':
model = DepthCompletionNet(args).cuda()
model_named_params = [p for _,p in model.named_parameters() if p.requires_grad]
elif args.arch == 'DRNSeg':
model = DRNSeg("drn_d_22", 1, pretrained_model=None,pretrained=False)
model_named_params = [p for _,p in model.named_parameters() if p.requires_grad]
elif args.arch == 'ERF':
model = ERF().cuda()
model_named_params = [p for _,p in model.named_parameters() if p.requires_grad]
elif args.arch == 'resnet18':
model = ResNet(layers=18, decoder=args.decoder, output_size=train_loader.dataset.output_size,
in_channels=in_channels, pretrained=args.pretrained)
print("=> model created.")
#optimizer = torch.optim.SGD(model.parameters(), args.lr, \
# momentum=args.momentum, weight_decay=args.weight_decay)
optimizer = torch.optim.Adam(model_named_params, lr=args.lr, weight_decay=args.weight_decay)
# model = torch.nn.DataParallel(model).cuda() # for multi-gpu training
model = model.cuda()
# define loss function (criterion) and optimizer
if args.criterion == 'l2':
criterion = criteria.MaskedMSELoss().cuda()
elif args.criterion == 'l1':
criterion = criteria.MaskedL1Loss().cuda()
smoothloss = criteria.SmoothnessLoss().cuda()
photometric_loss = criteria.PhotometricLoss().cuda()
# create results folder, if not already exists
output_directory = utils.get_output_directory(args)
if not os.path.exists(output_directory):
os.makedirs(output_directory)
train_csv = os.path.join(output_directory, 'train.csv')
test_csv = os.path.join(output_directory, 'test.csv')
best_txt = os.path.join(output_directory, 'best.txt')
# create new csv files with only header
if not args.resume:
with open(train_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
with open(test_csv, 'w') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
for epoch in range(start_epoch, args.epochs):
utils.adjust_learning_rate(optimizer, epoch, args.lr)
train(train_loader, model, criterion,smoothloss, photometric_loss, optimizer, epoch) # train for one epoch
result, img_merge = validate(val_loader, model, epoch) # evaluate on validation set
# remember best rmse and save checkpoint
is_best = result.rmse < best_result.rmse
if is_best:
best_result = result
with open(best_txt, 'w') as txtfile:
txtfile.write("epoch={}\nmse={:.3f}\nrmse={:.3f}\nabsrel={:.3f}\nlg10={:.3f}\nmae={:.3f}\ndelta1={:.3f}\nt_gpu={:.4f}\n".
format(epoch, result.mse, result.rmse, result.absrel, result.lg10, result.mae, result.delta1, 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,
'arch': args.arch,
'model': model,
'best_result': best_result,
'optimizer' : optimizer,
}, is_best, epoch, output_directory)
