train.py

import argparse
import os
import warnings
from pathlib import Path
from typing import List, Union

import numpy as np
import torch
import torch.distributed as dist
import torchvision.models.optical_flow
import torchvision.prototype.models.depth.stereo
import utils
import visualization

from parsing import make_dataset, make_eval_transform, make_train_transform, VALID_DATASETS
from torch import nn
from torchvision.transforms.functional import get_dimensions, InterpolationMode, resize
from utils.metrics import AVAILABLE_METRICS
from utils.norm import freeze_batch_norm


def make_stereo_flow(flow: Union[torch.Tensor, List[torch.Tensor]], model_out_channels: int) -> torch.Tensor:
    """Helper function to make stereo flow from a given model output"""
    if isinstance(flow, list):
        return [make_stereo_flow(flow_i, model_out_channels) for flow_i in flow]

    B, C, H, W = flow.shape
    # we need to add zero flow if the model outputs 2 channels
    if C == 1 and model_out_channels == 2:
        zero_flow = torch.zeros_like(flow)
        # by convention the flow is X-Y axis, so we need the Y flow last
        flow = torch.cat([flow, zero_flow], dim=1)
    return flow


def make_lr_schedule(args: argparse.Namespace, optimizer: torch.optim.Optimizer) -> np.ndarray:
    """Helper function to return a learning rate scheduler for CRE-stereo"""
    if args.decay_after_steps < args.warmup_steps:
        raise ValueError(f"decay_after_steps: {args.function} must be greater than warmup_steps: {args.warmup_steps}")

    warmup_steps = args.warmup_steps if args.warmup_steps else 0
    flat_lr_steps = args.decay_after_steps - warmup_steps if args.decay_after_steps else 0
    decay_lr_steps = args.total_iterations - flat_lr_steps

    max_lr = args.lr
    min_lr = args.min_lr

    schedulers = []
    milestones = []

    if warmup_steps > 0:
        if args.lr_warmup_method == "linear":
            warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
                optimizer, start_factor=args.lr_warmup_factor, total_iters=warmup_steps
            )
        elif args.lr_warmup_method == "constant":
            warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
                optimizer, factor=args.lr_warmup_factor, total_iters=warmup_steps
            )
        else:
            raise ValueError(f"Unknown lr warmup method {args.lr_warmup_method}")
        schedulers.append(warmup_lr_scheduler)
        milestones.append(warmup_steps)

    if flat_lr_steps > 0:
        flat_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=max_lr, total_iters=flat_lr_steps)
        schedulers.append(flat_lr_scheduler)
        milestones.append(flat_lr_steps + warmup_steps)

    if decay_lr_steps > 0:
        if args.lr_decay_method == "cosine":
            decay_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizer, T_max=decay_lr_steps, eta_min=min_lr
            )
        elif args.lr_decay_method == "linear":
            decay_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
                optimizer, start_factor=max_lr, end_factor=min_lr, total_iters=decay_lr_steps
            )
        elif args.lr_decay_method == "exponential":
            decay_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
                optimizer, gamma=args.lr_decay_gamma, last_epoch=-1
            )
        else:
            raise ValueError(f"Unknown lr decay method {args.lr_decay_method}")
        schedulers.append(decay_lr_scheduler)

    scheduler = torch.optim.lr_scheduler.SequentialLR(optimizer, schedulers, milestones=milestones)
    return scheduler


def shuffle_dataset(dataset):
    """Shuffle the dataset"""
    perm = torch.randperm(len(dataset))
    return torch.utils.data.Subset(dataset, perm)


def resize_dataset_to_n_steps(
    dataset: torch.utils.data.Dataset, dataset_steps: int, samples_per_step: int, args: argparse.Namespace
) -> torch.utils.data.Dataset:
    original_size = len(dataset)
    if args.steps_is_epochs:
        samples_per_step = original_size
    target_size = dataset_steps * samples_per_step

    dataset_copies = []
    n_expands, remainder = divmod(target_size, original_size)
    for idx in range(n_expands):
        dataset_copies.append(dataset)

    if remainder > 0:
        dataset_copies.append(torch.utils.data.Subset(dataset, list(range(remainder))))

    if args.dataset_shuffle:
        dataset_copies = [shuffle_dataset(dataset_copy) for dataset_copy in dataset_copies]

    dataset = torch.utils.data.ConcatDataset(dataset_copies)
    return dataset


