learning.Patchwise

learning.Patchwise#

class bartorch.learning.Patchwise#

Bases: Module

Network applied to an image patch by patch on its own device.

The image stays on its device, normally the host, and is cut into non-overlapping patches over its trailing len(patch) axes. Groups of batch patches are copied to device, passed through the network under automatic mixed precision and copied back into place, so a volume or a time series of volumes larger than the device’s memory passes through a network trained on patches. The device holds the network and batch patches, or under overlap two groups and the outputs not yet placed. The result is float32 on the input’s device and differentiable with respect to the input and the weights.

With shift the patch grid is offset by a random amount at every call, drawn from PyTorch’s generator, so seams fall in a different place at each iteration of a reconstruction, and Unrolled recomputes a checkpointed step exactly. The image is zero-padded by the offset in front and to a whole patch behind.

The network is moved to device at every call if it is elsewhere, so a training loop that moves the whole model does not move it; its parameters remain the same objects.

Parameters:

  • net (nn.Module) – Called on (n, channels, [frames,] *patch) as net(x) or net(x, sigma), with the keywords this module is called with.

  • patch (sequence of int) – Patch extent along each trailing axis, two or three of them.

  • device (torch.device or str, default=None) – Where the network runs; the first CUDA device when there is one, the host otherwise.

  • dtype (torch.dtype or None or "auto", default="auto") – Precision of automatic mixed precision; None disables it. "auto" is bfloat16 on a card that supports it, float16 on one that does not (such as a T4), and none on the host.

  • batch (int, default=1) – Patches per call of the network.

  • shift (bool, default=True) – Offset the patch grid at random at every call.

  • overlap (bool, default=True) – Without gradient, with the image on the host and the network on a card, copy one group of patches on a second stream while the network computes on another.

Raises:

ValueError – patch not of two or three positive extents, or batch below one.

Examples

>>> net = learning.Patchwise(learning.UNet(10), patch=(64, 64, 64), batch=4)
>>> prior = priors.ImplicitPrior(learning.ComplexNet(net, channels=1))
forward()#

Apply the network to every patch of x.

Parameters:

  • x (torch.Tensor) – (n, channels, ..., *patched), with the patched axes last.

  • sigma (float or torch.Tensor, default=None) – Passed to the network; one value per item of x is repeated for each of that item’s patches.

  • **keywords – Passed to the network, repeated per patch like sigma.

Returns:

Float32, of x’s shape, on x’s device.

Return type:

torch.Tensor

extra_repr()#

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.