generate_laf_data.py

from __future__ import print_function, division
import argparse
import os
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import time
from datasets import __datasets__
from models import __models__, model_loss
from utils import *
from torch.utils.data import DataLoader
import gc

cudnn.benchmark = True
# os.environ["CUDA_VISIBLE_DEVICES"] = "1"

parser = argparse.ArgumentParser(description='Group-wise Correlation Stereo Network (GwcNet)')
parser.add_argument('--model', default='gwcnet-g', help='select a model structure', choices=__models__.keys())
parser.add_argument('--maxdisp', type=int, default=192, help='maximum disparity') # sclar down disp with H and W
parser.add_argument('--dataset', required=True, help='dataset name', choices=__datasets__.keys())
parser.add_argument('--zoom', type=float, default=1.0, help='scaler for zoom in/out the image')
parser.add_argument('--crop_w', type=int, default=0, help='random crop width')
parser.add_argument('--crop_h', type=int, default=0, help='random crop height')
parser.add_argument('--datapath', required=True, help='data path')
parser.add_argument('--trainlist', required=True, help='training list')
parser.add_argument('--testlist', required=True, help='testing list')

parser.add_argument('--training', action='store_true', help='turn to training mode if presents.')
parser.add_argument('--lr', type=float, default=0.0001, help='base learning rate')
parser.add_argument('--batch_size', type=int, default=2, help='training batch size')
parser.add_argument('--test_batch_size', type=int, default=2, help='testing batch size')
parser.add_argument('--epochs', type=int, required=True, help='number of epochs to train')
parser.add_argument('--lrepochs', type=str, required=True, help='the epochs to decay lr: the downscale rate')
parser.add_argument('--inliers', type=int, default=3, help='how many std to include for inliers')
parser.add_argument('--bin_scale', type=str, default='line', help='how to create the distribution, line or log')
parser.add_argument('--n_bins', type=int, default=11, help='how many bins to create the distribution')
parser.add_argument('--loss_type', type=str, required=True, help='define the componet of loss')
parser.add_argument('--mask', type=str, default='soft', help='type of mask assignment',choices=['soft','hard'])

parser.add_argument('--logdir', required=True, help='the directory to save logs and checkpoints')
parser.add_argument('--loadckpt', help='load the weights from a specific checkpoint')
parser.add_argument('--resume', action='store_true', help='continue training the model')
parser.add_argument('--seed', type=int, default=999, metavar='S', help='random seed (default: 1)')

parser.add_argument('--summary_freq', type=int, default=10, help='the frequency of saving summary')
parser.add_argument('--save_freq', type=int, default=1, help='the frequency of saving checkpoint')
parser.add_argument('--save_test', action='store_true', help='save test outputs if presents.')

# parse arguments, set seeds
args = parser.parse_args()
torch.cuda.empty_cache()
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
os.makedirs(args.logdir, exist_ok=True)

# create summary logger
# print("creating new summary file")
# logger = SummaryWriter(args.logdir)

args.maxdisp = int(args.maxdisp * args.zoom)
# dataset, dataloader
StereoDataset = __datasets__[args.dataset]
train_dataset = StereoDataset(args.datapath, args.trainlist, True, args)
test_dataset = StereoDataset(args.datapath, args.testlist, False, args)
TrainImgLoader = DataLoader(train_dataset, args.batch_size, shuffle=True, num_workers=8, drop_last=True)
TestImgLoader = DataLoader(test_dataset, args.test_batch_size, shuffle=False, num_workers=4, drop_last=False)

# model, optimizer
model = __models__[args.model](args.maxdisp)
model = nn.DataParallel(model)
model.cuda()
optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999))

# load parameters
start_epoch = 0
if args.resume:
    if args.loadckpt:
        print("loading model {}".format(args.loadckpt))
        state_dict = torch.load(args.loadckpt)
        model.load_state_dict(state_dict['model'])
        optimizer.load_state_dict(state_dict['optimizer'],strict=False)
        start_epoch = state_dict['epoch'] + 1
    else:# find all checkpoints file and sort according to epoch id
        all_saved_ckpts = [fn for fn in os.listdir(args.logdir) if fn.endswith(".ckpt")]
        all_saved_ckpts = sorted(all_saved_ckpts, key=lambda x: int(x.split('_')[-1].split('.')[0]))
        # use the latest checkpoint file
        loadckpt = os.path.join(args.logdir, all_saved_ckpts[-1])
        print("loading the lastest model in logdir: {}".format(loadckpt))
        state_dict = torch.load(loadckpt)
        model.load_state_dict(state_dict['model'],strict=False)
        optimizer.load_state_dict(state_dict['optimizer'],strict=False)
        start_epoch = state_dict['epoch'] + 1