def get_train_dataset(dataset_root: str, args: argparse.Namespace) -> torch.utils.data.Dataset:
    datasets = []
    for dataset_name in args.train_datasets:
        transform = make_train_transform(args)
        dataset = make_dataset(dataset_name, dataset_root, transform)
        datasets.append(dataset)

    if len(datasets) == 0:
        raise ValueError("No datasets specified for training")

    samples_per_step = args.world_size * args.batch_size

    for idx, (dataset, steps_per_dataset) in enumerate(zip(datasets, args.dataset_steps)):
        datasets[idx] = resize_dataset_to_n_steps(dataset, steps_per_dataset, samples_per_step, args)

    dataset = torch.utils.data.ConcatDataset(datasets)
    if args.dataset_order_shuffle:
        dataset = shuffle_dataset(dataset)

    print(f"Training dataset: {len(dataset)} samples")
    return dataset


@torch.inference_mode()
def _evaluate(
    model,
    args,
    val_loader,
    *,
    padder_mode,
    print_freq=10,
    writer=None,
    step=None,
    iterations=None,
    batch_size=None,
    header=None,
):
    """Helper function to compute various metrics (epe, etc.) for a model on a given dataset."""
    model.eval()
    header = header or "Test:"
    device = torch.device(args.device)
    metric_logger = utils.MetricLogger(delimiter="  ")

    iterations = iterations or args.recurrent_updates

    logger = utils.MetricLogger()
    for meter_name in args.metrics:
        logger.add_meter(meter_name, fmt="{global_avg:.4f}")
    if "fl-all" not in args.metrics:
        logger.add_meter("fl-all", fmt="{global_avg:.4f}")

    num_processed_samples = 0
    with torch.cuda.amp.autocast(enabled=args.mixed_precision, dtype=torch.float16):
        for blob in metric_logger.log_every(val_loader, print_freq, header):
            image_left, image_right, disp_gt, valid_disp_mask = (x.to(device) for x in blob)
            padder = utils.InputPadder(image_left.shape, mode=padder_mode)
            image_left, image_right = padder.pad(image_left, image_right)

            disp_predictions = model(image_left, image_right, flow_init=None, num_iters=iterations)
            disp_pred = disp_predictions[-1][:, :1, :, :]
            disp_pred = padder.unpad(disp_pred)

            metrics, _ = utils.compute_metrics(disp_pred, disp_gt, valid_disp_mask, metrics=logger.meters.keys())
            num_processed_samples += image_left.shape[0]
            for name in metrics:
                logger.meters[name].update(metrics[name], n=1)

    num_processed_samples = utils.reduce_across_processes(num_processed_samples)

    print("Num_processed_samples: ", num_processed_samples)
    if (
        hasattr(val_loader.dataset, "__len__")
        and len(val_loader.dataset) != num_processed_samples
        and torch.distributed.get_rank() == 0
    ):
        warnings.warn(
            f"Number of processed samples {num_processed_samples} is different"
            f"from the dataset size {len(val_loader.dataset)}. This may happen if"
            "the dataset is not divisible by the batch size. Try lowering the batch size or GPU number for more accurate results."
        )

    if writer is not None and args.rank == 0:
        for meter_name, meter_value in logger.meters.items():
            scalar_name = f"{meter_name} {header}"
            writer.add_scalar(scalar_name, meter_value.avg, step)

    logger.synchronize_between_processes()
    print(header, logger)


def make_eval_loader(dataset_name: str, args: argparse.Namespace) -> torch.utils.data.DataLoader:
    if args.weights:
        weights = torchvision.models.get_weight(args.weights)
        trans = weights.transforms()

        def preprocessing(image_left, image_right, disp, valid_disp_mask):
            C_o, H_o, W_o = get_dimensions(image_left)
            image_left, image_right = trans(image_left, image_right)

            C_t, H_t, W_t = get_dimensions(image_left)
            scale_factor = W_t / W_o

            if disp is not None and not isinstance(disp, torch.Tensor):
                disp = torch.from_numpy(disp)
                if W_t != W_o:
                    disp = resize(disp, (H_t, W_t), mode=InterpolationMode.BILINEAR) * scale_factor
            if valid_disp_mask is not None and not isinstance(valid_disp_mask, torch.Tensor):
                valid_disp_mask = torch.from_numpy(valid_disp_mask)
                if W_t != W_o:
                    valid_disp_mask = resize(valid_disp_mask, (H_t, W_t), mode=InterpolationMode.NEAREST)
            return image_left, image_right, disp, valid_disp_mask

    else:
        preprocessing = make_eval_transform(args)

    val_dataset = make_dataset(dataset_name, args.dataset_root, transforms=preprocessing)
    if args.distributed:
        sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False, drop_last=False)
    else:
        sampler = torch.utils.data.SequentialSampler(val_dataset)

    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        sampler=sampler,
        batch_size=args.batch_size,
        pin_memory=True,
        num_workers=args.workers,
    )

    return val_loader


