zae_engine.examples package

Submodules

zae_engine.examples.classification1D module

class zae_engine.examples.classification1D.InferenceDataset(x)

기반 클래스: Dataset

class zae_engine.examples.classification1D.InferenceTrainer(model, device, mode, optimizer, scheduler)

기반 클래스: Trainer

test_step(batch: dict)

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.

매개변수:

batch (Union[tuple, dict]) – A batch of data.

반환:

A dictionary containing the results of the testing step, including the loss.

반환 형식:

Dict[str, torch.Tensor]

예제

>>> 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]

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.

매개변수:

batch (Union[tuple, dict]) – A batch of data.

반환:

A dictionary containing the results of the training step, including the loss.

반환 형식:

Dict[str, torch.Tensor]

예제

>>> 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

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.

매개변수:

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.

반환:

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

반환 형식:

np.ndarray

예외 발생:

• AssertionError – If the input array has 3 or more dimensions.

• Exception – If any error occurs during the inference process.

예제

>>> 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]

참고

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)

기반 클래스: Dataset

Custom Dataset class for MNIST data.

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

매개변수:

• 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)

기반 클래스: 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.

매개변수:

• 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

Forward pass for DDPM.

매개변수:

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

• t (torch.Tensor) – Timestep tensor of shape (batch_size,).

반환:

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

반환 형식:

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)

기반 클래스: 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]]

Generate new samples using the trained diffusion model.

매개변수:

• 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.

반환:

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

반환 형식:

Tuple[Tensor, List[Any]]

noise_scheduling(noise_scheduler: NoiseScheduler) → None

Update the noise scheduler used by the trainer.

매개변수:

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

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

Perform a testing step for DDPM.

매개변수:

batch (Dict[str, torch.Tensor]) – A batch of data containing ‘x_t’, ‘x0’, ‘t’, ‘noise’.

반환:

A dictionary containing the loss and the model’s output.

반환 형식:

Dict[str, torch.Tensor]

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

Perform a training step for DDPM.

매개변수:

batch (Dict[str, torch.Tensor]) – A batch of data containing ‘x_t’, ‘x0’, ‘t’, ‘noise’.

반환:

A dictionary containing the loss and the model’s output.

반환 형식:

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

Visualize generated samples and training progress.

매개변수:

• 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)

기반 클래스: 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.

매개변수:

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)

기반 클래스: 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.

매개변수:

• 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

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

매개변수:

• timesteps (int) – Total number of diffusion steps.

• s (float, optional) – Small offset to prevent beta from being exactly 0. Default is 0.008.

반환:

Beta schedule.

반환 형식:

torch.Tensor

get_sigma(t: Tensor, ddim: bool = False) → Tensor

Get the sigma value for a given timestep.

매개변수:

• t (torch.Tensor) – Timestep tensor.

• ddim (bool, optional) – Whether to use DDIM sampling. If True, uses sigma=0 for deterministic sampling.

반환:

Sigma value for the given timestep.

반환 형식:

torch.Tensor

linear_beta_schedule(timesteps: int, beta_start: float, beta_end: float) → Tensor

Linear schedule for beta.

매개변수:

• timesteps (int) – Total number of diffusion steps.

• beta_start (float) – Starting value of beta.

• beta_end (float) – Ending value of beta.

반환:

Beta schedule.

반환 형식:

torch.Tensor

class zae_engine.examples.ddpm_mnist.TimestepEmbedding(embed_dim)

기반 클래스: Module

Sinusoidal embedding for timesteps.

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

매개변수:

embed_dim (int) – Dimension of the embedding vector.

forward(t: Tensor) → Tensor

Forward pass for timestep embedding.

매개변수:

t (Tensor) – Timestep tensor of shape (batch_size,) or higher.

반환:

Embedded timestep tensor of shape (batch_size, embed_dim).

반환 형식:

Tensor

zae_engine.examples.segmentation1D module

class zae_engine.examples.segmentation1D.InferenceDataset(x)

기반 클래스: Dataset

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

기반 클래스: Trainer

test_step(batch: dict)

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.

매개변수:

batch (Union[tuple, dict]) – A batch of data.

반환:

A dictionary containing the results of the testing step, including the loss.

반환 형식:

Dict[str, torch.Tensor]

예제

>>> 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)

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.

매개변수:

batch (Union[tuple, dict]) – A batch of data.

반환:

A dictionary containing the results of the training step, including the loss.

반환 형식:

Dict[str, torch.Tensor]

예제

>>> 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

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

매개변수:

x (Union[np.ndarray, torch.Tensor]) – The input 1-D array to be processed.

반환:

The result of the inference process as a numpy array with the predicted class indices.

반환 형식:

np.ndarray

예외 발생:

AssertionError – If the input array is not 1-D.

예제

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

참고

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.

Module contents