TENNs modules

Spatiotemporal blocks

class tenns_modules.SpatialBlock(in_channels, out_channels, kernel_size, stride=1, bias=True, depthwise=False)[source]

A spatial (potentially separable) convolution.

BatchNormalization and ReLU activation are included in this block.

Parameters:

  • in_channels (int) – number of channels in the input

  • out_channels (int) – number of channels produced by the block

  • kernel_size (int) – size of the kernel

  • stride (int, optional) – stride of the convolution. Defaults to 1.

  • bias (bool, optional) – if True, adds a learnable bias to the output. Defaults to True.

  • depthwise (bool, optional) – if True, the block will be separable. Defaults to False.

class tenns_modules.TemporalBlock(in_channels, out_channels, kernel_size, bias=True, depthwise=False, use_pleiades=False)[source]

A temporal (potentially separable) convolution.

BatchNormalization and ReLU activation are included in this block.

Parameters:

  • in_channels (int) – number of channels in the input

  • out_channels (int) – number of channels produced by the block

  • kernel_size (int) – size of the kernel

  • bias (bool, optional) – if True, adds a learnable bias to the output. Defaults to True.

  • depthwise (bool, optional) – if True, the block will be separable. Defaults to False.

  • use_pleiades (bool, optional) – if True, the first conv3d is a PleiadesLayer. Defaults to False.

class tenns_modules.SpatioTemporalBlock(in_channels, med_channels, out_channels, t_kernel_size, s_kernel_size, s_stride=1, bias=True, t_depthwise=False, s_depthwise=False, use_pleiades=False)[source]

A combination of temporal and spatial convolutions.

This first applies a temporal convolution (potentially separable) to process the input over the temporal dimension, followed by a spatial convolution (potentially separable) to process the output over the spatial dimension.

Parameters:

  • in_channels (int) – number of channels in the input

  • med_channels (int) – number of channels produced by the TemporalBlock

  • out_channels (int) – number of channels produced by the SpatialBlock

  • t_kernel_size (int) – size of the TemporalBlock kernel

  • s_kernel_size (int) – size of the SpatialBlock kernel

  • s_stride (int, optional) – stride of the SpatialBlock convolution. Defaults to 1.

  • bias (bool, optional) – if True, adds a learnable bias to the output. Defaults to True.

  • t_depthwise (bool, optional) – if True, the TemporalBlock will be separable. Defaults to False.

  • s_depthwise (bool, optional) – if True, the SpatialBlock will be separable. Defaults to False.

  • use_pleiades (bool, optional) – if True, the first conv3d of the TemporalBlock is a PleiadesLayer. Defaults to False.

Export

tenns_modules.export_to_onnx(model, input_shape, out_path='model.onnx')[source]

Exports a PyTorch model to ONNX format, and saves the optimized ONNX model to a specified file.

Parameters:

  • model (torch.nn.Module) – The input PyTorch model to be converted.

  • input_shape (tuple) – The shape of the input tensor for the model.

  • out_path (str, optional) – The output file path for the ONNX model. Defaults to ‘model.onnx’.