BaseModel

Abstract model to be used as a parent of concrete optical flow networks.

In order to use this module, the concrete model should just need to: 1. Specify the network structure in their __init__() method and call super().__init__() with the required args. 2. Implement a forward method which receives inputs and outputs as specified. See the forward method docs for more details. 3. Optionally, define a loss function, if the model is able to be trained.

class ptlflow.models.base_model.base_model.BaseModel(args: Namespace, loss_fn: Callable, output_stride: int)

A base abstract optical flow model.

__init__(args: Namespace, loss_fn: Callable, output_stride: int) → None

Initialize BaseModel.

Parameters:

argsNamespace

A namespace with the required arguments. Typically, this can be gotten from add_model_specific_args().

loss_fnCallable

A function to be used to compute the loss for the training. The input of this function must match the output of the forward() method. The output of this function must be a tensor with a single value.

output_strideint

How many times the output of the network is smaller than the input.

static add_model_specific_args(parent_parser: ArgumentParser | None = None) → ArgumentParser

Generate a parser for the arguments required by this model.

Parameters:

parent_parserArgumentParser

An existing parser, to be extended with the arguments from this model.

Returns:

ArgumentParser

The parser after extending its arguments.

Notes

If the concrete model needs to add more arguments than these defined in this BaseModel, then it should create its own method defined as follows:

>>>
@staticmethod
def add_model_specific_args(parent_parser):
    parent_parser = BaseModel.add_model_specific_args(parent_parser)
    parser = ArgumentParser(parents=[parent_parser], add_help=False)
    parser.add_argument(...add new arguments...)
    return parser

configure_optimizers() → Dict[str, Any]

Initialize the optimizers and LR schedulers.

This function is called internally by Pytorch-Lightning at the beginning of the training.

Returns:

Dict[str, Any]

A dict with two keys: - ‘optimizer’: an optimizer from PyTorch. - ‘lr_scheduler’: Dict[‘str’, Any], a dict with the selected scheduler and its required arguments.

abstract forward(*args: Any, **kwargs: Any) → Dict[str, Tensor]

Forward the inputs through the network and produce the predictions.

The method inputs can be anything, up to the implementation of the concrete model. However, the recommended input is to receive only dict[str, Any] as argument. This dict should contain everything required for one pass of the network (images, etc.). Arguments which do not change for each forward should be defined as arguments in the parser (see add_model_specific_args()).

Parameters:

argsAny

Any arguments.

kwargsAny

Any named arguments.

Returns:

Dict[str, torch.Tensor]

For compatibility with the framework, the output should be a dict containing at least the following keys:

• ‘flows’: a 5D tensor BNCHW containing the predicted flows. The flows must be at the original scale, i.e., their size is the same as the input images, and the flow magnitudes are scaled accordingly. Most networks only produce a single 2D optical flow per batch element, so the output shape will be B12HW. N can be larger than one if the network produces predictions for a larger temporal window.

• ‘occs’: optional, and only included if the network also predicts occlusion masks. It is a 5D tensor following the same structure as ‘flows’.

• ‘mbs’: same as ‘occs’ but for occlusion masks.

• ‘confs’: same as ‘occs’ but for flow confidence predictions.

The keys above are used by other parts of PTLFlow (if available). The output can also include any other keys as well. For example, by default, the output of forward() will be passed to the loss function. So the output may include keys which are going to be used for computing the training loss.

parse_dataset_selection(dataset_selection: str) → List[Tuple[str, int]]

Parse the input string into the selected dataset and their multipliers and parameters.

For example, ‘chairs-train+3*sintel-clean-trainval+kitti-2012-train*5’ will be parsed into [(1, ‘chairs’, ‘train’), (3, ‘sintel’, ‘clean’, ‘trainval’), (5, ‘kitti’, ‘2012’, ‘train’)].

Parameters:

dataset_selectionstr

