batch_size=opt.testBatchSize,
shuffle=False)
testing_data_loader_kitti = DataLoader(dataset=test_set_kitti,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
testing_data_loader_eth = DataLoader(dataset=test_set_eth,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
# Summaries
tb_writer = SummaryWriter(log_dir='./result/records/')
print('===> Building model')
model = GANet(opt.max_disp)
criterion = MyLoss2(thresh=3, alpha=2)
if cuda:
model = torch.nn.DataParallel(model).cuda()
optimizer = optim.Adam(model.parameters(), lr=opt.lr, betas=(0.9, 0.999))
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
# optimizer.load_state_dict(checkpoint['optimizer'])
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
def main():
opt = parse_args()
if opt.model == "GANet11":
from models.GANet11 import GANet
elif opt.model == "GANet_deep":
from models.GANet_deep import GANet
else:
raise Exception("No suitable model found...")
cuda = opt.cuda
# cuda = True
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
distributed = num_gpus > 1
if distributed:
assert cuda, "Distributed inference only works with GPUs"
torch.cuda.set_device(opt.local_rank)
torch.distributed.init_process_group(
backend="nccl", init_method="env://"
)
dist.barrier()
# torch.manual_seed(opt.seed)
# if cuda:
# torch.cuda.manual_seed(opt.seed)
# print('===> Loading datasets')
print("===> Building model")
model = GANet(opt.max_disp)
if distributed:
model.to('cuda')
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[opt.local_rank], output_device=[opt.local_rank])
elif cuda:
model = torch.nn.DataParallel(model)
model.to('cuda')
model.eval()
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
model.load_state_dict(checkpoint["state_dict"], strict=False)
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# Inference
file_path = opt.data_path
file_list = opt.test_list
with open(file_list, "r") as f:
filelist = [line.strip() for line in f.readlines()]
if distributed:
filelist = filelist[opt.local_rank::num_gpus]
for current_file in tqdm(filelist):
if opt.kitti2015:
leftname = os.path.join(file_path, "image_2", current_file)
rightname = os.path.join(file_path, "image_3", current_file)
savename = os.path.join(opt.save_path, current_file)
if opt.kitti:
leftname = os.path.join(file_path, "colored_0", current_file)
rightname = os.path.join(file_path, "colored_1", current_file)
savename = os.path.join(opt.save_path, current_file)
if opt.cityscapes:
file_id = current_file.split("_leftImg8bit.png")[0]
leftname = os.path.join(
file_path, "leftImg8bit", file_id + "_leftImg8bit.png"
)
rightname = os.path.join(
file_path, "rightImg8bit", file_id + "_rightImg8bit.png"
)
savename = os.path.join(
opt.save_path, os.path.basename(file_id) + "_Disp16bit.png"
)
test(model, leftname, rightname, savename, opt.crop_height, opt.crop_width, cuda)
raise Exception("No GPU found, please run without --cuda")
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
train_set = get_training_set(opt.data_path, opt.training_list, [opt.crop_height, opt.crop_width], opt.left_right,
opt.kitti, opt.kitti2015, opt.shift)
test_set = get_test_set(opt.data_path, opt.val_list, [256, 512], opt.left_right, opt.kitti, opt.kitti2015)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True,
drop_last=True)
testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)
print('===> Building model')
model = GANet(opt.max_disp)
# skip layers for Hierarchical training
if opt.model == 'MyGANet3' or opt.model == 'MyGANet4' or opt.model == 'MyGANet4_8'\
or opt.model == 'MyGANet4_8_rf_t3'or opt.model == 'MyGANet9_t3':
open_layers = ['edge_refine']
for name, module in model.named_children():
if name in open_layers:
module.train()
for p in module.parameters():
p.requires_grad = True
else:
module.eval()
for p in module.parameters():
p.requires_grad = False
else:
raise Exception("No suitable model found ...")
cuda = opt.cuda
#cuda = True
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
#torch.manual_seed(opt.seed)
#if cuda:
# torch.cuda.manual_seed(opt.seed)
#print('===> Loading datasets')
print('===> Building model')
model = GANet(opt.max_disp)
if cuda:
model = torch.nn.DataParallel(model).cuda()
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
model.load_state_dict(checkpoint['state_dict'], strict=False)
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
def add_noise(img, height, width, rmeans=None, rstdevs=None, mod_savename=None):
def apply_noise(image, noise):
[opt.crop_height, opt.crop_width], opt.left_right,
opt.kitti, opt.kitti2015, opt.shift)
test_set = get_test_set(opt.data_path, opt.val_list, [576, 960],
opt.left_right, opt.kitti, opt.kitti2015)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True,
drop_last=True)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
print('===> Building model')
model = GANet(opt.max_disp)
adaptive_loss = DecoupledAdaptiveLossFunction(num_dims=16,
float_dtype=np.float32,
device='cuda:0')
# criterion = MyLoss2(thresh=3, alpha=2)
if cuda:
model = torch.nn.DataParallel(model).cuda()
model_parameters = list(model.parameters())
loss_parameters = list(adaptive_loss.parameters())
optimizer = optim.Adam(model_parameters + loss_parameters,
lr=opt.lr,
betas=(0.9, 0.999))
if opt.resume:
if os.path.isfile(opt.resume):