zae_engine.examples package¶

Submodules¶

zae_engine.examples.classification1D module¶

class zae_engine.examples.classification1D.InferenceDataset(x)[source]¶: Bases: Dataset

class zae_engine.examples.classification1D.InferenceTrainer(model, device, mode, optimizer, scheduler)[source]¶

Bases: Trainer

test_step(batch: dict)[source]¶

Perform a testing step.

This method must be implemented by subclasses. This unused dummy function exists to provide I/O format information.

The batch is a part of the dataset in the dataloader fetched via the __getitem__ method. The dictionary consists of {‘str’: torch.tensor(value)}, and must include ‘loss’. Note that this function works only in ‘test’ mode, and the backward() method is not necessary.

Parameters:: batch (Union[tuple, dict]) – A batch of data.
Returns:: A dictionary containing the results of the testing step, including the loss.
Return type:: Dict[str, torch.Tensor]

Examples

>>> x, y, fn = batch['x'], batch['y'], batch['fn']  # or x, y, fn = batch
>>> outputs = self.model(x)
>>> loss = criteria(outputs)
>>> return {"loss": loss, "mean_output": torch.mean(outputs)}

train_step(batch: tuple | dict) → Dict[str, Tensor][source]¶

Perform a training step.

This method must be implemented by subclasses. This unused dummy function exists to provide I/O format information.

The batch is a part of the dataset in the dataloader fetched via the __getitem__ method. The dictionary consists of {‘str’: torch.tensor(value)}, and must include ‘loss’. Note that this function works only in ‘train’ mode, hence the backward() method in nn.Module is necessary.

Parameters:: batch (Union[tuple, dict]) – A batch of data.
Returns:: A dictionary containing the results of the training step, including the loss.
Return type:: Dict[str, torch.Tensor]

Examples

>>> x, y, fn = batch['x'], batch['y'], batch['fn']  # or x, y, fn = batch
>>> outputs = self.model(x)
>>> loss = criteria(outputs)
>>> return {"loss": loss, "mean_output": torch.mean(outputs)}

zae_engine.examples.classification1D.core(x: ndarray | Tensor) → list[source]¶

Perform inference on a given input array using a predefined CNN model and data pipeline.

This function processes the input data through a series of preprocessing steps, feeds it into a Convolutional Neural Network (CNN) model, and returns the predicted class indices for each input segment.

Parameters:

x (np.ndarray) – Input array to be processed. The array should be less than 3-D (i.e., 1-D or 2-D). For 2-D inputs, the shape should be (num_samples, 2048), where 2048 is the expected segment length.

Returns:

A NumPy array containing the predicted class indices for each input sample.

Return type:

np.ndarray

Raises:

AssertionError – If the input array has 3 or more dimensions.
Exception – If any error occurs during the inference process.

Example

>>> import numpy as np
>>> x = np.zeros(20480).reshape(-1, 2048)  # 10 samples of length 2048 each
>>> predictions = core(x)
>>> print(predictions)
[0 1 0 2 1 0 1 0 1 2]

Notes

This function performs the following steps: 1. Validates the input array dimensions. 2. Sets up the computation device (CPU or GPU). 3. Initializes the data pipeline, including dataset and data loader with preprocessing modules. 4. Sets up the CNN model, optimizer, and learning rate scheduler. 5. Executes the inference process using the model and returns the predicted class indices.

The model used is a CNN-based architecture (CNNBase with BasicBlock), converted from 2D to 1D. The input data is preprocessed using filtering, scaling, and one-hot encoding before being fed into the model. Inference is performed in batches to handle large input arrays efficiently.

zae_engine.examples.ddpm_mnist module¶

class zae_engine.examples.ddpm_mnist.CustomMNISTDataset(root, train=True, transform=None, download=False)[source]¶

Bases: Dataset

Custom Dataset class for MNIST data.

This class wraps the torchvision MNIST dataset and returns a dictionary containing the image.

Parameters:

