x-unet

Install the x-unet package using pip to access the U-net implementation with efficient attention and nested structures.

$ pip install x-unet

Initialize XUnet for 2D images by specifying the input dimensions and nested depths. The nested_unet_depths parameter follows the U-Net squared architecture.

Key parameters:

• dim: The base dimension.
• channels: Number of input channels.
• dim_mults: A tuple of dimension multipliers for each level.
• nested_unet_depths: Depths for the nested U-Nets.
• consolidate_upsample_fmaps: If True, consolidates outputs from all upsample blocks (as used in U-Net squared).

import torch
from x_unet import XUnet

unet = XUnet(
    dim = 64,
    channels = 3,
    dim_mults = (1, 2, 4, 8),
    nested_unet_depths = (7, 4, 2, 1),
    consolidate_upsample_fmaps = True,
)

img = torch.randn(1, 3, 256, 256)
out = unet(img) # (1, 3, 256, 256)

To process 3D data such as video sequences or medical scans, initialize XUnet with a frame_kernel_size greater than 1. The input tensor shape should be (batch, channels, frames, height, width).

Key parameters for 3D:

• frame_kernel_size: Set to a value > 1 to enable 3D processing.
• weight_standardize: Boolean to enable weight standardization.

import torch
from x_unet import XUnet

unet = XUnet(
    dim = 64,
    frame_kernel_size = 3,                 # set this to greater than 1
    channels = 3,
    dim_mults = (1, 2, 4, 8),
    nested_unet_depths = (5, 4, 2, 1),
    consolidate_upsample_fmaps = True,
    weight_standardize = True
)

video = torch.randn(1, 3, 10, 128, 128)    # (batch, channels, frames, height, width)
out = unet(video) # (1, 3, 10, 128, 128)