CNNModel

class CNNModel(backbone, num_classes=1)[source]

Retrieve the model backbone and attach an extra FCN to perform classification.

This class initializes a Convolutional Neural Network (CNN) model with a specified backbone and attaches a fully connected layer for classification tasks.

Parameters:

  • backbone (str) – Name of the CNN model backbone (e.g., “resnet18”, “densenet121”).

  • num_classes (int) – Number of classes output by model. Defaults to 1.

num_classes

Number of classes output by the model.

Type:

int

feat_extract

Backbone CNN model.

Type:

nn.Module

pool

Type of pooling applied after feature extraction.

Type:

nn.Module

classifier

Linear classifier module used to map the features to the output.

Type:

nn.Module

Example

>>> model = CNNModel("resnet18", num_classes=2)
>>> output = model(torch.randn(1, 3, 224, 224))
>>> print(output.shape)

Initialize CNNModel.

Methods

forward

Pass input data through the model.

infer_batch

Run inference on an input batch.

postproc

Define the post-processing of this class of model.

Attributes

training

forward(imgs)[source]

Pass input data through the model.

Parameters:
Returns:

The output logits after passing through the model.

Return type:

torch.Tensor

static infer_batch(model, batch_data, device='cpu')[source]

Run inference on an input batch.

Contains logic for forward operation as well as i/o aggregation.

Parameters:

  • model (nn.Module) – PyTorch defined model.

  • batch_data (torch.Tensor) – A batch of data generated by torch.utils.data.DataLoader.

  • device (str) – Transfers model to the specified device. Default is “cpu”.

Return type:

dict[str, ndarray]

Example

>>> output = _infer_batch(model, batch_data, "cuda")
>>> print(output)
static postproc(image)[source]

Define the post-processing of this class of model.

This simply applies argmax along last axis of the input.

Parameters:

image (np.ndarray) – The input image array.

Returns:

The post-processed image array.

Return type:

np.ndarray