root (str) – Root directory of dataset where MNIST exists or will be saved.
train (bool, optional) – If True, creates dataset from training set, otherwise from test set.
transform (callable, optional) – A function/transform that takes in an image and returns a transformed version.
download (bool, optional) – If True, downloads the dataset from the internet and puts it in root directory.

class zae_engine.examples.ddpm_mnist.DDPM(block: ~typing.Type[~zae_engine.nn_night.blocks.unet_block.UNetBlock | ~torch.nn.modules.module.Module], ch_in: int, ch_out: int, width: int, layers: ~typing.Sequence[int], groups: int = 1, dilation: int = 1, norm_layer: ~typing.Callable[[...], ~torch.nn.modules.module.Module] = <class 'torch.nn.modules.batchnorm.BatchNorm2d'>, skip_connect: bool = False, timestep_embed_dim: int = 256)[source]¶

Bases: AutoEncoder

Denoising Diffusion Probabilistic Model (DDPM) implemented as an AutoEncoder.

This model integrates timestep embeddings into the bottleneck of the AutoEncoder architecture, allowing the model to condition on the diffusion timestep during training.

Parameters:

block (Type[Union[UNetBlock, nn.Module]]) – The block type to use in the AutoEncoder (e.g., UNetBlock).
ch_in (int) – Number of input channels.
ch_out (int) – Number of output channels.
width (int) – Base width of the network.
layers (Sequence[int]) – Number of layers in each block.
groups (int, optional) – Number of groups for group normalization, by default 1.
dilation (int, optional) – Dilation rate for convolutions, by default 1.
norm_layer (Callable[..., nn.Module], optional) – Normalization layer to use, by default nn.BatchNorm2d.
skip_connect (bool, optional) – Whether to use skip connections, by default False.
timestep_embed_dim (int, optional) – Dimension of the timestep embedding, by default 256.

timestep_embedding¶

Module for generating timestep embeddings.

Type:: TimestepEmbedding

t_embed_proj¶

Linear layer to project timestep embeddings to match the bottleneck dimensions.

Type:: nn.Linear

forward(x: Tensor, t: Tensor) → Tensor[source]¶

Forward pass for DDPM.

Parameters:

x (torch.Tensor) – Noised image (x_t) of shape (batch_size, channels, height, width).
t (torch.Tensor) – Timestep tensor of shape (batch_size,).

Returns:

Reconstructed image tensor of shape (batch_size, channels, height, width).

Return type:

torch.Tensor

class zae_engine.examples.ddpm_mnist.DDPMTrainer(model, device: device | Sequence[device], mode: str, optimizer: Optimizer, scheduler: LRScheduler | T | None, *, log_bar: bool = True, scheduler_step_on_batch: bool = False, gradient_clip: float = 0.0)[source]¶

Bases: Trainer

Trainer class specialized for DDPM training and sampling.

Inherits from the abstract Trainer class and implements the train_step and test_step methods. Additionally, it includes methods for generating and visualizing samples using the trained model.

generate(n_samples: int, channels: int, height: int, width: int, intermediate: int = 0, ddim: bool = False) → Tuple[Tensor, List[Any]][source]¶

Generate new samples using the trained diffusion model.

Parameters:

n_samples (int) – Number of samples to generate.
channels (int) – Number of channels in the generated images.
height (int) – Height of the generated images.
width (int) – Width of the generated images.
intermediate (int, optional) – Number of intermediate samples to save during generation, by default 0.
ddim (bool, optional) – Whether to use DDIM sampling. If True, uses DDIM; otherwise, uses DDPM.

Returns:

Generated samples tensor of shape (n_samples, channels, height, width). List of intermediate samples if specified.

Return type:

Tuple[Tensor, List[Any]]

noise_scheduling(noise_scheduler: NoiseScheduler) → None[source]¶

Update the noise scheduler used by the trainer.

Parameters:: noise_scheduler (NoiseScheduler) – New noise scheduler to be used.

test_step(batch: Dict[str, Tensor]) → Dict[str, Tensor][source]¶

Perform a testing step for DDPM.

