qbm4eo.lbae module

This file is a part of qbm4eo.rst project.

It has been modified as a part of the EuroHPC PL project funded at the Smart Growth Operational Programme 2014-2020, Measure 4.2 under the grant agreement no. POIR.04.02.00-00-D014/20-00.

class qbm4eo.lbae.LBAE(input_size: Tuple[int, ...], out_channels: int, latent_space_size: int, num_layers: int, *args: Dict[str, Any], quantize: bool, **kwargs: Dict[str, Any])[source]

Bases: LightningModule

A class implementing the Latent Bernoulli Autoencoder (LBAE) model.

A default constructor for the LBAE class.

Parameters:

input_size – The size of the input image.
out_channels – The number of channels in the output image.
latent_space_size – The size of the latent space.
num_layers – The number of layers in the encoder and decoder.
quantize – The epochs during which the output of the encoder should be quantized.
args – Additional arguments.
kwargs – Additional keyword arguments.

A function for configuring the optimizers for the LBAE model.

Returns:: A dictionary containing the configured optimizers.

forward(x: Tensor) → Tensor[source]

A forward pass through the LBAE model.

Parameters:: x – The input image.
Returns:: The reconstructed image.

predict_step(batch: Any, batch_idx: int, dataloader_idx: int = Ellipsis) → Any[source]

A function for predicting the output of the model.

Parameters:

batch – A batch of images.
batch_idx – The index of the batch.
dataloader_idx – The index of the dataloader.

Returns:

The output of the model, the input images, and the labels.

test_step(batch: Any, batch_idx: int) → Any[source]

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

Parameters:

batch – The output of your data iterable, normally a DataLoader.
batch_idx – The index of this batch.
dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

Tensor - The loss tensor
dict - A dictionary. Can include any keys, but must include the key 'loss'.
None - Skip to the next batch.

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...


# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

training_step(batch: Tuple[Tensor, Tensor], batch_idx: int) → Tensor[source]

A training step for the LBAE model.

Parameters:

batch – A batch of images.
batch_idx – The index of the batch.

Returns:

The loss value.

qbm4eo.lbae.loss(xr: Tensor, x: Tensor) → Tensor[source]

Loss function for LBAE. Uses MSE.

Parameters:

xr – Reconstructed image.
x – Original image.

Returns:

MSE loss value.