hyrax.models.image_dcae

Classes

ArcsinhActivation	Helper module for ImageDCAE to use the arcsinh function
ImageDCAE	This is an autoencoder with skipconnections that should work with

Module Contents

class ArcsinhActivation(*args, **kwargs)

Bases: torch.nn.Module

Helper module for ImageDCAE to use the arcsinh function

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

class ImageDCAE(config, data_sample=None)

Bases: torch.nn.Module

This is an autoencoder with skipconnections that should work with arbitarily sized images with arbitrary number of channels.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

input_shape

config

latent_dim

base_channel_size

conv_output_size

encoder1

encoder2

encoder3

encoder4

pool

latent_encoder

latent_decoder

decoder4

decoder3

decoder2

decoder1

activation

_calculate_conv_output_size()

Calculate the output size after all convolutional layers for the linear bottleneck.

encode(x)

Encode input to latent space with skip connections.

decode(latent, skip_connections, encoded_shape)

Decode from latent space to image with skip connections.

forward(x)

Forward pass - returns latent representation for anomaly detection.

reconstruct(x)

Full reconstruction for evaluation and anomaly detection.

train_step(batch)

This function contains the logic for a single training step.

Parameters:

batch (tuple) – A tuple containing the two values the loss function

Returns:

Current loss value – Dictionary containing the loss value for the current batch.

Return type:

dict

_optimizer()

Default optimizer configuration.

static to_tensor(data_dict)

Convert structured data to tensor format.