frarch.models.classification.cnn.mnistCNN module#

class frarch.models.classification.cnn.mnistCNN.MNISTCNN(input_channels: int = 1, embedding_size: int = 256)[source]#

Bases: torch.nn.modules.module.Module

Small CNN network for MNIST dataset.

Parameters

  • input_channels (int) – Number of input channels in the input tensor. Defaults to 1.

  • embedding_size (int) – Size of the output embedding for the feature extraction network. Defaults to 256.

forward(x: torch.Tensor) → torch.Tensor[source]#

Define the computation performed at every call.

forward computation for MNISTCNN.

Parameters

x (torch.Tensor) – input to the model.

Returns

output of the model.

Return type

torch.Tensor

training: bool#
class frarch.models.classification.cnn.mnistCNN.MNISTClassifier(embedding_size: int = 256, num_classes: int = 10)[source]#

Bases: torch.nn.modules.module.Module

Classifier network for MNISTCNN.

Parameters

  • embedding_size (int) – embedding size from MNISTCNN network. Defaults to 256.

  • classes (int) – number of output classes for the classifier. Defaults to 10.

forward(x: torch.Tensor) → torch.Tensor[source]#

Define the computation performed at every call.

forward computation for MNISTClassifier.

Parameters

x (torch.Tensor) – input to the model.

Returns

output of the model.

Return type

torch.Tensor

training: bool#