hyrax.models.hyrax_autoencoderv2

Attributes

logger

Classes

ArcsinhActivation	Helper module for HyraxAutoencoderV2 to use the arcsinh function
HyraxAutoencoderV2	This is tweaked version of HyraxAutoencoder and is designed to work with a wide range of imaging datasets.

Module Contents

logger

class ArcsinhActivation(*args, **kwargs)

Bases: torch.nn.Module

Helper module for HyraxAutoencoderV2 to use the arcsinh function

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

forward(x)

class HyraxAutoencoderV2(config, data_sample=None)

Bases: torch.nn.Module

This is tweaked version of HyraxAutoencoder and is designed to work with a wide range of imaging datasets.

V2 improvements: - Configurable final layer activation - Uses criterion and optimizer from config variables

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

config

c_hid

latent_dim

conv_end_w

conv_end_h

conv2d_multi_layer(input_size, num_applications, **kwargs) → int

conv2d_output_size(input_size, kernel_size, padding=0, stride=1, dilation=1) → int

_init_encoder()

_eval_encoder(x)

_init_decoder()

_eval_decoder(x)

forward(batch)

train_step(batch)

This function contains the logic for a single training step. i.e. the contents of the inner loop of a ML training process.

Parameters:

batch (tuple) – A tuple containing the inputs and labels for the current batch.

Returns:

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

Return type:

dict

static to_tensor(data_dict) → tuple[torch.Tensor]

This function converts structured data to the input tensor we need to run

Parameters:

data_dict (dict) – The dictionary returned from our data source