elif args.loadckpt:
    # load the checkpoint file specified by args.loadckpt
    print("loading model {}".format(args.loadckpt))
    state_dict = torch.load(args.loadckpt)
    model.load_state_dict(state_dict['model'],strict=False)
print("start at epoch {}".format(start_epoch))


def train():
    for epoch_idx in range(start_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch_idx, args.lr, args.lrepochs)
        # testing
        avg_test_scalars = AverageMeterDict()
        for batch_idx, sample in enumerate(TestImgLoader):
            save_outputs = {}
            global_step = len(TestImgLoader) * epoch_idx + batch_idx
            start_time = time.time()
            do_summary = global_step % args.summary_freq == 0
            losses, scalar_outputs, image_outputs = test_sample(sample, compute_metrics=do_summary)
            avg_test_scalars.update(scalar_outputs)
            if args.save_test:
                if epoch_idx == (args.epochs - 1):
                    save_outputs["imgL"] = image_outputs["imgL"].cpu()
                    save_outputs["disp_est"] = image_outputs["disp_est"][0].cpu()
                    save_outputs["disp_gt"] = image_outputs["disp_gt"].cpu()
                    # save_outputs["uncert_est"] = image_outputs["uncert_est"][0].cpu()
                    save_outputs["cost3"]= image_outputs["cost"][0].cpu()
                    if args.model in ['gwcnet-gcs']:
                        save_outputs["uncert_est"].append(image_outputs["uncert_est"][0].cpu())
            del scalar_outputs, image_outputs
            print('Epoch {}/{}, Iter {}/{}, test loss = {:.3f}, time = {:3f}'.format(epoch_idx, args.epochs,
                                                                             batch_idx,
                                                                                     len(TestImgLoader), losses["loss"],
                                                                                     time.time() - start_time))
            if args.save_test:
                torch.save(save_outputs, "{}/sample_{:0>6}.ckpt".format(args.logdir, batch_idx))
        avg_test_scalars = avg_test_scalars.mean()
        # save_scalars(logger, 'fulltest', avg_test_scalars, len(TrainImgLoader) * (epoch_idx + 1))
        print("avg_test_scalars", avg_test_scalars)
        gc.collect()

# test one sample
@make_nograd_func
def test_sample(sample, compute_metrics=True):
    model.eval()

    imgL, imgR, disp_gt = sample['left'], sample['right'], sample['disparity']
    imgL = imgL.cuda()
    imgR = imgR.cuda()
    disp_gt = disp_gt.cuda()

    output = model(imgL, imgR)
    disp_ests = output['disp']
    costs = output['cost']
    mask  = (disp_gt < args.maxdisp) & (disp_gt > 0)
    losses = model_loss(output, disp_gt, mask, args)

    scalar_outputs = {}
    for key in losses.keys():
        scalar_outputs[key] = losses[key]
    image_outputs = {"disp_est": disp_ests,
                     "disp_gt": disp_gt,
                     "imgL": imgL,
                     "imgR": imgR,
                     "cost": costs
                     }

    if args.model in ['gwcnet-gcs']:
        uncert_ests = output['uncert']
        image_outputs["uncert_est"] = uncert_ests

    scalar_outputs["D1"] = [D1_metric(disp_est, disp_gt, mask) for disp_est in disp_ests]
    scalar_outputs["EPE"] = [EPE_metric(disp_est, disp_gt, mask) for disp_est in disp_ests]
    scalar_outputs["Thres1"] = [Thres_metric(disp_est, disp_gt, mask, 1.0) for disp_est in disp_ests]
    scalar_outputs["Thres2"] = [Thres_metric(disp_est, disp_gt, mask, 2.0) for disp_est in disp_ests]
    scalar_outputs["Thres3"] = [Thres_metric(disp_est, disp_gt, mask, 3.0) for disp_est in disp_ests]

    if compute_metrics:
        image_outputs["errormap"] = [disp_error_image_func()(disp_est, disp_gt) for disp_est in disp_ests]

    return tensor2float(losses), tensor2float(scalar_outputs), image_outputs


if __name__ == '__main__':
    train()