neuralop.training.patching.MultigridPatching2D

class neuralop.training.patching.MultigridPatching2D(model: Module, levels: int = 0, padding_fraction: float = 0, use_distributed: bool = False, stitching: bool = True)[source]

MultigridPatching2D wraps a model in multi-grid domain decomposition and patching.

Parameters:

modelnn.Module

model to wrap

levelsint, optional

number of levels of patching to use, by default 0

padding_fractionfloat, optional

fraction of input size to add as padding around patches, by default 0

use_distributedbool, optional

whether patching is performed in distributed mode, by default False

stitchingbool, optional

whether patches are to be stitched back together in spatial dimensions during training, by default True. Stitching is always performed during evaluation.

Methods

patch(x, y)	use multi-grid domain decomposition to split x and y into patches.
unpatch(x, y[, evaluation])	unpatch tensors created by self.patch.

patch(x, y)[source]

use multi-grid domain decomposition to split x and y into patches. If in a distributed scheme, scatters patches across processes.

Parameters:

xtorch.tensor

model input function

y_type_

model output function

unpatch(x, y, evaluation=False)[source]

unpatch tensors created by self.patch. Stitch patches together if in evaluation mode, or if stitching is applied.

Parameters:

xtorch.tensor

tensor with patching structure created patching input x. May be either inputs x or raw model outputs (same shape/patching structure) Shape (b * n^2, c, h / n + 2 * pad_h, w / n + 2 * pad_w)

ytorch.tensor

tensor of patched ground-truth y. Shape (b * n^2, c, h / n, w / n) or (b, c, h, w) when not stitched

evaluationbool, optional

whether in evaluation mode, by default False. If True, x and y are both evaluated after stitching, regardless of other settings.