The string defining the dataset selection. Each dataset is separated by a ‘+’ sign. The multiplier must be either in the beginning or the end of one dataset string, connected to a ‘*’ sign. The remaining content must be separated by ‘-’ symbols.

Returns:

List[Tuple[str, int]]

The parsed choice of datasets and their number of repetitions.

Raises:

ValueError

If the given string is invalid.

postprocess_predictions(prediction: Tensor, image_resizer: InputPadder | InputScaler | None, is_flow: bool) → Tensor

Simple resizing post-processing. Just use image_resizer to revert the resizing operations.

Parameters:

predictiontorch.Tensor

A tensor with at least 3 dimensions in this order: […, C, H, W].

image_resizerOptional[Union[InputPadder, InputScaler]]

An instance of InputPadder or InputScaler that will be used to reverse the resizing done to the inputs. Typically, this will be the instance returned by self.preprocess_images().

is_flowbool

Indicates if prediction is an optical flow prediction of not. Only used if image_resizer is an instance of InputScaler, in which case the flow values need to be scaled.

Returns:

torch.Tensor

A copy of the prediction after reversing the resizing.

preprocess_images(images: Tensor, bgr_add: float | Tuple[float, float, float] | ndarray | Tensor = 0, bgr_mult: float | Tuple[float, float, float] | ndarray | Tensor = 1, bgr_to_rgb: bool = False, image_resizer: InputPadder | InputScaler | None = None, resize_mode: str = 'pad', target_size: Tuple[int, int] | None = None, pad_mode: str = 'replicate', pad_value: float = 0.0, pad_two_side: bool = True, interpolation_mode: str = 'bilinear', interpolation_align_corners: bool = True) → Tuple[Tensor, InputPadder | InputScaler]

Applies basic pre-processing to the images.

The pre-processing is done in this order: 1. images = images + bgr_add 2. images = images * bgr_mult 3. (optional) Convert BGR channels to RGB 4. Pad or resize the input to the closest larger size multiple of self.output_stride

Parameters:

imagestorch.Tensor

A tensor with at least 3 dimensions in this order: […, 3, H, W].

bgr_addUnion[float, Tuple[float, float, float], np.ndarray, torch.Tensor], default 0

BGR values to be added to the images. It can be a single value, a triple, or a tensor with a shape compatible with images.

bgr_multUnion[float, Tuple[float, float, float], np.ndarray, torch.Tensor], default 1

BGR values to be multiplied by the images. It can be a single value, a triple, or a tensor with a shape compatible with images.

bgr_to_rgbbool, default False

If True, flip the channels to convert from BGR to RGB.

image_resizerOptional[Union[InputPadder, InputScaler]]

An instance of InputPadder or InputScaler that will be used to resize the images. If not provided, a new one will be created based on the given resize_mode.

resize_modestr, default “pad”

How to resize the input. Accepted values are “pad” and “interpolation”.

target_sizeOptional[Tuple[int, int]], default None

If given, the images will be resized to this size, instead of calculating a multiple of self.output_stride.

pad_modestr, default “replicate”

Used if resize_mode == “pad”. How to pad the input. Must be one of the values accepted by the ‘mode’ argument of torch.nn.functional.pad.

pad_valuefloat, default 0.0

Used if resize_mode == “pad” and pad_mode == “constant”. The value to fill in the padded area.

pad_two_sidebool, default True

Used if resize_mode == “pad”. If True, half of the padding goes to left/top and the rest to right/bottom. Otherwise, all the padding goes to the bottom right.

interpolation_modestr, default “bilinear”

Used if resize_mode == “interpolation”. How to interpolate the input. Must be one of the values accepted by the ‘mode’ argument of torch.nn.functional.interpolate.

interpolation_align_cornersbool, default True

Used if resize_mode == “interpolation”. See ‘align_corners’ in https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html.

Returns:

torch.Tensor

A copy of the input images after applying all of the pre-processing steps.

Union[InputPadder, InputScaler]

