im2sim.layers.ImageConvBlock

Contents

im2sim.layers.ImageConvBlock#

class ImageConvBlock(in_channels, out_channels, rank, cfg)[source]#

Bases: Module

A configurable image convolutional block that consists of a sequence of convolutional layers, normalization layers, dropout layers, and attention layers. The block supports residual connections and allows for flexible configuration of its components.

Parameters:

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • rank (int) – The rank of the convolutional layers (e.g., 2 for 2D convolutions).

  • cfg (ImageConvBlockConfig) – Configuration object that defines the parameters of the block.

Examples

To create an ImageConvBlock with a depth of 3, ReLU activation, and softmax output activation, you can use the following code:

cfg = ImageConvBlockConfig(depth=3, activation="ReLU", out_activation="softmax")
model = ImageConvBlock(
       rank=2,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )

Since the configs are rankless, you could use the same config for a 1D, 2D, or 3D convolutional block by changing the rank parameter when creating the ImageConvBlock instance.

cfg = ImageConvBlockConfig(depth=3, activation="ReLU", out_activation="softmax")
model1D = ImageConvBlock(
       rank=1,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )
model2D = ImageConvBlock(
       rank=2,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )
model3D = ImageConvBlock(
       rank=3,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )

Models can be saved and loaded using the standard PyTorch methods:

torch.save(model.state_dict(), "model.pth")
model.load_state_dict(torch.load("model.pth"))

Configs can also be saved and loaded using the methods provided in the im2sim.configs.ImageConvBlockConfig class:

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

Contents