hyrax.models

hyrax.models#

Submodules#

Classes#

`HSCAutoencoder`	This autoencoder is designed to work with datasets that are prepared with Hyrax's HSC Data Set class.
`HSCDCAE`	This autoencoder is designed to work with datasets that are prepared with Hyrax's HSC Data Set class.
`ImageDCAE`	This is an autoencoder with skipconnections that should work with
`HyraxAutoencoder`	This autoencoder is designed to work with a wide range of image datasets to allow testing.
`HyraxAutoencoderV2`	This is tweaked version of HyraxAutoencoder and is designed to work with a wide range of imaging datasets.
`HyraxCNN`	This CNN is designed to work with datasets that are prepared with Hyrax's HSC Data Set class.
`HyraxLoopback`	Simple model for testing which returns its own input
`SimCLR`	SimCLR model. Implementation based on Chen, 2020

Functions#

hyrax_model(cls)

Decorator to register a model with the model registry, and to add common interface functions

Package Contents#

class HSCAutoencoder(config, data_sample=None)[source]#

Bases: torch.nn.Module

This autoencoder is designed to work with datasets that are prepared with Hyrax’s HSC Data Set class.

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

encoder#

decoder#

config#

forward(x)[source]#

train_step(batch)[source]#

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 two values the loss function
Returns:: Current loss value – Dictionary containing the loss value for the current batch.
Return type:: dict

class HSCDCAE(config, data_sample=None)[source]#

Bases: torch.nn.Module

This autoencoder is designed to work with datasets that are prepared with Hyrax’s HSC Data Set class.

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

encoder1#

encoder2#

encoder3#

encoder4#

pool#

decoder4#

decoder3#

decoder2#

decoder1#

activation#

config#

forward(x)[source]#

train_step(batch)[source]#

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 two values the loss function
Returns:: Current loss value – Dictionary containing the loss value for the current batch.
Return type:: dict

class ImageDCAE(config, data_sample=None)[source]#

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()[source]#: Calculate the output size after all convolutional layers for the linear bottleneck.

encode(x)[source]#: Encode input to latent space with skip connections.

decode(latent, skip_connections, encoded_shape)[source]#: Decode from latent space to image with skip connections.

forward(x)[source]#: Forward pass - returns latent representation for anomaly detection.

reconstruct(x)[source]#: Full reconstruction for evaluation and anomaly detection.

train_step(batch)[source]#

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

static to_tensor(data_dict)[source]#: Convert structured data to tensor format.

class HyraxAutoencoder(config, data_sample=None)[source]#

Bases: torch.nn.Module

This autoencoder is designed to work with a wide range of image datasets to allow testing.

This example model is taken from this autoenocoder tutorial

The train function has been converted into train_step for use with pytorch-ignite.

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[source]#

conv2d_output_size(input_size, kernel_size, padding=0, stride=1, dilation=1) → int[source]#

_init_encoder()[source]#

_eval_encoder(x)[source]#

_init_decoder()[source]#

_eval_decoder(x)[source]#

forward(batch)[source]#

train_step(batch)[source]#

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[source]#

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

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

_optimizer()[source]#

class HyraxAutoencoderV2(config, data_sample=None)[source]#

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[source]#

conv2d_output_size(input_size, kernel_size, padding=0, stride=1, dilation=1) → int[source]#

_init_encoder()[source]#

_eval_encoder(x)[source]#

_init_decoder()[source]#

_eval_decoder(x)[source]#

forward(batch)[source]#

train_step(batch)[source]#

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][source]#

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

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

class HyraxCNN(config, data_sample=None)[source]#

Bases: torch.nn.Module

This CNN is designed to work with datasets that are prepared with Hyrax’s HSC Data Set class.

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

config#

conv1#

pool#

conv2#

fc1#

fc2#

fc3#

conv2d_output_size(input_size, kernel_size, padding=0, stride=1, dilation=1) → int[source]#

pool2d_output_size(input_size, kernel_size, stride, padding=0, dilation=1) → int[source]#

forward(x)[source]#

train_step(batch)[source]#

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[source]#: Does NOT convert to PyTorch Tensors. This works exclusively with numpy data types and returns a tuple of numpy data types.

class HyraxLoopback(config, data_sample=None)[source]#

Bases: torch.nn.Module

Simple model for testing which returns its own input

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

unused_module#

config#

load#

forward(x)[source]#: We simply return our input

train_step(batch)[source]#: Training is a noop

hyrax_model(cls)[source]#

Decorator to register a model with the model registry, and to add common interface functions

Returns:: The class with additional interface functions.
Return type:: type

class SimCLR(config, shape)[source]#

Bases: torch.nn.Module

SimCLR model. Implementation based on Chen, 2020

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

config#

shape#

backbone#

projection_head#

criterion#

forward(x)[source]#

train_step(x)[source]#