def evaluate(model, loaders, args, writer=None, step=None):
    for loader_name, loader in loaders.items():
        _evaluate(
            model,
            args,
            loader,
            iterations=args.recurrent_updates,
            padder_mode=args.padder_type,
            header=f"{loader_name} evaluation",
            batch_size=args.batch_size,
            writer=writer,
            step=step,
        )


def run(model, optimizer, scheduler, train_loader, val_loaders, logger, writer, scaler, args):
    device = torch.device(args.device)
    # wrap the loader in a logger
    loader = iter(logger.log_every(train_loader))
    # output channels
    model_out_channels = model.module.output_channels if args.distributed else model.output_channels

    torch.set_num_threads(args.threads)

    sequence_criterion = utils.SequenceLoss(
        gamma=args.gamma,
        max_flow=args.max_disparity,
        exclude_large_flows=args.flow_loss_exclude_large,
    ).to(device)

    if args.consistency_weight:
        consistency_criterion = utils.FlowSequenceConsistencyLoss(
            args.gamma,
            resize_factor=0.25,
            rescale_factor=0.25,
            rescale_mode="bilinear",
        ).to(device)
    else:
        consistency_criterion = None

    if args.psnr_weight:
        psnr_criterion = utils.PSNRLoss().to(device)
    else:
        psnr_criterion = None

    if args.smoothness_weight:
        smoothness_criterion = utils.SmoothnessLoss().to(device)
    else:
        smoothness_criterion = None

    if args.photometric_weight:
        photometric_criterion = utils.FlowPhotoMetricLoss(
            ssim_weight=args.photometric_ssim_weight,
            max_displacement_ratio=args.photometric_max_displacement_ratio,
            ssim_use_padding=False,
        ).to(device)
    else:
        photometric_criterion = None

    for step in range(args.start_step + 1, args.total_iterations + 1):
        data_blob = next(loader)
        optimizer.zero_grad()

        # unpack the data blob
        image_left, image_right, disp_mask, valid_disp_mask = (x.to(device) for x in data_blob)
        with torch.cuda.amp.autocast(enabled=args.mixed_precision, dtype=torch.float16):
            disp_predictions = model(image_left, image_right, flow_init=None, num_iters=args.recurrent_updates)
            # different models have different outputs, make sure we get the right ones for this task
            disp_predictions = make_stereo_flow(disp_predictions, model_out_channels)
            # should the architecture or training loop require it, we have to adjust the disparity mask
            # target to possibly look like an optical flow mask
            disp_mask = make_stereo_flow(disp_mask, model_out_channels)
            # sequence loss on top of the model outputs

        loss = sequence_criterion(disp_predictions, disp_mask, valid_disp_mask) * args.flow_loss_weight

        if args.consistency_weight > 0:
            loss_consistency = consistency_criterion(disp_predictions)
            loss += loss_consistency * args.consistency_weight

        if args.psnr_weight > 0:
            loss_psnr = 0.0
            for pred in disp_predictions:
                # predictions might have 2 channels
                loss_psnr += psnr_criterion(
                    pred * valid_disp_mask.unsqueeze(1),
                    disp_mask * valid_disp_mask.unsqueeze(1),
                ).mean()  # mean the psnr loss over the batch
            loss += loss_psnr / len(disp_predictions) * args.psnr_weight

        if args.photometric_weight > 0:
            loss_photometric = 0.0
            for pred in disp_predictions:
                # predictions might have 1 channel, therefore we need to inpute 0s for the second channel
                if model_out_channels == 1:
                    pred = torch.cat([pred, torch.zeros_like(pred)], dim=1)

                loss_photometric += photometric_criterion(
                    image_left, image_right, pred, valid_disp_mask
                )  # photometric loss already comes out meaned over the batch
            loss += loss_photometric / len(disp_predictions) * args.photometric_weight

        if args.smoothness_weight > 0:
            loss_smoothness = 0.0
            for pred in disp_predictions:
                # predictions might have 2 channels
                loss_smoothness += smoothness_criterion(
                    image_left, pred[:, :1, :, :]
                ).mean()  # mean the smoothness loss over the batch
            loss += loss_smoothness / len(disp_predictions) * args.smoothness_weight

        with torch.no_grad():
            metrics, _ = utils.compute_metrics(
                disp_predictions[-1][:, :1, :, :],  # predictions might have 2 channels
                disp_mask[:, :1, :, :],  # so does the ground truth
                valid_disp_mask,
                args.metrics,
            )

        metrics.pop("fl-all", None)
        logger.update(loss=loss, **metrics)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            if args.clip_grad_norm:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=args.clip_grad_norm)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            if args.clip_grad_norm:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=args.clip_grad_norm)
            optimizer.step()

        scheduler.step()

        if not dist.is_initialized() or dist.get_rank() == 0:
            if writer is not None and step % args.tensorboard_log_frequency == 0:
                # log the loss and metrics to tensorboard

                writer.add_scalar("loss", loss, step)
                for name, value in logger.meters.items():
                    writer.add_scalar(name, value.avg, step)
                # log the images to tensorboard
                pred_grid = visualization.make_training_sample_grid(
                    image_left, image_right, disp_mask, valid_disp_mask, disp_predictions
                )
                writer.add_image("predictions", pred_grid, step, dataformats="HWC")

                # second thing we want to see is how relevant the iterative refinement is
                pred_sequence_grid = visualization.make_disparity_sequence_grid(disp_predictions, disp_mask)
                writer.add_image("sequence", pred_sequence_grid, step, dataformats="HWC")

        if step % args.save_frequency == 0:
            if not args.distributed or args.rank == 0:
                model_without_ddp = (
                    model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model
                )
                checkpoint = {
                    "model": model_without_ddp.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "scheduler": scheduler.state_dict(),
                    "step": step,
                    "args": args,
                }
                os.makedirs(args.checkpoint_dir, exist_ok=True)
                torch.save(checkpoint, Path(args.checkpoint_dir) / f"{args.name}_{step}.pth")
                torch.save(checkpoint, Path(args.checkpoint_dir) / f"{args.name}.pth")

        if step % args.valid_frequency == 0:
            evaluate(model, val_loaders, args, writer, step)
            model.train()
            if args.freeze_batch_norm:
                if isinstance(model, nn.parallel.DistributedDataParallel):
                    freeze_batch_norm(model.module)
                else:
                    freeze_batch_norm(model)

    # one final save at the end
    if not args.distributed or args.rank == 0:
        model_without_ddp = model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model
        checkpoint = {
            "model": model_without_ddp.state_dict(),
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict(),
            "step": step,
            "args": args,
        }
        os.makedirs(args.checkpoint_dir, exist_ok=True)
        torch.save(checkpoint, Path(args.checkpoint_dir) / f"{args.name}_{step}.pth")
        torch.save(checkpoint, Path(args.checkpoint_dir) / f"{args.name}.pth")