An instance of InputPadder or InputScaler that was used to resize the images. Can be used to reverse the resizing operations.

test_dataloader() → List[DataLoader]

Initialize and return the list of test dataloaders.

The list of datasets to load is provided by the –test_dataset argument.

Returns:

Optional[List[DataLoader]]

A list of dataloaders each for one dataset.

Raises:

ValueError

If –test_dataset is not provided.

train_dataloader() → DataLoader

Initialize and return the training dataloader.

self.args.train_dataset will be parsed into the selected datasets and their parameters. parse_dataset_selection() is used to parse, and the parsed outputs are used as follows: (dataset_multiplier, dataset_name, dataset_params…).

A method called _get_[dataset_name]_dataset(dataset_params) will be called to get the instance of each dataset.

Each dataset will be repeated dataset_multiplier times and then mixed together into a single dataloader with samples of all the datasets.

Returns:

DataLoader

A single dataloader with all the selected training datasets.

See also

parse_dataset_selection

This function does the parsing of self.args.train_dataset. You can read its documentation to learn how to write valid dataset strings.

val_dataloader

Similar to this method, but val_dataloader returns one dataloader for each dataset.

Notes

This function could be easily modified to return one dataloader for each dataset, which may be preferable in some cases. However, this would require that the loading function have some logic to mix all the dataloaders internally. In order to keep the whole BaseModel more generic, we use just a single dataloader randomly mixed to train.

training_step(batch: Dict[str, Any], batch_idx: int) → Dict[str, Any]

Perform one step of the training.

This function is called internally by Pytorch-Lightning during training.

Parameters:

batchDict[str, Any]

One batch of data, that is going to be given as input to the network.

batch_idxint

The index of the current batch.

Returns:

Dict[str, Any]

A dict with two keys:

• ‘loss’: torch.Tensor, containing the loss value. Required by Pytorch-Lightning for the optimization step.

• ‘dataset_name’: str, a string representing the name of the dataset from where this batch came from. Used only for logging purposes.

val_dataloader() → List[DataLoader] | None

Initialize and return the list of validation dataloaders.

self.args.val_dataset will be parsed into the selected datasets and their parameters. parse_dataset_selection() is used to parse, and the parsed outputs are used as follows: (dataset_multiplier, dataset_name, dataset_params…).

A method called _get_[dataset_name]_dataset(dataset_params) will be called to get the instance of each dataset.

If there is a dataset_multiplier, it will be ignored. Each dataset will be returned in a separated dataloader.

Returns:

Optional[List[DataLoader]]

A list of dataloaders each for one dataset.

See also

parse_dataset_selection

This function does the parsing of self.args.val_dataset. You can read its documentation to learn how to write valid dataset strings.

validation_epoch_end(outputs: List[Dict[str, Tensor]] | List[List[Dict[str, Tensor]]]) → None

Perform operations at the end of one validation epoch.

This function is called internally by Pytorch-Lightning during validation.

Parameters:

outputsUnion[List[Dict[str, Any]], List[List[Dict[str, Any]]]]

A list in which each element is the output of validation_step().

See also

validation_step

validation_step(batch: Dict[str, Any], batch_idx: int, dataloader_idx: int = 0) → Dict[str, Any]

Perform one step of the validation.

This function is called internally by Pytorch-Lightning during validation.

Parameters:

batchDict[str, Any]

One batch of data, that is going to be given as input to the network.

batch_idxint

The index of the current batch.

dataloader_idxint, default 0

When using multiple loaders, indicate from which loader this input is coming from.

Returns:

Dict[str, Any]

A dict with two keys:

• ‘metrics’: dict[str, torch.Tensor], the metrics computed during this validation step.

• ‘dataset_name’: str, a string representing the name of the dataset from where this batch came from. Used only for logging purposes.

See also

ptlflow.utils.flow_metrics.FlowMetrics

class to manage and compute the optical flow metrics.