Parameters:: batch (Dict[str, torch.Tensor]) – A batch of data containing ‘x_t’, ‘x0’, ‘t’, ‘noise’.
Returns:: A dictionary containing the loss and the model’s output.
Return type:: Dict[str, torch.Tensor]

train_step(batch: Dict[str, Tensor]) → Dict[str, Tensor][source]¶

Perform a training step for DDPM.

Parameters:: batch (Dict[str, torch.Tensor]) – A batch of data containing ‘x_t’, ‘x0’, ‘t’, ‘noise’.
Returns:: A dictionary containing the loss and the model’s output.
Return type:: Dict[str, torch.Tensor]

visualize_samples(final_samples: Tensor, intermediate_images: List[Any] = None, train_losses: List[float] = None, valid_losses: List[float] = None, lr_history: List[float] = None) → None[source]¶

Visualize generated samples and training progress.

Parameters:

final_samples (torch.Tensor) – Final generated samples. Shape: (n_samples, channels, height, width).
intermediate_images (List[Any], optional) – Intermediate images for selected samples. Shape: (num_selected, channels, height, width), by default None.
train_losses (List[float], optional) – Training loss history, by default None.
valid_losses (List[float], optional) – Validation loss history, by default None.
lr_history (List[float], optional) – Learning rate history, by default None.

class zae_engine.examples.ddpm_mnist.ForwardDiffusion(noise_scheduler: NoiseScheduler)[source]¶

Bases: object

Class for performing the forward diffusion process by adding noise to the input data.

This class adds noise to the input data at a randomly sampled timestep, producing the noised data x_t, along with the original data x0, the timestep t, and the noise added.

Parameters:: noise_scheduler (NoiseScheduler) – Instance of NoiseScheduler managing the noise levels.

class zae_engine.examples.ddpm_mnist.NoiseScheduler(timesteps=1000, schedule='linear', beta_start=0.0001, beta_end=0.02)[source]¶

Bases: object

Scheduler for managing the noise levels in DDPM.

This class defines the noise schedule used in the diffusion process, supporting both linear and cosine schedules.

Parameters:

timesteps (int, optional) – Total number of diffusion steps. Default is 1000.
schedule (str, optional) – Type of noise schedule (‘linear’, ‘cosine’). Default is ‘linear’.
beta_start (float, optional) – Starting value of beta. Default is 1e-4.
beta_end (float, optional) – Ending value of beta. Default is 0.02.

beta¶

Noise levels for each timestep.

Type:: torch.Tensor

alpha¶

1 - beta for each timestep.

Type:: torch.Tensor

alpha_bar¶

Cumulative product of alpha up to each timestep.

Type:: torch.Tensor

sqrt_alpha_bar¶

Square root of alpha_bar.

Type:: torch.Tensor

sqrt_one_minus_alpha_bar¶

Square root of (1 - alpha_bar).

Type:: torch.Tensor

posterior_variance¶

Variance used in the posterior distribution.

Type:: torch.Tensor

timesteps¶

Total number of diffusion steps.

Type:: int

cosine_beta_schedule(timesteps: int, s: float = 0.008) → Tensor[source]¶

Cosine schedule for beta as proposed in https://arxiv.org/abs/2102.09672.

Parameters:

timesteps (int) – Total number of diffusion steps.
s (float, optional) – Small offset to prevent beta from being exactly 0. Default is 0.008.

Returns:

Beta schedule.

Return type:

torch.Tensor

get_sigma(t: Tensor, ddim: bool = False) → Tensor[source]¶

Get the sigma value for a given timestep.

Parameters:

t (torch.Tensor) – Timestep tensor.
ddim (bool, optional) – Whether to use DDIM sampling. If True, uses sigma=0 for deterministic sampling.

Returns:

Sigma value for the given timestep.

Return type:

torch.Tensor

linear_beta_schedule(timesteps: int, beta_start: float, beta_end: float) → Tensor[source]¶

Linear schedule for beta.

Parameters:

timesteps (int) – Total number of diffusion steps.
beta_start (float) – Starting value of beta.
beta_end (float) – Ending value of beta.