def main(args):
    args.total_iterations = sum(args.dataset_steps)

    # initialize DDP setting
    utils.setup_ddp(args)
    print(args)

    args.test_only = args.train_datasets is None

    # set the appropriate devices
    if args.distributed and args.device == "cpu":
        raise ValueError("The device must be cuda if we want to run in distributed mode using torchrun")
    device = torch.device(args.device)

    # select model architecture
    model = torchvision.prototype.models.depth.stereo.__dict__[args.model](weights=args.weights)

    # convert to DDP if need be
    if args.distributed:
        model = model.to(args.gpu)
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        model_without_ddp = model.module
    else:
        model.to(device)
        model_without_ddp = model

    os.makedirs(args.checkpoint_dir, exist_ok=True)

    val_loaders = {name: make_eval_loader(name, args) for name in args.test_datasets}

    # EVAL ONLY configurations
    if args.test_only:
        evaluate(model, val_loaders, args)
        return

    # Sanity check for the parameter count
    print(f"Parameter Count: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

    # Compose the training dataset
    train_dataset = get_train_dataset(args.dataset_root, args)

    # initialize the optimizer
    if args.optimizer == "adam":
        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    elif args.optimizer == "sgd":
        optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, momentum=0.9)
    else:
        raise ValueError(f"Unknown optimizer {args.optimizer}. Please choose between adam and sgd")

    # initialize the learning rate schedule
    scheduler = make_lr_schedule(args, optimizer)

    # load them from checkpoint if needed
    args.start_step = 0
    if args.resume_path is not None:
        checkpoint = torch.load(args.resume_path, map_location="cpu", weights_only=True)
        if "model" in checkpoint:
            # this means the user requested to resume from a training checkpoint
            model_without_ddp.load_state_dict(checkpoint["model"])
            # this means the user wants to continue training from where it was left off
            if args.resume_schedule:
                optimizer.load_state_dict(checkpoint["optimizer"])
                scheduler.load_state_dict(checkpoint["scheduler"])
                args.start_step = checkpoint["step"] + 1
                # modify starting point of the dat
                sample_start_step = args.start_step * args.batch_size * args.world_size
                train_dataset = train_dataset[sample_start_step:]

        else:
            # this means the user wants to finetune on top of a model state dict
            # and that no other changes are required
            model_without_ddp.load_state_dict(checkpoint)

    torch.backends.cudnn.benchmark = True

    # enable training mode
    model.train()
    if args.freeze_batch_norm:
        freeze_batch_norm(model_without_ddp)

    # put dataloader on top of the dataset
    # make sure to disable shuffling since the dataset is already shuffled
    # in order to guarantee quasi randomness whilst retaining a deterministic
    # dataset consumption order
    if args.distributed:
        # the train dataset is preshuffled in order to respect the iteration order
        sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=False, drop_last=True)
    else:
        # the train dataset is already shuffled, so we can use a simple SequentialSampler
        sampler = torch.utils.data.SequentialSampler(train_dataset)

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        sampler=sampler,
        batch_size=args.batch_size,
        pin_memory=True,
        num_workers=args.workers,
    )

    # initialize the logger
    if args.tensorboard_summaries:
        from torch.utils.tensorboard import SummaryWriter

        tensorboard_path = Path(args.checkpoint_dir) / "tensorboard"
        os.makedirs(tensorboard_path, exist_ok=True)

        tensorboard_run = tensorboard_path / f"{args.name}"
        writer = SummaryWriter(tensorboard_run)
    else:
        writer = None

    logger = utils.MetricLogger(delimiter="  ")

    scaler = torch.cuda.amp.GradScaler() if args.mixed_precision else None
    # run the training loop
    # this will perform optimization, respectively logging and saving checkpoints
    # when need be
    run(
        model=model,
        optimizer=optimizer,
        scheduler=scheduler,
        train_loader=train_loader,
        val_loaders=val_loaders,
        logger=logger,
        writer=writer,
        scaler=scaler,
        args=args,
    )


