default=False,
help='enables CUDA training')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
all_left_img, all_right_img, all_left_disp, test_left_img, test_right_img, test_left_disp = kl2012.dataloader(
args.datapath)
TrainImgLoader = torch.utils.data.DataLoader(kl.myImageFloder(
all_left_img, all_right_img, all_left_disp, True),
batch_size=3,
shuffle=True,
num_workers=8,
drop_last=False)
TestImgLoader = torch.utils.data.DataLoader(kl.myImageFloder(
test_left_img, test_right_img, test_left_disp, False),
batch_size=3,
shuffle=False,
num_workers=4,
drop_last=False)
def init_dataloader(input_args):
batch_size = input_args.batchsize
scale_factor = input_args.maxdisp / 384. # controls training resolution
hrvs_folder = '%s/hrvs/carla-highres/trainingF' % input_args.database
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
hrvs_folder)
loader_carla = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=2)
middlebury_folder = '%s/middlebury/mb-ex-training/trainingF' % input_args.database
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
middlebury_folder)
loader_mb = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=0)
rand_scale = [0.9, 2.4 * scale_factor]
all_left_img, all_right_img, all_left_disp, all_right_disp = lt.dataloader(
'%s/sceneflow/' % input_args.database)
loader_scene = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=rand_scale,
order=2)
# change to trainval when finetuning on KITTI
all_left_img, all_right_img, all_left_disp, _, _, _ = lk15.dataloader(
'%s/kitti15/training/' % input_args.database, split='train')
loader_kitti15 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=rand_scale,
order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader(
'%s/kitti12/training/' % input_args.database)
loader_kitti12 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=rand_scale,
order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader(
'%s/eth3d/' % input_args.database)
loader_eth3d = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=rand_scale,
order=0)
all_left_img, all_right_img, all_left_disp, all_right_disp = lidar_dataloader(
'%s/lidar-hdsm-dataset/' % input_args.database)
loader_lidar = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.5, 1.1 * scale_factor],
rand_bright=[0.8, 1.2],
order=2,
flip_disp_ud=True,
occlusion_size=[10, 25])
data_inuse = torch.utils.data.ConcatDataset(
[loader_carla] * 10 + [loader_mb] * 150 + # 71 pairs
[loader_scene] + # 39K pairs 960x540
[loader_kitti15] + [loader_kitti12] * 24 + [loader_eth3d] * 300 +
[loader_lidar]) # 25K pairs
# airsim ~750
train_dataloader = torch.utils.data.DataLoader(data_inuse,
batch_size=batch_size,
shuffle=True,
num_workers=batch_size,
drop_last=True,
worker_init_fn=_init_fn)
print('%d batches per epoch' % (len(data_inuse) // batch_size))
return train_dataloader
from dataloader import MiddleburyLoader as DA
batch_size = args.batchsize
scale_factor = args.maxdisp / 384. # controls training resolution
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader('%s/carla-highres/trainingF'%args.database)
loader_carla = DA.myImageFloder(all_left_img,all_right_img,all_left_disp,right_disparity=all_right_disp, rand_scale=[0.225,0.6*scale_factor], order=2)
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader('%s/mb-ex-training/trainingF'%args.database) # mb-ex
loader_mb = DA.myImageFloder(all_left_img,all_right_img,all_left_disp,right_disparity=all_right_disp, rand_scale=[0.225,0.6*scale_factor], order=0)
all_left_img, all_right_img, all_left_disp, all_right_disp = lt.dataloader('%s/sceneflow/'%args.database)
loader_scene = DA.myImageFloder(all_left_img,all_right_img,all_left_disp,right_disparity=all_right_disp, rand_scale=[0.9,2.4*scale_factor], order=2)
all_left_img, all_right_img, all_left_disp,_,_,_ = lk15.dataloader('%s/kitti_scene/training/'%args.database,typ='train') # trainval
loader_kitti15 = DA.myImageFloder(all_left_img,all_right_img,all_left_disp, rand_scale=[0.9,2.4*scale_factor], order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader('%s/data_stereo_flow/training/'%args.database)
loader_kitti12 = DA.myImageFloder(all_left_img,all_right_img,all_left_disp, rand_scale=[0.9,2.4*scale_factor], order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader('%s/eth3d/'%args.database)
loader_eth3d = DA.myImageFloder(all_left_img,all_right_img,all_left_disp, rand_scale=[0.9,2.4*scale_factor],order=0)
data_inuse = torch.utils.data.ConcatDataset([loader_carla]*40 + [loader_mb]*500 + [loader_scene] + [loader_kitti15] + [loader_kitti12]*80 + [loader_eth3d]*1000)
TrainImgLoader = torch.utils.data.DataLoader(
data_inuse,
batch_size= batch_size, shuffle= True, num_workers=batch_size, drop_last=True, worker_init_fn=_init_fn)
print('%d batches per epoch'%(len(data_inuse)//batch_size))
def main():
global args
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
train_left_img, train_right_img, train_left_disp, test_left_img, test_right_img, test_left_disp = ls.dataloader(
args.datapath)
TrainImgLoader = torch.utils.data.DataLoader(DA.myImageFloder(
train_left_img, train_right_img, train_left_disp, True),
batch_size=args.train_bsize,
shuffle=True,
num_workers=1,
drop_last=False)
TestImgLoader = torch.utils.data.DataLoader(DA.myImageFloder(
test_left_img, test_right_img, test_left_disp, False),
batch_size=args.test_bsize,
shuffle=False,
num_workers=4,
drop_last=False)
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path)
log = logger.setup_logger(args.save_path + 'training.log')
for key, value in sorted(vars(args).items()):
log.info(str(key) + ':' + str(value))
model = StereoNet(maxdisp=args.maxdisp)
model = nn.DataParallel(model).cuda()
model.apply(weights_init)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=[200],
gamma=args.gamma)
log.info('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
args.start_epoch = 0
if args.resume:
if os.path.isfile(args.resume):
log.info("=> loading checkpoint '{}'".format((args.resume)))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
log.info("=> loaded checkpoint '{}' (epeoch {})".format(
args.resume, checkpoint['epoch']))
else:
log.info("=> no checkpoint found at '{}'".format(args.resume))
log.info("=> will start from scratch.")
else:
log.info("Not Resume")
min_erro = 100000
max_epo = 0
start_full_time = time.time()
for epoch in range(args.start_epoch, args.epoch):
log.info('This is {}-th epoch'.format(epoch))
train(TrainImgLoader, model, optimizer, log, epoch)
scheduler.step()
erro = test(TestImgLoader, model, log)
if erro < min_erro:
max_epo = epoch
min_erro = erro
savefilename = args.save_path + 'finetune_checkpoint_{}.pth'.format(
max_epo)
torch.save({
'epoch': epoch,
'state_dict': model.state_dict()
}, savefilename)
log.info('MIN epoch %d total test erro = %.3f' % (max_epo, min_erro))
log.info('full training time = {: 2f} Hours'.format(
(time.time() - start_full_time) / 3600))
def get_training_dataloader(maxdisp, dataset_folder):
scale_factor = maxdisp / 384.
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
'%s/hrvs/carla-highres/trainingF' % dataset_folder)
loader_carla = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=2)
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
'%s/middlebury/mb-ex-training/trainingF' % dataset_folder) # mb-ex
loader_mb = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=0)
all_left_img, all_right_img, all_left_disp, all_right_disp = lt.dataloader(
'%s/sceneflow/' % dataset_folder)
loader_scene = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=2)
all_left_img, all_right_img, all_left_disp, _, _, _ = lk15.dataloader(
'%s/kitti15/training/' % dataset_folder,
typ='train') # change to trainval when finetuning on KITTI
loader_kitti15 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader(
'%s/kitti12/training/' % dataset_folder)
loader_kitti12 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader(
'%s/eth3d/' % dataset_folder)
loader_eth3d = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp = lld.dataloader(
'%s/lidar_dataset/train' % dataset_folder)
loader_lidar = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.5, 1.25 * scale_factor],
rand_bright=[0.8, 1.2],
order=0)
all_dataloaders = [{
'name': 'lidar',
'dl': loader_lidar,
'count': 1
}, {
'name': 'hrvs',
'dl': loader_carla,
'count': 1
}, {
'name': 'middlebury',
'dl': loader_mb,
'count': 1
}, {
'name': 'sceneflow',
'dl': loader_scene,
'count': 1
}, {
'name': 'kitti12',
'dl': loader_kitti12,
'count': 1
}, {
'name': 'kitti15',
'dl': loader_kitti15,
'count': 1
}, {
'name': 'eth3d',
'dl': loader_eth3d,
'count': 1
}]
max_count = 0
for dataloader in all_dataloaders:
max_count = max(max_count, len(dataloader['dl']))
print('=' * 80)
concat_dataloaders = []
for dataloader in all_dataloaders:
dataloader['count'] = max(1, max_count // len(dataloader['dl']))
concat_dataloaders += [dataloader['dl']] * dataloader['count']
print('{name}: {size} (x{count})'.format(name=dataloader['name'],
size=len(dataloader['dl']),
count=dataloader['count']))
data_inuse = torch.utils.data.ConcatDataset(concat_dataloaders)
print('Total dataset size: {}'.format(len(data_inuse)))
print('=' * 80)
return data_inuse
loader_scene = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=2)
all_left_img, all_right_img, all_left_disp, _, _, _ = lk15.dataloader(
'%s/kitti15/training/' % args.database,
typ='train') # change to trainval when finetuning on KITTI
loader_kitti15 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader(
'%s/kitti12/training/' % args.database)
loader_kitti12 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader('%s/eth3d/' %
args.database)
loader_eth3d = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp = lld.dataloader(