ml.models.architectures.bifpn

Defines a BiFPN model architecture.

The BiFPN model takes in a list of feature maps from a backbone and outputs a list of feature maps that have been fused together using a bidirectional feature pyramid network. This is a good general-purpose model for multi-scale per-pixel classification or prediction.

bsz = image.shape[0]
x1, x2, x3, x4 = resnet_backbone(image)
assert x1.shape[1] == 256
assert x2.shape[1] == 512
assert x3.shape[1] == 1024
assert x4.shape[1] == 2048
bifpn = BiFPN([256, 512, 1024, 2048], feature_size=32)
f1, f2, f3, f4 = bifpn([x1, x2, x3, x4])
assert f1.shape == (bsz, 32, *x1.shape[2:])
assert f2.shape == (bsz, 32, *x2.shape[2:])
assert f3.shape == (bsz, 32, *x3.shape[2:])
assert f4.shape == (bsz, 32, *x4.shape[2:])

class ml.models.architectures.bifpn.DepthwiseConvBlock(in_channels: int, out_channels: int, kernel_size: int = 1, stride: int = 1, padding: int = 0, dilation: int = 1)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.architectures.bifpn.ConvBlock(in_channels: int, out_channels: int, kernel_size: int = 1, stride: int = 1, padding: int = 0, dilation: int = 1)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.architectures.bifpn.BiFPNBlock(feature_size: int, num_inputs: int, scale_factor: float)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: list[torch.Tensor]) → list[torch.Tensor]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.architectures.bifpn.BiFPN(input_size: list[int], feature_size: int = 64, num_layers: int = 2, scale_factor: float = 2.0)

Bases: Module

Defines the BiFPN module.

This implementation assumes a constant downsampling factor of 2, and that the inputs are sorted in order from least to most downsampled, meaning that the first input should have the largest height and width.

The inputs to the BiFPN are typically the outputs of some feature extractor backbone like ResNet, and the outputs of the BiFPN are typically used as inputs to some head like a classification head. This can be used for multi-scale per-pixel classification or prediction.

Parameters:

• input_size – The number of channels for each input.

• feature_size – The number of channels for the BiFPN features. All of the BiFPN output features will have this number of channels.

• num_layers – The number of layers to use for each BiFPN block.

• scale_factor – The downsampling factor between input images. This is assumed to be the same for all inputs.

Inputs:

inputs: A list of tensors with shape ``(batch_size, channels, height,

width)``, where the tensors are sorted in order from least to most downsampled.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: list[torch.Tensor]) → list[torch.Tensor]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.