flambe.vision.classification.model

Module Contents

class flambe.vision.classification.model.ImageClassifier(encoder: Module, output_layer: Module)[source]

Bases: flambe.nn.Module

Implements a simple image classifier.

This classifier consists of an encocder module, followed by a fully connected output layer that outputs a probability distribution.

encoder

The encoder layer

Type:Moodule
output_layer

The output layer, yields a probability distribution over targets

Type:Module
forward(self, data: Tensor, target: Optional[Tensor] = None)[source]

Run a forward pass through the network.

Parameters:
  • data (Tensor) – The input data
  • target (Tensor, optional) – The input targets, optional
Returns:

The output predictions, and optionally the targets

Return type:

Union[Tensor, Tuple[Tensor, Tensor]