unet

U-Net

Overview

U-Net is a type of convolutional neural network (CNN) that is commonly used for segmentation tasks. U-Net is a fully convolutional network that consists of a series of convolutional layers and pooling layers. The pooling layers are used to downsample the input while the convolutional layers are used to upsample the input. The skip connections between the pooling layers and convolutional layers allow U-Net to preserve spatial/temporal information while also allowing for faster training and inference times.

For more info, refer to the original paper U-Net: Convolutional Networks for Biomedical Image Segmentation.

Parameters are in helia_edge.models.unet_params.

Functions

Name	Description
build	U-Net model from UNetParams
unet_layer	Generate functional U-Net model

Additions

The U-Net architecture has been modified to allow the following:

• Enable 1D and 2D variants.
• Convolutional pairs can factorized into depthwise separable convolutions.
• Specifiy the number of convolutional layers per block both downstream and upstream.
• Normalization can be set between batch normalization and layer normalization.
• ReLU is replaced with the approximated ReLU6.

Usage

from helia_edge.models import ModelSpec, UNetBlockParams, UNetParams, build

block = dict(depth=2, ddepth=1, kernel=(1, 5), pool=(1, 3), strides=(1, 2), skip=True, seperable=True)
params = UNetParams(
    blocks=[UNetBlockParams(filters=f, **block) for f in (12, 24, 32, 48)],
    output_kernel_size=(1, 5),
    include_top=True,
    use_logits=True,
    num_classes=5,
)
model = build(ModelSpec(params=params, input_shape=(1, 800, 1)))

Machine-readable model

• unet_layerfunctionCreate UNet TF functional model
• buildfunctionBuild a UNet model.

function

unet_layer

Python

Create UNet TF functional model

helia_edge/models/unet.py:51

unet_layer(x: keras.KerasTensor, params: UNetParams) -> keras.KerasTensor

Create UNet TF functional model

Parameters

Parameters of unet_layer

Name	Type	Default	Description
x	keras.KerasTensor	Required	Input tensor
params	UNetParams	Required	Model parameters.

Returns

Returns of unet_layer

Type	Description
keras.KerasTensor	keras.KerasTensor: Output tensor

function

build

Python

Build a UNet model.

helia_edge/models/unet.py:264

build(
    params: UNetParams,
    input_shape: tuple[int | None, ...],
    *,
    batch_size: int | None = None,
    name: str | None = None,
) -> keras.Model

Build a UNet model.

Parameters

Parameters of build

Name	Type	Default	Description
params	UNetParams	Required	Model parameters.
input_shape	tuple[int | None, ...]	Required	Input shape without the batch axis; None for a variable axis.
batch_size	int | None	None	Static batch size; None for a dynamic batch.
name	str | None	None	Model name; the family when None.

Returns

Returns of build

Type	Description
keras.Model	keras.Model: The model, named ``unet`` unless ``name`` is given.