def get_args_parser(add_help=True):
    import argparse

    parser = argparse.ArgumentParser(description="PyTorch Stereo Matching Training", add_help=add_help)
    # checkpointing
    parser.add_argument("--name", default="crestereo", help="name of the experiment")
    parser.add_argument("--resume", type=str, default=None, help="from which checkpoint to resume")
    parser.add_argument("--checkpoint-dir", type=str, default="checkpoints", help="path to the checkpoint directory")

    # dataset
    parser.add_argument("--dataset-root", type=str, default="", help="path to the dataset root directory")
    parser.add_argument(
        "--train-datasets",
        type=str,
        nargs="+",
        default=["crestereo"],
        help="dataset(s) to train on",
        choices=list(VALID_DATASETS.keys()),
    )
    parser.add_argument(
        "--dataset-steps", type=int, nargs="+", default=[300_000], help="number of steps for each dataset"
    )
    parser.add_argument(
        "--steps-is-epochs", action="store_true", help="if set, dataset-steps are interpreted as epochs"
    )
    parser.add_argument(
        "--test-datasets",
        type=str,
        nargs="+",
        default=["middlebury2014-train"],
        help="dataset(s) to test on",
        choices=["middlebury2014-train"],
    )
    parser.add_argument("--dataset-shuffle", type=bool, help="shuffle the dataset", default=True)
    parser.add_argument("--dataset-order-shuffle", type=bool, help="shuffle the dataset order", default=True)
    parser.add_argument("--batch-size", type=int, default=2, help="batch size per GPU")
    parser.add_argument("--workers", type=int, default=4, help="number of workers per GPU")
    parser.add_argument(
        "--threads",
        type=int,
        default=16,
        help="number of CPU threads per GPU. This can be changed around to speed-up transforms if needed. This can lead to worker thread contention so use with care.",
    )

    # model architecture
    parser.add_argument(
        "--model",
        type=str,
        default="crestereo_base",
        help="model architecture",
        choices=["crestereo_base", "raft_stereo"],
    )
    parser.add_argument("--recurrent-updates", type=int, default=10, help="number of recurrent updates")
    parser.add_argument("--freeze-batch-norm", action="store_true", help="freeze batch norm parameters")

    # loss parameters
    parser.add_argument("--gamma", type=float, default=0.8, help="gamma parameter for the flow sequence loss")
    parser.add_argument("--flow-loss-weight", type=float, default=1.0, help="weight for the flow loss")
    parser.add_argument(
        "--flow-loss-exclude-large",
        action="store_true",
        help="exclude large flow values from the loss. A large value is defined as a value greater than the ground truth flow norm",
        default=False,
    )
    parser.add_argument("--consistency-weight", type=float, default=0.0, help="consistency loss weight")
    parser.add_argument(
        "--consistency-resize-factor",
        type=float,
        default=0.25,
        help="consistency loss resize factor to account for the fact that the flow is computed on a downsampled image",
    )
    parser.add_argument("--psnr-weight", type=float, default=0.0, help="psnr loss weight")
    parser.add_argument("--smoothness-weight", type=float, default=0.0, help="smoothness loss weight")
    parser.add_argument("--photometric-weight", type=float, default=0.0, help="photometric loss weight")
    parser.add_argument(
        "--photometric-max-displacement-ratio",
        type=float,
        default=0.15,
        help="Only pixels with a displacement smaller than this ratio of the image width will be considered for the photometric loss",
    )
    parser.add_argument("--photometric-ssim-weight", type=float, default=0.85, help="photometric ssim loss weight")

    # transforms parameters
    parser.add_argument("--gpu-transforms", action="store_true", help="use GPU transforms")
    parser.add_argument(
        "--eval-size", type=int, nargs="+", default=[384, 512], help="size of the images for evaluation"
    )
    parser.add_argument("--resize-size", type=int, nargs=2, default=None, help="resize size")
    parser.add_argument("--crop-size", type=int, nargs=2, default=[384, 512], help="crop size")
    parser.add_argument("--scale-range", type=float, nargs=2, default=[0.6, 1.0], help="random scale range")
    parser.add_argument("--rescale-prob", type=float, default=1.0, help="probability of resizing the image")
    parser.add_argument(
        "--scaling-type", type=str, default="linear", help="scaling type", choices=["exponential", "linear"]
    )
    parser.add_argument("--flip-prob", type=float, default=0.5, help="probability of flipping the image")
    parser.add_argument(
        "--norm-mean", type=float, nargs="+", default=[0.5, 0.5, 0.5], help="mean for image normalization"
    )
    parser.add_argument(
        "--norm-std", type=float, nargs="+", default=[0.5, 0.5, 0.5], help="std for image normalization"
    )
    parser.add_argument(
        "--use-grayscale", action="store_true", help="use grayscale images instead of RGB", default=False
    )
    parser.add_argument("--max-disparity", type=float, default=None, help="maximum disparity")
    parser.add_argument(
        "--interpolation-strategy",
        type=str,
        default="bilinear",
        help="interpolation strategy",
        choices=["bilinear", "bicubic", "mixed"],
    )
    parser.add_argument("--spatial-shift-prob", type=float, default=1.0, help="probability of shifting the image")
    parser.add_argument(
        "--spatial-shift-max-angle", type=float, default=0.1, help="maximum angle for the spatial shift"
    )
    parser.add_argument(
        "--spatial-shift-max-displacement", type=float, default=2.0, help="maximum displacement for the spatial shift"
    )
    parser.add_argument("--gamma-range", type=float, nargs="+", default=[0.8, 1.2], help="range for gamma correction")
    parser.add_argument(
        "--brightness-range", type=float, nargs="+", default=[0.8, 1.2], help="range for brightness correction"
    )
    parser.add_argument(
        "--contrast-range", type=float, nargs="+", default=[0.8, 1.2], help="range for contrast correction"
    )
    parser.add_argument(
        "--saturation-range", type=float, nargs="+", default=0.0, help="range for saturation correction"
    )
    parser.add_argument("--hue-range", type=float, nargs="+", default=0.0, help="range for hue correction")
    parser.add_argument(
        "--asymmetric-jitter-prob",
        type=float,
        default=1.0,
        help="probability of using asymmetric jitter instead of symmetric jitter",
    )
    parser.add_argument("--occlusion-prob", type=float, default=0.5, help="probability of occluding the rightimage")
    parser.add_argument(
        "--occlusion-px-range", type=int, nargs="+", default=[50, 100], help="range for the number of occluded pixels"
    )
    parser.add_argument("--erase-prob", type=float, default=0.0, help="probability of erasing in both images")
    parser.add_argument(
        "--erase-px-range", type=int, nargs="+", default=[50, 100], help="range for the number of erased pixels"
    )
    parser.add_argument(
        "--erase-num-repeats", type=int, default=1, help="number of times to repeat the erase operation"
    )

    # optimizer parameters
    parser.add_argument("--optimizer", type=str, default="adam", help="optimizer", choices=["adam", "sgd"])
    parser.add_argument("--lr", type=float, default=4e-4, help="learning rate")
    parser.add_argument("--weight-decay", type=float, default=0.0, help="weight decay")
    parser.add_argument("--clip-grad-norm", type=float, default=0.0, help="clip grad norm")

    # lr_scheduler parameters
    parser.add_argument("--min-lr", type=float, default=2e-5, help="minimum learning rate")
    parser.add_argument("--warmup-steps", type=int, default=6_000, help="number of warmup steps")
    parser.add_argument(
        "--decay-after-steps", type=int, default=180_000, help="number of steps after which to start decay the lr"
    )
    parser.add_argument(
        "--lr-warmup-method", type=str, default="linear", help="warmup method", choices=["linear", "cosine"]
    )
    parser.add_argument("--lr-warmup-factor", type=float, default=0.02, help="warmup factor for the learning rate")
    parser.add_argument(
        "--lr-decay-method",
        type=str,
        default="linear",
        help="decay method",
        choices=["linear", "cosine", "exponential"],
    )
    parser.add_argument("--lr-decay-gamma", type=float, default=0.8, help="decay factor for the learning rate")

    # deterministic behaviour
    parser.add_argument("--seed", type=int, default=42, help="seed for random number generators")

    # mixed precision training
    parser.add_argument("--mixed-precision", action="store_true", help="use mixed precision training")

    # logging
    parser.add_argument("--tensorboard-summaries", action="store_true", help="log to tensorboard")
    parser.add_argument("--tensorboard-log-frequency", type=int, default=100, help="log frequency")
    parser.add_argument("--save-frequency", type=int, default=1_000, help="save frequency")
    parser.add_argument("--valid-frequency", type=int, default=1_000, help="validation frequency")
    parser.add_argument(
        "--metrics",
        type=str,
        nargs="+",
        default=["mae", "rmse", "1px", "3px", "5px", "relepe"],
        help="metrics to log",
        choices=AVAILABLE_METRICS,
    )

    # distributed parameters
    parser.add_argument("--world-size", type=int, default=8, help="number of distributed processes")
    parser.add_argument("--dist-url", type=str, default="env://", help="url used to set up distributed training")
    parser.add_argument("--device", type=str, default="cuda", help="device to use for training")

    # weights API
    parser.add_argument("--weights", type=str, default=None, help="weights API url")
    parser.add_argument(
        "--resume-path", type=str, default=None, help="a path from which to resume or start fine-tuning"
    )
    parser.add_argument("--resume-schedule", action="store_true", help="resume optimizer state")

    # padder parameters
    parser.add_argument("--padder-type", type=str, default="kitti", help="padder type", choices=["kitti", "sintel"])
    return parser


if __name__ == "__main__":
    args = get_args_parser().parse_args()
    main(args)