Returns:

Beta schedule.

Return type:

torch.Tensor

class zae_engine.examples.ddpm_mnist.TimestepEmbedding(embed_dim)[source]¶

Bases: Module

Sinusoidal embedding for timesteps.

This module generates sinusoidal embeddings for each timestep, similar to positional encodings used in Transformers.

Parameters:: embed_dim (int) – Dimension of the embedding vector.

forward(t: Tensor) → Tensor[source]¶

Forward pass for timestep embedding.

Parameters:: t (Tensor) – Timestep tensor of shape (batch_size,) or higher.
Returns:: Embedded timestep tensor of shape (batch_size, embed_dim).
Return type:: Tensor

zae_engine.examples.segmentation1D module¶

class zae_engine.examples.segmentation1D.InferenceDataset(x)[source]¶: Bases: Dataset

class zae_engine.examples.segmentation1D.InferenceTrainer(model, device, mode: str, optimizer: Optimizer = None, scheduler=None)[source]¶

Bases: Trainer

test_step(batch: dict)[source]¶

Perform a testing step.

This method must be implemented by subclasses. This unused dummy function exists to provide I/O format information.

The batch is a part of the dataset in the dataloader fetched via the __getitem__ method. The dictionary consists of {‘str’: torch.tensor(value)}, and must include ‘loss’. Note that this function works only in ‘test’ mode, and the backward() method is not necessary.

Parameters:: batch (Union[tuple, dict]) – A batch of data.
Returns:: A dictionary containing the results of the testing step, including the loss.
Return type:: Dict[str, torch.Tensor]

Examples

>>> x, y, fn = batch['x'], batch['y'], batch['fn']  # or x, y, fn = batch
>>> outputs = self.model(x)
>>> loss = criteria(outputs)
>>> return {"loss": loss, "mean_output": torch.mean(outputs)}

train_step(batch: dict)[source]¶

Perform a training step.

This method must be implemented by subclasses. This unused dummy function exists to provide I/O format information.

The batch is a part of the dataset in the dataloader fetched via the __getitem__ method. The dictionary consists of {‘str’: torch.tensor(value)}, and must include ‘loss’. Note that this function works only in ‘train’ mode, hence the backward() method in nn.Module is necessary.

Parameters:: batch (Union[tuple, dict]) – A batch of data.
Returns:: A dictionary containing the results of the training step, including the loss.
Return type:: Dict[str, torch.Tensor]

Examples

>>> x, y, fn = batch['x'], batch['y'], batch['fn']  # or x, y, fn = batch
>>> outputs = self.model(x)
>>> loss = criteria(outputs)
>>> return {"loss": loss, "mean_output": torch.mean(outputs)}

zae_engine.examples.segmentation1D.core(x: ndarray | Tensor) → list[source]¶

Core function to perform inference on a given 1-D input array using a predefined model and data pipeline.

Parameters:: x (Union[np.ndarray, torch.Tensor]) – The input 1-D array to be processed.
Returns:: The result of the inference process as a numpy array with the predicted class indices.
Return type:: np.ndarray
Raises:: AssertionError – If the input array is not 1-D.

Example

>>> x = np.zeros(2048000)
>>> result = core(x)
>>> print(result)

Notes

This function performs the following steps: 1. Checks if the input array is 1-D. 2. Sets up the device for computation (CPU or GPU). 3. Initializes the data pipeline, including dataset and data loader with necessary preprocessing steps. 4. Sets up the model, optimizer, and learning rate scheduler. 5. Runs the inference process using the model and returns the predicted class indices.

The model used is an autoencoder with a UNetBlock structure, converted from 2D to 1D. The data is split into smaller segments using the Spliter class before being fed into the model. The inference is performed in batches to handle large input arrays efficiently.

zae_engine.examples package¶

Submodules¶

zae_engine.examples.classification1D module¶

zae_engine.examples.ddpm_mnist module¶

zae_engine.examples.segmentation1D module¶

Module contents¶