AtomAI Models
Segmentor
- class atomai.models.Segmentor(model='Unet', nb_classes=1, **kwargs)[source]
Bases:
SegTrainer
Model for semantic segmentation-based analysis of images
- Parameters:
model (
Type
[Union
[str
,Module
]]) – Type of model to train: ‘Unet’, ‘ResHedNet’ or ‘dilnet’ (Default: ‘Unet’). See atomai.nets for more details. One can also pass here a custom fully convolutional neural network model.nb_classes (
int
) – Number of classes in classification scheme (last NN layer)**batch_norm (bool) – Apply batch normalization after each convolutional layer (Default: True)
**dropout (bool) – Apply dropouts to the three inner blocks in the middle of a network (Default: False)
**upsampling_mode (str) – “bilinear” or “nearest” upsampling method (Default: “bilinear”)
**nb_filters (int) – Number of convolutional filters (aka “kernels”) in the first convolutional block (this number doubles in the consequtive block(s), see definition of Unet and dilnet models for details)
**with_dilation (bool) – Use dilated convolutions in the bottleneck of Unet (Default: False)
**layers (list) – List with a number of layers in each block. For UNet the first 4 elements in the list are used to determine the number of layers in each block of the encoder (including bottleneck layer), and the number of layers in the decoder is chosen accordingly (to maintain symmetry between encoder and decoder)
Example:
>>> # Initialize model >>> model = aoi.models.Segmentor(nb_classes=3) >>> # Train >>> model.fit(images, masks, images_test, masks_test, >>> training_cycles=300, compute_accuracy=True, swa=True) >>> # Predict with trained model >>> nn_output, coordinates = model.predict(expdata)
- fit(X_train, y_train, X_test=None, y_test=None, loss='ce', optimizer=None, training_cycles=1000, batch_size=32, compute_accuracy=False, full_epoch=False, swa=False, perturb_weights=False, **kwargs)[source]
Compiles a trainer and performs model training
- Parameters:
X_train (
Union
[ndarray
,Tensor
]) – 4D numpy array or pytorch tensor of training images (n_samples, 1, height, width). One can also pass a regular 3D image stack without a channel dimension of 1 which will be added automaticallyy_train (
Union
[ndarray
,Tensor
]) – 4D (binary) / 3D (multiclass) numpy array or pytorch tensor of training masks (aka ground truth) stacked along the first dimension. The reason why in the multiclass case the X_train is 4-dimensional and the y_train is 3-dimensional is because of how the cross-entropy loss is calculated in PyTorch (see https://pytorch.org/docs/stable/nn.html#nllloss).X_test (
Union
[ndarray
,Tensor
,None
]) – 4D numpy array or pytorch tensor of test images (stacked along the first dimension)y_test (
Union
[ndarray
,Tensor
,None
]) – 4D (binary) / 3D (multiclass) numpy array or pytorch tensor of training masks (aka ground truth) stacked along the first dimension.loss (
str
) – loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss)optimizer (
Optional
[Type
[Optimizer
]]) – weights optimizer (defaults to Adam optimizer with lr=1e-3)training_cycles (
int
) – Number of training ‘epochs’. If full_epoch argument is set to False, 1 epoch == 1 mini-batch. Otherwise, each cycle corresponds to all mini-batches of data passing through a NN.batch_size (
int
) – Size of training and test mini-batchescompute_accuracy (
bool
) – Computes accuracy function at each training cyclefull_epoch (
bool
) – If True, passes all mini-batches of training/test data at each training cycle and computes the average loss. If False, passes a single (randomly chosen) mini-batch at each cycle.swa (
bool
) – Saves the recent stochastic weights and averages them at the end of trainingperturb_weights (
bool
) – Time-dependent weight perturbation, \(w\leftarrow w + a / (1 + e)^\gamma\), where parameters a and gamma can be passed as a dictionary, together with parameter e_p determining every n-th epoch at which a perturbation is applied**lr_scheduler (list of floats) – List with learning rates for each training iteration/epoch. If the length of list is smaller than the number of training iterations, the last values in the list is used for the remaining iterations.
**print_loss (int) – Prints loss every n-th epoch
**accuracy_metrics (str) – Accuracy metrics (used only for printing training statistics)
**filename (str) – Filename for model weights (appended with “_test_weights_best.pt” and “_weights_final.pt”)
**plot_training_history (bool) – Plots training and test curves vs. training cycles at the end of training
**kwargs – One can also pass kwargs for utils.datatransform class to perform the augmentation “on-the-fly” (e.g. rotation=True, gauss_nois=[20, 60], etc.)
- predict(imgdata, refine=False, logits=True, resize=None, compute_coords=True, **kwargs)[source]
Apply (trained) model to new data
- Parameters:
image_data – 3D image stack or a single 2D image (all greyscale)
refine (
bool
) – Atomic positions refinement with 2d Gaussian peak fitting (may take some time)logits (
bool
) – Indicates whether the features are passed through a softmax/sigmoid layer at the end of a neural network (logits=True for AtomAI models)resize (
Optional
[Tuple
[int
,int
]]) – Resizes input data to (new_height, new_width) before passing to a neural networkcompute_coords (bool) – Computes centers of the mass of individual blobs in the segmented images (Default: True)
**thresh (float) – Value between 0 and 1 for thresholding the NN output (Default: 0.5)
**d (int) – half-side of a square around each atomic position used for refinement with 2d Gaussian peak fitting. Defaults to 1/4 of average nearest neighbor atomic distance
**num_batches (int) – number of batches (Default: 10)
**norm (bool) – Normalize data to (0, 1) during pre-processing
**verbose (bool) – verbosity
- Return type:
Tuple
[ndarray
,Dict
[int
,ndarray
]]- Returns:
Semantically segmented image and coordinates of (atomic) objects
ImSpec
- class atomai.models.ImSpec(in_dim, out_dim, latent_dim=2, **kwargs)[source]
Bases:
ImSpecTrainer
Model for predicting spectra from images and vice versa
- Parameters:
in_dim (
Tuple
[int
]) – Input data dimensions. (height, width) for images or (length,) for spectraout_dim (
Tuple
[int
]) – Output dimensions. (length,) for spectra or (height, width) for imageslatent_dim (
int
) – Dimensionality of the latent space (number of neurons in a fully connected “bottleneck” layer)**seed (int) – Seed used when initializng model weights (Default: 1)
**nblayers_encoder (int) – Number of convolutional layers in the encoder
**nblayers_decoder (int) – Number of convolutional layers in the decoder
**nbfilters_encoder (int) – number of convolutional filters in each layer of the encoder
**nbfilters_decoder (int) – Number of convolutional filters in each layer of the decoder
**batch_norm (bool) – Apply batch normalization after each convolutional layer (Default: True)
**encoder_downsampling (int) – Downsamples input data by this factor before passing to convolutional layers (Default: no downsampling)
**decoder_upsampling (bool) – Performs upsampling+convolution operation twice on the reshaped latent vector (starting from image/spectra dims 4x smaller than the target dims) before passing to the decoder
Example:
>>> in_dim = (16, 16) # Input dimensions (images) >>> out_dim = (64,) # Output dimensions (spectra) >>> # Initialize and train model >>> model = aoi.models.ImSpec(in_dim, out_dim, latent_dim=10) >>> model.fit(imgs_train, spectra_train, imgs_test, spectra_test, >>> full_epoch=True, training_cycles=120, swa=True) >>> # Make a prediction with the trained model >>> prediction = model.predict(imgs_test, norm=False)
- fit(X_train, y_train, X_test=None, y_test=None, loss='mse', optimizer=None, training_cycles=1000, batch_size=64, compute_accuracy=False, full_epoch=False, swa=False, perturb_weights=False, **kwargs)[source]
Compiles a trainer and performs model training
- Parameters:
X_train (
Union
[ndarray
,Tensor
]) – 4D numpy array or torch tensor with image data (n_samples x 1 x height x width) or 3D array/tensor with spectral data (n_samples x 1 x signal_length). It is also possible to pass 3D and 2D arrays by ignoring the channel dim of 1, which will be added automatically. The X_train is typically referred to as ‘features’y_train (
Union
[ndarray
,Tensor
]) – 3D numpy array or torch tensor with spectral data (n_samples x 1 x signal_length) or 4D array/tensor with image data (n_samples x 1 x height x width). It is also possible to pass 2D and 3D arrays by ignoring the channel dim of 1, which will be added automatically. Note that if your X_train data are images, then your y_train must be spectra and vice versa. The y_train is typicaly referred to as “targets”X_test (
Union
[ndarray
,Tensor
,None
]) – Test data (features) of the same dimesnionality (except for the number of samples) as X_trainy_test (
Union
[ndarray
,Tensor
,None
]) – Test data (targets) of the same dimesnionality (except for the number of samples) as y_trainloss (
str
) – Loss function (currently works only with ‘mse’)optimizer (
Optional
[Type
[Optimizer
]]) – weights optimizer (defaults to Adam optimizer with lr=1e-3)training_cycles (
int
) – Number of training ‘epochs’. If full_epoch argument is set to False, 1 epoch == 1 mini-batch. Otherwise, each cycle corresponds to all mini-batches of data passing through a NN.batch_size (
int
) – Size of training and test mini-batchesfull_epoch (
bool
) – If True, passes all mini-batches of training/test data at each training cycle and computes the average loss. If False, passes a single (randomly chosen) mini-batch at each cycle.swa (
bool
) – Saves the recent stochastic weights and averages them at the end of trainingperturb_weights (
bool
) – Time-dependent weight perturbation, \(w\leftarrow w + a / (1 + e)^\gamma\), where parameters a and gamma can be passed as a dictionary, together with parameter e_p determining every n-th epoch at which a perturbation is applied**print_loss (int) – Prints loss every n-th epoch
**filename (str) – Filename for model weights (appended with “_test_weights_best.pt” and “_weights_final.pt”)
**plot_training_history (bool) – Plots training and test curves vs. training cycles at the end of training
**kwargs – One can also pass kwargs for utils.datatransform class to perform the augmentation “on-the-fly” (e.g. rotation=True, gauss_nois=[20, 60], etc.)
- predict(data, **kwargs)[source]
Apply (trained model) to new data
- Parameters:
data (
ndarray
) – Input image/spectrum or batch of images/spectra**num_batches (int) – number of batches (Default: 10)
**norm (bool) – Normalize data to (0, 1) during pre-processing
**verbose (bool) – verbosity (Default: True)
- Return type:
ndarray
Variational Autoencoder (VAE)
- class atomai.models.VAE(in_dim=None, latent_dim=2, nb_classes=0, seed=0, **kwargs)[source]
Bases:
BaseVAE
Implements a standard Variational Autoencoder (VAE)
- Parameters:
in_dim (
Optional
[int
]) – Input dimensions for image data passed as (heigth, width) for grayscale data or (height, width, channels) for multichannel datalatent_dim (
int
) – Number of VAE latent dimensionsnb_classes (
int
) – Number of classes for class-conditional rVAEseed (
int
) – seed for torch and numpy (pseudo-)random numbers generators**conv_encoder (bool) – use convolutional layers in encoder
**conv_decoder (bool) – use convolutional layers in decoder
**numlayers_encoder (int) – number of layers in encoder (Default: 2)
**numlayers_decoder (int) – number of layers in decoder (Default: 2)
**numhidden_encoder (int) – number of hidden units OR conv filters in encoder (Default: 128)
**numhidden_decoder (int) – number of hidden units OR conv filters in decoder (Default: 128)
Example:
>>> input_dim = (28, 28) # Input data dimensions (without n_samples) >>> # Intitialize model >>> vae = aoi.models.VAE(input_dim) >>> # Train >>> vae.fit(imstack_train, training_cycles=100, batch_size=100) >>> # Visualize learned manifold (for 2 latent dimesnions) >>> vae.manifold2d(origin="upper", cmap="gnuplot2)
One can also pass labels to train a class-conditioned VAE
>>> # Intitialize model >>> vae = aoi.models.VAE(input_dim, nb_classes=10) >>> # Train >>> vae.fit(imstack_train, labels_train, training_cycles=100, batch_size=100) >>> # Visualize learned manifold for class 1 >>> vae.manifold2d(label=1, origin="upper", cmap="gnuplot2")
- forward_compute_elbo(x, y=None, mode='train')[source]
VAE’s forward pass with training/test loss computation
- Return type:
Tensor
- compile_trainer(train_data, test_data=None, optimizer=None, elbo_fn=None, training_cycles=100, batch_size=32, **kwargs)
Compiles model’s trainer
- Parameters:
train_data (
Tuple
[Union
[Tensor
,ndarray
]]) – Train data and (optionally) corresponding targets or labelstrain_data – Test data and (optionally) corresponding targets or labels
optimizer (
Optional
[Type
[Optimizer
]]) – Weights optimizer. Defaults to Adam with learning rate 1e-4elbo_fn (
Optional
[Callable
]) – function that calculates elbo losstraining_cycles (
int
) – Number of training iterations (aka “epochs”)batch_size (
int
) – Size of mini-batch for training**kwargs – Additional keyword arguments are ‘filename’ (for saving model) and ‘memory alloc’ (threshold for keeping data on GPU)
- Return type:
None
- decode(z_sample, y=None)
Takes a point in latent space and maps it to data space via the learned generative model
- Parameters:
z_sample (
Union
[ndarray
,Tensor
]) – coordinates in latent spacey (
Union
[int
,ndarray
,Tensor
,None
]) – label (optional)
- Return type:
ndarray
- Returns:
Generated (“decoded”) image(s)
- encode(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Mean and SD of the encoded continuous distribution, and alphas (“class probabilities”) for the encoded discrete distribution(s) (if any). For rVAE, the output is (z_mean, z_sd). For jVAE and jrVAE, the output is (z_mean, z_sd, alphas). In all the cases z_mean consists of the encoded angle as 1st dimension, encoded x- and y-shift as 2nd and 3rd dimensions (if translation is set to True), and standard VAE latent variables as 4th, 5th, …, n-th dimensions (if translation is set to True; otherwise, 2nd, 3rd, … n-th dimensions)
- encode_(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Concatenated array of encoded vectors
- encode_image_(img, **kwargs)
Crops and encodes a subimage around each pixel in the input image. The size of subimage is determined by size of images in VAE training data.
- Parameters:
img (
ndarray
) – 2D numpy array**num_batches (int) – number of batches for encoding subimages
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image and encoded array (cropping is due to finite window size)
- encode_images(imgdata, **kwargs)
Encodes every pixel of every image in image stack
- Parameters:
imgdata (
ndarray
) – 3D numpy array of images. Can also be a single 2D image**num_batches (int) – number of batches for for encoding pixels of a single image
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image stack and encoded array (cropping is due to finite window size)
- encode_trajectories(imgdata, coord_class_dict, window_size, min_length, rmax, **kwargs)
Calculates trajectories and latent variable value for each point in a trajectory.
- Parameters:
imgdata (
ndarray
) – NN output (preferable) or raw datacoord_class_dict (
Dict
[int
,ndarray
]) – atomic/defect/particle coordinateswindow_size (
int
) – size of subimages to cropmin_length (
int
) – minimum length of trajectory to be includedrmax (
int
) – maximum allowed distance (projected on xy plane) between defect in one frame and the position of its nearest neigbor in the next one**num_batches (int) – number of batches for self.encode (Default: 10)
- Return type:
Tuple
[List
[ndarray
],List
[ndarray
]]- Returns:
List of encoded trajectories and corresponding movie frame numbers
- evaluate_model()
Evaluates model on test data
- fit(X_train, y_train=None, X_test=None, y_test=None, loss='mse', **kwargs)[source]
Trains VAE model
- Parameters:
X_train (
Union
[ndarray
,Tensor
]) – For images, 3D or 4D stack of training images with dimensions (n_images, height, width) for grayscale data or or (n_images, height, width, channels) for multi-channel data. For spectra, 2D stack of spectra with dimensions (length,)y_train (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of training images/spectraX_test (
Union
[ndarray
,Tensor
,None
]) – 3D or 4D stack of test images or 2D stack of spectra with the same dimensions as for the X_train (Default: None)y_test (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of test images/spectraloss (
str
) – reconstruction loss function, “ce” or “mse” (Default: “mse”)**capacity (list) – List containing (max_capacity, num_iters, gamma) parameters to control the capacity of the latent channel. Based on https://arxiv.org/pdf/1804.03599.pdf
**filename (str) – file path for saving model aftereach training cycle (“epoch”)
- Return type:
None
- kld_normal(z, q_param, p_param=None)
Calculates KL divergence term between two normal distributions or (if p_param = None) between normal and standard normal distributions
- Parameters:
z (
Tensor
) – latent vector (reparametrized)q_param (
Tuple
[Tensor
]) – tuple with mean and SD of the 1st distributionp_param (
Optional
[Tuple
[Tensor
]]) – tuple with mean and SD of the 2nd distribution (optional)
- Return type:
Tensor
- load_weights(filepath)
Loads saved weights
- Return type:
None
- classmethod log_normal(x, mu, log_sd)
Computes log-pdf for a normal distribution
- Return type:
Tensor
- classmethod log_unit_normal(x)
Computes log-pdf of a unit normal distribution
- Return type:
Tensor
- manifold2d(**kwargs)
Performs mapping from latent space to data space allowing the learned manifold to be visualized. This works only for 2d latent variable (not counting angle & translation dimensions)
- Parameters:
**d (int) – grid size
**l1 (list) – range of 1st latent variable
**l2 (list) – range of 2nd latent variable
**label (int) – label in class-conditioned (r)VAE
**disc_idx (int) – discrete “class”
**cmap (str) – color map (Default: gnuplot)
**draw_grid (bool) – plot semi-transparent grid
**origin (str) – plot origin (e.g. ‘lower’)
- Return type:
None
- manifold_traversal(cont_idx, d=10, cont_idx_fixed=0, plot=True, **kwargs)
Latent space traversals for joint continuous and discrete latent representations
- Return type:
ndarray
- print_statistics(e)
Prints training and (optionally) test loss after each training cycle
- reconstruct(x_new, **kwargs)
Forward prediction with uncertainty quantification by sampling from the encoded mean and std. Works only for regular VAE (and not for rVAE)
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**label (int) – class to be reconstructed (for cVAE, crVAE, jVAE, and jrVAE)
**num_samples (int) – number of samples to generate from normal distribution
- Return type:
ndarray
- Returns:
Ensemble of “decoded” images
- classmethod reparameterize(z_mean, z_sd)
Reparameterization trick for continuous distributions
- Return type:
Tensor
- classmethod reparameterize_discrete(alpha, tau)
Reparameterization trick for discrete gumbel-softmax distributions
- save_model(*args)
Saves trained weights and the key model parameters
- Return type:
None
- save_weights(*args)
Saves trained weights
- Return type:
None
- set_data(X_train, y_train=None, X_test=None, y_test=None, memory_alloc=4)
Initializes train and (optionally) test data loaders
- Return type:
None
- set_decoder(decoder_net)
Sets a decoder network only
- Return type:
None
- set_encoder(encoder_net)
Sets an encoder network only
- Return type:
None
- set_model(encoder_net, decoder_net)
Sets encoder and decoder models
- Return type:
None
- train_epoch()
Trains a single epoch
- classmethod visualize_manifold_learning(frames_dir, **kwargs)
Creates and stores a video showing evolution of learned 2D manifold during rVAE’s training
- Parameters:
frames_dir (
str
) – directory with snapshots of manifold as .png files (the files should be named as “1.png”, “2.png”, etc.)**moviename (str) – name of the movie
**frame_duration (int) – duration of each movie frame
- Return type:
None
Rotational Variational Autoencoder (rVAE)
- class atomai.models.rVAE(in_dim=None, latent_dim=2, nb_classes=0, translation=True, seed=0, **kwargs)[source]
Bases:
BaseVAE
Implements rotationally and translationally invariant Variational Autoencoder (VAE) based on the idea of “spatial decoder” by Bepler et al. in arXiv:1909.11663. In addition, this class allows implementating the class-conditioned VAE and skip-VAE (arXiv:1807.04863) with rotational and translational variance.
- Parameters:
in_dim (
Optional
[int
]) – Input dimensions for image data passed as (heigth, width) for grayscale data or (height, width, channels) for multichannel datalatent_dim (
int
) – Number of VAE latent dimensions associated with image contentnb_classes (
int
) – Number of classes for class-conditional rVAEtranslation (
bool
) – account for xy shifts of image content (Default: True)seed (
int
) – seed for torch and numpy (pseudo-)random numbers generators**conv_encoder (bool) – use convolutional layers in encoder
**numlayers_encoder (int) – number of layers in encoder (Default: 2)
**numlayers_decoder (int) – number of layers in decoder (Default: 2)
**numhidden_encoder (int) – number of hidden units OR conv filters in encoder (Default: 128)
**numhidden_decoder (int) – number of hidden units in decoder (Default: 128)
**skip (bool) – uses generative skip model with residual paths between latents and decoder layers (Default: False)
Example:
>>> input_dim = (28, 28) # input dimensions >>> # Intitialize model >>> rvae = aoi.models.rVAE(input_dim) >>> # Train >>> rvae.fit(imstack_train, training_cycles=100, batch_size=100, rotation_prior=np.pi/2) >>> rvae.manifold2d(origin="upper", cmap="gnuplot2")
One can also pass labels to train a class-conditioned rVAE
>>> # Intitialize model >>> rvae = aoi.models.rVAE(input_dim, nb_classes=10) >>> # Train >>> rvae.fit(imstack_train, labels_train, training_cycles=100, >>> batch_size=100, rotation_prior=np.pi/2) >>> # Visualize learned manifold for class 1 >>> rvae.manifold2d(label=1, origin="upper", cmap="gnuplot2")
- forward_compute_elbo(x, y=None, mode='train')[source]
rVAE’s forward pass with training/test loss computation
- Return type:
Tensor
- fit(X_train, y_train=None, X_test=None, y_test=None, loss='mse', **kwargs)[source]
Trains rVAE model
- Parameters:
X_train (
Union
[ndarray
,Tensor
]) – 3D or 4D stack of training images with dimensions (n_images, height, width) for grayscale data or or (n_images, height, width, channels) for multi-channel datay_train (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of training imagesX_test (
Union
[ndarray
,Tensor
,None
]) – 3D or 4D stack of test images with the same dimensions as for the X_train (Default: None)y_test (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of test imagesloss (
str
) – reconstruction loss function, “ce” or “mse” (Default: “mse”)**translation_prior (float) – translation prior
**rotation_prior (float) – rotational prior
**capacity (list) – List containing (max_capacity, num_iters, gamma) parameters to control the capacity of the latent channel. Based on https://arxiv.org/pdf/1804.03599.pdf
**filename (str) – file path for saving model aftereach training cycle (“epoch”)
**recording (bool) – saves a learned 2d manifold at each training step
- Return type:
None
- compile_trainer(train_data, test_data=None, optimizer=None, elbo_fn=None, training_cycles=100, batch_size=32, **kwargs)
Compiles model’s trainer
- Parameters:
train_data (
Tuple
[Union
[Tensor
,ndarray
]]) – Train data and (optionally) corresponding targets or labelstrain_data – Test data and (optionally) corresponding targets or labels
optimizer (
Optional
[Type
[Optimizer
]]) – Weights optimizer. Defaults to Adam with learning rate 1e-4elbo_fn (
Optional
[Callable
]) – function that calculates elbo losstraining_cycles (
int
) – Number of training iterations (aka “epochs”)batch_size (
int
) – Size of mini-batch for training**kwargs – Additional keyword arguments are ‘filename’ (for saving model) and ‘memory alloc’ (threshold for keeping data on GPU)
- Return type:
None
- decode(z_sample, y=None)
Takes a point in latent space and maps it to data space via the learned generative model
- Parameters:
z_sample (
Union
[ndarray
,Tensor
]) – coordinates in latent spacey (
Union
[int
,ndarray
,Tensor
,None
]) – label (optional)
- Return type:
ndarray
- Returns:
Generated (“decoded”) image(s)
- encode(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Mean and SD of the encoded continuous distribution, and alphas (“class probabilities”) for the encoded discrete distribution(s) (if any). For rVAE, the output is (z_mean, z_sd). For jVAE and jrVAE, the output is (z_mean, z_sd, alphas). In all the cases z_mean consists of the encoded angle as 1st dimension, encoded x- and y-shift as 2nd and 3rd dimensions (if translation is set to True), and standard VAE latent variables as 4th, 5th, …, n-th dimensions (if translation is set to True; otherwise, 2nd, 3rd, … n-th dimensions)
- encode_(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Concatenated array of encoded vectors
- encode_image_(img, **kwargs)
Crops and encodes a subimage around each pixel in the input image. The size of subimage is determined by size of images in VAE training data.
- Parameters:
img (
ndarray
) – 2D numpy array**num_batches (int) – number of batches for encoding subimages
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image and encoded array (cropping is due to finite window size)
- encode_images(imgdata, **kwargs)
Encodes every pixel of every image in image stack
- Parameters:
imgdata (
ndarray
) – 3D numpy array of images. Can also be a single 2D image**num_batches (int) – number of batches for for encoding pixels of a single image
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image stack and encoded array (cropping is due to finite window size)
- encode_trajectories(imgdata, coord_class_dict, window_size, min_length, rmax, **kwargs)
Calculates trajectories and latent variable value for each point in a trajectory.
- Parameters:
imgdata (
ndarray
) – NN output (preferable) or raw datacoord_class_dict (
Dict
[int
,ndarray
]) – atomic/defect/particle coordinateswindow_size (
int
) – size of subimages to cropmin_length (
int
) – minimum length of trajectory to be includedrmax (
int
) – maximum allowed distance (projected on xy plane) between defect in one frame and the position of its nearest neigbor in the next one**num_batches (int) – number of batches for self.encode (Default: 10)
- Return type:
Tuple
[List
[ndarray
],List
[ndarray
]]- Returns:
List of encoded trajectories and corresponding movie frame numbers
- evaluate_model()
Evaluates model on test data
- kld_normal(z, q_param, p_param=None)
Calculates KL divergence term between two normal distributions or (if p_param = None) between normal and standard normal distributions
- Parameters:
z (
Tensor
) – latent vector (reparametrized)q_param (
Tuple
[Tensor
]) – tuple with mean and SD of the 1st distributionp_param (
Optional
[Tuple
[Tensor
]]) – tuple with mean and SD of the 2nd distribution (optional)
- Return type:
Tensor
- load_weights(filepath)
Loads saved weights
- Return type:
None
- classmethod log_normal(x, mu, log_sd)
Computes log-pdf for a normal distribution
- Return type:
Tensor
- classmethod log_unit_normal(x)
Computes log-pdf of a unit normal distribution
- Return type:
Tensor
- manifold2d(**kwargs)
Performs mapping from latent space to data space allowing the learned manifold to be visualized. This works only for 2d latent variable (not counting angle & translation dimensions)
- Parameters:
**d (int) – grid size
**l1 (list) – range of 1st latent variable
**l2 (list) – range of 2nd latent variable
**label (int) – label in class-conditioned (r)VAE
**disc_idx (int) – discrete “class”
**cmap (str) – color map (Default: gnuplot)
**draw_grid (bool) – plot semi-transparent grid
**origin (str) – plot origin (e.g. ‘lower’)
- Return type:
None
- manifold_traversal(cont_idx, d=10, cont_idx_fixed=0, plot=True, **kwargs)
Latent space traversals for joint continuous and discrete latent representations
- Return type:
ndarray
- print_statistics(e)
Prints training and (optionally) test loss after each training cycle
- reconstruct(x_new, **kwargs)
Forward prediction with uncertainty quantification by sampling from the encoded mean and std. Works only for regular VAE (and not for rVAE)
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**label (int) – class to be reconstructed (for cVAE, crVAE, jVAE, and jrVAE)
**num_samples (int) – number of samples to generate from normal distribution
- Return type:
ndarray
- Returns:
Ensemble of “decoded” images
- classmethod reparameterize(z_mean, z_sd)
Reparameterization trick for continuous distributions
- Return type:
Tensor
- classmethod reparameterize_discrete(alpha, tau)
Reparameterization trick for discrete gumbel-softmax distributions
- save_model(*args)
Saves trained weights and the key model parameters
- Return type:
None
- save_weights(*args)
Saves trained weights
- Return type:
None
- set_data(X_train, y_train=None, X_test=None, y_test=None, memory_alloc=4)
Initializes train and (optionally) test data loaders
- Return type:
None
- set_decoder(decoder_net)
Sets a decoder network only
- Return type:
None
- set_encoder(encoder_net)
Sets an encoder network only
- Return type:
None
- set_model(encoder_net, decoder_net)
Sets encoder and decoder models
- Return type:
None
- train_epoch()
Trains a single epoch
- classmethod visualize_manifold_learning(frames_dir, **kwargs)
Creates and stores a video showing evolution of learned 2D manifold during rVAE’s training
- Parameters:
frames_dir (
str
) – directory with snapshots of manifold as .png files (the files should be named as “1.png”, “2.png”, etc.)**moviename (str) – name of the movie
**frame_duration (int) – duration of each movie frame
- Return type:
None
Joint Variational Autoencoder (jVAE)
- class atomai.models.jVAE(in_dim=None, latent_dim=2, discrete_dim=[2], nb_classes=0, seed=0, **kwargs)[source]
Bases:
BaseVAE
VAE for joint (continuous + discrete) latent representations
- Parameters:
in_dim (
Optional
[int
]) – Input dimensions for image data passed as (heigth, width) for grayscale data or (height, width, channels) for multichannel datalatent_dim (
int
) – Number of latent dimensions associated with image contentdiscrete_dim (
List
[int
]) – List specifying dimensionalities of discrete (Gumbel-Softmax) latent variables associated with image contentnb_classes (
int
) – Number of classes for class-conditional VAE (leave it at 0 to learn discrete latent reprenetations)seed (
int
) – seed for torch and numpy (pseudo-)random numbers generators**conv_encoder (bool) – use convolutional layers in encoder
**conv_decoder (bool) – use convolutional layers in decoder
**numlayers_encoder (int) – number of layers in encoder (Default: 2)
**numlayers_decoder (int) – number of layers in decoder (Default: 2)
**numhidden_encoder (int) – number of hidden units OR conv filters in encoder (Default: 128)
**numhidden_decoder (int) – number of hidden units in decoder (Default: 128)
**skip (bool) – uses generative skip model with residual paths between latents and decoder layers (Default: False)
Example:
>>> input_dim = (28, 28) # intput dimensions >>> # Intitialize model >>> jvae = aoi.models.jVAE(input_dim, latent_dim=2, discrete_dim=[10], >>> numlayers_encoder=3, numhidden_encoder=512, >>> numlayers_decoder=3, numhidden_decoder=512) >>> # Train >>> jvae.fit(imstack_train, training_cycles=100, batch_size=100) >>> # View a traversal of the learned manifold >>> jvae.manifold_traversal(cont_idx=1, origin="upper", cmap="gnuplot2")
- forward_compute_elbo(x, y=None, mode='train')[source]
Joint VAE’s forward pass with training/test loss computation
- Return type:
Tensor
- fit(X_train, y_train=None, X_test=None, y_test=None, loss='mse', **kwargs)[source]
Trains joint VAE model
- Parameters:
X_train (
Union
[ndarray
,Tensor
]) – For images, 3D or 4D stack of training images with dimensions (n_images, height, width) for grayscale data or or (n_images, height, width, channels) for multi-channel data. For spectra, 2D stack of spectra with dimensions (length,)y_train (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of training images/spectraX_test (
Union
[ndarray
,Tensor
,None
]) – 3D or 4D stack of test images or 2D stack of spectra with the same dimensions as for the X_train (Default: None)y_test (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of test images/spectraloss (
str
) – reconstruction loss function, “ce” or “mse” (Default: “mse”)**cont_capacity (list) – List containing (max_capacity, num_iters, gamma) parameters to control the capacity of the continuous latent channel. Default values: [5.0, 25000, 30]. Based on https://arxiv.org/pdf/1804.03599.pdf & https://arxiv.org/abs/1804.00104
**disc_capacity (list) – List containing (max_capacity, num_iters, gamma) parameters to control the capacity of the discrete latent channel(s). Default values: [5.0, 25000, 30]. Based on https://arxiv.org/pdf/1804.03599.pdf & https://arxiv.org/abs/1804.00104
**filename (str) – file path for saving model aftereach training cycle (“epoch”)
- Return type:
None
- compile_trainer(train_data, test_data=None, optimizer=None, elbo_fn=None, training_cycles=100, batch_size=32, **kwargs)
Compiles model’s trainer
- Parameters:
train_data (
Tuple
[Union
[Tensor
,ndarray
]]) – Train data and (optionally) corresponding targets or labelstrain_data – Test data and (optionally) corresponding targets or labels
optimizer (
Optional
[Type
[Optimizer
]]) – Weights optimizer. Defaults to Adam with learning rate 1e-4elbo_fn (
Optional
[Callable
]) – function that calculates elbo losstraining_cycles (
int
) – Number of training iterations (aka “epochs”)batch_size (
int
) – Size of mini-batch for training**kwargs – Additional keyword arguments are ‘filename’ (for saving model) and ‘memory alloc’ (threshold for keeping data on GPU)
- Return type:
None
- decode(z_sample, y=None)
Takes a point in latent space and maps it to data space via the learned generative model
- Parameters:
z_sample (
Union
[ndarray
,Tensor
]) – coordinates in latent spacey (
Union
[int
,ndarray
,Tensor
,None
]) – label (optional)
- Return type:
ndarray
- Returns:
Generated (“decoded”) image(s)
- encode(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Mean and SD of the encoded continuous distribution, and alphas (“class probabilities”) for the encoded discrete distribution(s) (if any). For rVAE, the output is (z_mean, z_sd). For jVAE and jrVAE, the output is (z_mean, z_sd, alphas). In all the cases z_mean consists of the encoded angle as 1st dimension, encoded x- and y-shift as 2nd and 3rd dimensions (if translation is set to True), and standard VAE latent variables as 4th, 5th, …, n-th dimensions (if translation is set to True; otherwise, 2nd, 3rd, … n-th dimensions)
- encode_(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Concatenated array of encoded vectors
- encode_image_(img, **kwargs)
Crops and encodes a subimage around each pixel in the input image. The size of subimage is determined by size of images in VAE training data.
- Parameters:
img (
ndarray
) – 2D numpy array**num_batches (int) – number of batches for encoding subimages
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image and encoded array (cropping is due to finite window size)
- encode_images(imgdata, **kwargs)
Encodes every pixel of every image in image stack
- Parameters:
imgdata (
ndarray
) – 3D numpy array of images. Can also be a single 2D image**num_batches (int) – number of batches for for encoding pixels of a single image
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image stack and encoded array (cropping is due to finite window size)
- encode_trajectories(imgdata, coord_class_dict, window_size, min_length, rmax, **kwargs)
Calculates trajectories and latent variable value for each point in a trajectory.
- Parameters:
imgdata (
ndarray
) – NN output (preferable) or raw datacoord_class_dict (
Dict
[int
,ndarray
]) – atomic/defect/particle coordinateswindow_size (
int
) – size of subimages to cropmin_length (
int
) – minimum length of trajectory to be includedrmax (
int
) – maximum allowed distance (projected on xy plane) between defect in one frame and the position of its nearest neigbor in the next one**num_batches (int) – number of batches for self.encode (Default: 10)
- Return type:
Tuple
[List
[ndarray
],List
[ndarray
]]- Returns:
List of encoded trajectories and corresponding movie frame numbers
- evaluate_model()
Evaluates model on test data
- kld_normal(z, q_param, p_param=None)
Calculates KL divergence term between two normal distributions or (if p_param = None) between normal and standard normal distributions
- Parameters:
z (
Tensor
) – latent vector (reparametrized)q_param (
Tuple
[Tensor
]) – tuple with mean and SD of the 1st distributionp_param (
Optional
[Tuple
[Tensor
]]) – tuple with mean and SD of the 2nd distribution (optional)
- Return type:
Tensor
- load_weights(filepath)
Loads saved weights
- Return type:
None
- classmethod log_normal(x, mu, log_sd)
Computes log-pdf for a normal distribution
- Return type:
Tensor
- classmethod log_unit_normal(x)
Computes log-pdf of a unit normal distribution
- Return type:
Tensor
- manifold2d(**kwargs)
Performs mapping from latent space to data space allowing the learned manifold to be visualized. This works only for 2d latent variable (not counting angle & translation dimensions)
- Parameters:
**d (int) – grid size
**l1 (list) – range of 1st latent variable
**l2 (list) – range of 2nd latent variable
**label (int) – label in class-conditioned (r)VAE
**disc_idx (int) – discrete “class”
**cmap (str) – color map (Default: gnuplot)
**draw_grid (bool) – plot semi-transparent grid
**origin (str) – plot origin (e.g. ‘lower’)
- Return type:
None
- manifold_traversal(cont_idx, d=10, cont_idx_fixed=0, plot=True, **kwargs)
Latent space traversals for joint continuous and discrete latent representations
- Return type:
ndarray
- print_statistics(e)
Prints training and (optionally) test loss after each training cycle
- reconstruct(x_new, **kwargs)
Forward prediction with uncertainty quantification by sampling from the encoded mean and std. Works only for regular VAE (and not for rVAE)
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**label (int) – class to be reconstructed (for cVAE, crVAE, jVAE, and jrVAE)
**num_samples (int) – number of samples to generate from normal distribution
- Return type:
ndarray
- Returns:
Ensemble of “decoded” images
- classmethod reparameterize(z_mean, z_sd)
Reparameterization trick for continuous distributions
- Return type:
Tensor
- classmethod reparameterize_discrete(alpha, tau)
Reparameterization trick for discrete gumbel-softmax distributions
- save_model(*args)
Saves trained weights and the key model parameters
- Return type:
None
- save_weights(*args)
Saves trained weights
- Return type:
None
- set_data(X_train, y_train=None, X_test=None, y_test=None, memory_alloc=4)
Initializes train and (optionally) test data loaders
- Return type:
None
- set_decoder(decoder_net)
Sets a decoder network only
- Return type:
None
- set_encoder(encoder_net)
Sets an encoder network only
- Return type:
None
- set_model(encoder_net, decoder_net)
Sets encoder and decoder models
- Return type:
None
- train_epoch()
Trains a single epoch
- classmethod visualize_manifold_learning(frames_dir, **kwargs)
Creates and stores a video showing evolution of learned 2D manifold during rVAE’s training
- Parameters:
frames_dir (
str
) – directory with snapshots of manifold as .png files (the files should be named as “1.png”, “2.png”, etc.)**moviename (str) – name of the movie
**frame_duration (int) – duration of each movie frame
- Return type:
None
Joint Rotational Variational Autoencoder (jrVAE)
- class atomai.models.jrVAE(in_dim=None, latent_dim=2, discrete_dim=[2], nb_classes=0, translation=True, seed=0, **kwargs)[source]
Bases:
BaseVAE
Rotationally-invariant VAE for joint continuous and discrete latent representations.
- Parameters:
in_dim (
Optional
[int
]) – Input dimensions for image data passed as (heigth, width) for grayscale data or (height, width, channels) for multichannel datalatent_dim (
int
) – Number of latent dimensions associated with image contentdiscrete_dim (
List
[int
]) – List specifying dimensionalities of discrete (Gumbel-Softmax) latent variables associated with image contentnb_classes (
int
) – Number of classes for class-conditional VAE. (leave it at 0 to learn discrete latent reprenetations)translation (
bool
) – account for xy shifts of image content (Default: True)seed (
int
) – seed for torch and numpy (pseudo-)random numbers generators**conv_encoder (bool) – use convolutional layers in encoder
**numlayers_encoder (int) – number of layers in encoder (Default: 2)
**numlayers_decoder (int) – number of layers in decoder (Default: 2)
**numhidden_encoder (int) – number of hidden units OR conv filters in encoder (Default: 128)
**numhidden_decoder (int) – number of hidden units in decoder (Default: 128)
**skip (bool) – uses generative skip model with residual paths between latents and decoder layers (Default: False)
Example:
>>> input_dim = (28, 28) # intput dimensions >>> # Intitialize model >>> jrvae = aoi.models.jVAE(input_dim, latent_dim=2, discrete_dim=[10], >>> numlayers_encoder=3, numhidden_encoder=512, >>> numlayers_decoder=3, numhidden_decoder=512) >>> # Train >>> jrvae.fit(imstack_train, training_cycles=100, batch_size=100, rotation_prior=np.pi/4) >>> jrvae.manifold2d(origin="upper", cmap="gnuplot2")
- forward_compute_elbo(x, y=None, mode='train')[source]
Joint rVAE’s forward pass with training/test loss computation
- Return type:
Tensor
- fit(X_train, y_train=None, X_test=None, y_test=None, loss='mse', verbose='True', **kwargs)[source]
Trains joint rVAE model
- Parameters:
X_train (
Union
[ndarray
,Tensor
]) – 3D or 4D stack of training images with dimensions (n_images, height, width) for grayscale data or or (n_images, height, width, channels) for multi-channel datay_train (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of training imagesX_test (
Union
[ndarray
,Tensor
,None
]) – 3D or 4D stack of test images with the same dimensions as for the X_train (Default: None)y_test (
Union
[ndarray
,Tensor
,None
]) – Vector with labels of dimension (n_images,), where n_images is a number of test imagesloss (
str
) – reconstruction loss function, “ce” or “mse” (Default: “mse”)**translation_prior (float) – translation prior
**rotation_prior (float) – rotational prior
**temperature (float) – Relaxation parameter for Gumbel-Softmax distribution
**cont_capacity (list) – List containing (max_capacity, num_iters, gamma) parameters to control the capacity of the continuous latent channel. Default values: [5.0, 25000, 30]. Based on https://arxiv.org/pdf/1804.03599.pdf & https://arxiv.org/abs/1804.00104
**disc_capacity (list) – List containing (max_capacity, num_iters, gamma) parameters to control the capacity of the discrete latent channel(s). Default values: [5.0, 25000, 30]. Based on https://arxiv.org/pdf/1804.03599.pdf & https://arxiv.org/abs/1804.00104
**filename (str) – file path for saving model after each training cycle (“epoch”)
verbose (
str
) – display training output, “True” or “False” (Default: “True”)
- Return type:
None
- compile_trainer(train_data, test_data=None, optimizer=None, elbo_fn=None, training_cycles=100, batch_size=32, **kwargs)
Compiles model’s trainer
- Parameters:
train_data (
Tuple
[Union
[Tensor
,ndarray
]]) – Train data and (optionally) corresponding targets or labelstrain_data – Test data and (optionally) corresponding targets or labels
optimizer (
Optional
[Type
[Optimizer
]]) – Weights optimizer. Defaults to Adam with learning rate 1e-4elbo_fn (
Optional
[Callable
]) – function that calculates elbo losstraining_cycles (
int
) – Number of training iterations (aka “epochs”)batch_size (
int
) – Size of mini-batch for training**kwargs – Additional keyword arguments are ‘filename’ (for saving model) and ‘memory alloc’ (threshold for keeping data on GPU)
- Return type:
None
- decode(z_sample, y=None)
Takes a point in latent space and maps it to data space via the learned generative model
- Parameters:
z_sample (
Union
[ndarray
,Tensor
]) – coordinates in latent spacey (
Union
[int
,ndarray
,Tensor
,None
]) – label (optional)
- Return type:
ndarray
- Returns:
Generated (“decoded”) image(s)
- encode(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Mean and SD of the encoded continuous distribution, and alphas (“class probabilities”) for the encoded discrete distribution(s) (if any). For rVAE, the output is (z_mean, z_sd). For jVAE and jrVAE, the output is (z_mean, z_sd, alphas). In all the cases z_mean consists of the encoded angle as 1st dimension, encoded x- and y-shift as 2nd and 3rd dimensions (if translation is set to True), and standard VAE latent variables as 4th, 5th, …, n-th dimensions (if translation is set to True; otherwise, 2nd, 3rd, … n-th dimensions)
- encode_(x_new, **kwargs)
Encodes input image data using a trained VAE’s encoder
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**num_batches (int) – number of batches (Default: 10)
- Return type:
Tuple
[ndarray
]- Returns:
Concatenated array of encoded vectors
- encode_image_(img, **kwargs)
Crops and encodes a subimage around each pixel in the input image. The size of subimage is determined by size of images in VAE training data.
- Parameters:
img (
ndarray
) – 2D numpy array**num_batches (int) – number of batches for encoding subimages
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image and encoded array (cropping is due to finite window size)
- encode_images(imgdata, **kwargs)
Encodes every pixel of every image in image stack
- Parameters:
imgdata (
ndarray
) – 3D numpy array of images. Can also be a single 2D image**num_batches (int) – number of batches for for encoding pixels of a single image
- Return type:
Tuple
[ndarray
,ndarray
]- Returns:
Cropped original image stack and encoded array (cropping is due to finite window size)
- encode_trajectories(imgdata, coord_class_dict, window_size, min_length, rmax, **kwargs)
Calculates trajectories and latent variable value for each point in a trajectory.
- Parameters:
imgdata (
ndarray
) – NN output (preferable) or raw datacoord_class_dict (
Dict
[int
,ndarray
]) – atomic/defect/particle coordinateswindow_size (
int
) – size of subimages to cropmin_length (
int
) – minimum length of trajectory to be includedrmax (
int
) – maximum allowed distance (projected on xy plane) between defect in one frame and the position of its nearest neigbor in the next one**num_batches (int) – number of batches for self.encode (Default: 10)
- Return type:
Tuple
[List
[ndarray
],List
[ndarray
]]- Returns:
List of encoded trajectories and corresponding movie frame numbers
- evaluate_model()
Evaluates model on test data
- kld_normal(z, q_param, p_param=None)
Calculates KL divergence term between two normal distributions or (if p_param = None) between normal and standard normal distributions
- Parameters:
z (
Tensor
) – latent vector (reparametrized)q_param (
Tuple
[Tensor
]) – tuple with mean and SD of the 1st distributionp_param (
Optional
[Tuple
[Tensor
]]) – tuple with mean and SD of the 2nd distribution (optional)
- Return type:
Tensor
- load_weights(filepath)
Loads saved weights
- Return type:
None
- classmethod log_normal(x, mu, log_sd)
Computes log-pdf for a normal distribution
- Return type:
Tensor
- classmethod log_unit_normal(x)
Computes log-pdf of a unit normal distribution
- Return type:
Tensor
- manifold2d(**kwargs)
Performs mapping from latent space to data space allowing the learned manifold to be visualized. This works only for 2d latent variable (not counting angle & translation dimensions)
- Parameters:
**d (int) – grid size
**l1 (list) – range of 1st latent variable
**l2 (list) – range of 2nd latent variable
**label (int) – label in class-conditioned (r)VAE
**disc_idx (int) – discrete “class”
**cmap (str) – color map (Default: gnuplot)
**draw_grid (bool) – plot semi-transparent grid
**origin (str) – plot origin (e.g. ‘lower’)
- Return type:
None
- manifold_traversal(cont_idx, d=10, cont_idx_fixed=0, plot=True, **kwargs)
Latent space traversals for joint continuous and discrete latent representations
- Return type:
ndarray
- print_statistics(e)
Prints training and (optionally) test loss after each training cycle
- reconstruct(x_new, **kwargs)
Forward prediction with uncertainty quantification by sampling from the encoded mean and std. Works only for regular VAE (and not for rVAE)
- Parameters:
x_new (
Union
[ndarray
,Tensor
]) – image array to encode**label (int) – class to be reconstructed (for cVAE, crVAE, jVAE, and jrVAE)
**num_samples (int) – number of samples to generate from normal distribution
- Return type:
ndarray
- Returns:
Ensemble of “decoded” images
- classmethod reparameterize(z_mean, z_sd)
Reparameterization trick for continuous distributions
- Return type:
Tensor
- classmethod reparameterize_discrete(alpha, tau)
Reparameterization trick for discrete gumbel-softmax distributions
- save_model(*args)
Saves trained weights and the key model parameters
- Return type:
None
- save_weights(*args)
Saves trained weights
- Return type:
None
- set_data(X_train, y_train=None, X_test=None, y_test=None, memory_alloc=4)
Initializes train and (optionally) test data loaders
- Return type:
None
- set_decoder(decoder_net)
Sets a decoder network only
- Return type:
None
- set_encoder(encoder_net)
Sets an encoder network only
- Return type:
None
- set_model(encoder_net, decoder_net)
Sets encoder and decoder models
- Return type:
None
- train_epoch()
Trains a single epoch
- classmethod visualize_manifold_learning(frames_dir, **kwargs)
Creates and stores a video showing evolution of learned 2D manifold during rVAE’s training
- Parameters:
frames_dir (
str
) – directory with snapshots of manifold as .png files (the files should be named as “1.png”, “2.png”, etc.)**moviename (str) – name of the movie
**frame_duration (int) – duration of each movie frame
- Return type:
None
Deep Kernel Learning
- class atomai.models.dklGPR(indim, embedim=2, shared_embedding_space=True, **kwargs)[source]
Bases:
dklGPTrainer
Deep kernel learning (DKL)-based Gaussian process regression (GPR)
- Parameters:
indim (
int
) – input feature dimensionembedim (
int
) – embedding dimension (determines dimensionality of kernel space)shared_embedding_space (
bool
) – use one embedding space for all target outputs
- Keyword Arguments:
device – Sets device to which model and data will be moved. Defaults to ‘cuda:0’ if a GPU is available and to CPU otherwise.
precision – Sets tensor types for ‘single’ (torch.float32) or ‘double’ (torch.float64) precision
seed – Seed for enforcing reproducibility
Examples
Train a DKL-GPR model with high-dimensional inputs X and outputs y:
>>> data_dim = X.shape[-1] # X dimensions are n_samples x d >>> dklgp = aoi.models.dklGPR(data_dim, embedim=2, precision="double") >>> dklgp.fit(X, y, training_cycles=100, lr=1e-2)
Make a prediction on new data (mean and variance for each ‘test’ point):
>>> mean, var = dklgp.predict(X_test, batch_size=len(X_test))
Alternatively, one can obtain a prediction as follows:
>>> samples = dklgp.sample_from_posterior(X_test, num_samples=1000) >>> mean, var = samples.mean(0), samples.var(0)
- fit(X, y, training_cycles=1, **kwargs)[source]
Initializes and trains a deep kernel GP model
- Parameters:
X (
Union
[Tensor
,ndarray
]) – Input training data (aka features) of N x input_dim dimensionsy (
Union
[Tensor
,ndarray
]) – Output targets of batch_size x N or N (if batch_size=1) dimensionstraining_cycles (
int
) – Number of training epochs
- Keyword Arguments:
feature_extractor – (Optional) Custom neural network for feature extractor. Must take input/feature dims and embedding dims as its arguments.
freeze_weights – Freezes weights of feature extractor, that is, they are not passed to the optimizer. Used for a transfer learning.
lr – learning rate (Default: 0.01)
print_loss – print loss at every n-th training cycle (epoch)
- Return type:
None
- fit_ensemble(X, y, training_cycles=1, n_models=5, **kwargs)[source]
Initializes and trains an ensemble of deep kernel GP model
- Parameters:
X (
Union
[Tensor
,ndarray
]) – Input training data (aka features) of N x input_dim dimensionsy (
Union
[Tensor
,ndarray
]) – Output targets of batch_size x N or N (if batch_size=1) dimensionstraining_cycles (
int
) – Number of training epochsn_models (
int
) – Number of models in ensemble
- Keyword Arguments:
feature_extractor – (Optional) Custom neural network for feature extractor. Must take input/feature dims and embedding dims as its arguments.
freeze_weights – Freezes weights of feature extractor, that is, they are not passed to the optimizer. Used for a transfer learning.
lr – learning rate (Default: 0.01)
print_loss – print loss at every n-th training cycle (epoch)
- Return type:
None
- sample_from_posterior(X, num_samples=1000)[source]
Computes the posterior over model outputs at the provided points (X) and samples from it
- Return type:
ndarray
- thompson(X_cand, scalarize_func=None, maximize=True)[source]
Thompson sampling for selecting the next measurement point
- Return type:
Tuple
[ndarray
,int
]
- predict(x_new, **kwargs)[source]
Prediction of mean and variance using the trained model
- Return type:
Tuple
[ndarray
]
- embed(x_new, **kwargs)[source]
Embeds the input data to a “latent” space using a trained feature extractor NN.
- Return type:
Tensor
- decode(z_emb)[source]
“Decodes” the latent space variables into the target space using a trained Gaussian process model (i.e., “GP layer” of DKL-GP)
- Return type:
Tuple
[Tensor
]
- compile_multi_model_trainer(X, y, training_cycles=1, **kwargs)
Initializes deep kernel (feature extractor NNs + base kernels), sets optimizer and “loss” function. For vector-valued functions (multiple outputs), it assumes one latent space per output, that is, the number of neural networks is equal to the number of Gaussian processes. For example, if the outputs are spectra of length 128, one will have 128 neural networks and 128 GPs trained in parallel. It can be also used for training an ensembles of models for the same scalar output.
- Return type:
None
- compile_trainer(X, y, training_cycles=1, **kwargs)
Initializes deep kernel (feature extractor NN + base kernel), sets optimizer and “loss” function. For vector-valued functions (multiple outputs), it assumes a shared latent space, that is, a single neural network is connected to multiple Gaussian processes.
- Parameters:
X (
Union
[Tensor
,ndarray
]) – Input training data (aka features) of N x input_dim dimensionsy (
Union
[Tensor
,ndarray
]) – Output targets of batch_size x N or N (if batch_size=1) dimensionstraining_cycles (
int
) – Number of training epochs
- Keyword Arguments:
feature_extractor – (Optional) Custom neural network for feature extractor. Must take input/feature dims and embedding dims as its arguments.
grid_size – Grid size for structured kernel interpolation (Default: 50)
freeze_weights – Freezes weights of feature extractor, that is, they are not passed to the optimizer. Used for a transfer learning.
lr – learning rate (Default: 0.01)
- Return type:
None
- print_statistics(e)
- run(X=None, y=None, training_cycles=1, **kwargs)
Initializes and trains a deep kernel GP model
- Parameters:
X (
Union
[Tensor
,ndarray
,None
]) – Input training data (aka features) of N x input_dim dimensionsy (
Union
[Tensor
,ndarray
,None
]) – Output targets of batch_size x N or N (if batch_size=1) dimensionstraining_cycles (
int
) – Number of training epochs
- Keyword Arguments:
feature_extractor – (Optional) Custom neural network for feature extractor
freeze_weights – Freezes weights of feature extractor, that is, they are not passed to the optimizer. Used for a transfer learning.
grid_size – Grid size for structured kernel interpolation (Default: 50)
lr – learning rate (Default: 0.01)
print_loss – print loss at every n-th training cycle (epoch)
- Return type:
Type
[ExactGP
]
- save_weights(filename)
Saves weights of the feature extractor.
- Return type:
None
- set_data(x, y=None, device=None)
Data preprocessing. Casts data array to a selected tensor type and moves it to a selected devive.
- Return type:
Tuple
[tensor
]
- train_step()
Single training step with backpropagation to computegradients and optimizes weights.
- Return type:
None
Load trained models
- atomai.models.load_ensemble(filepath)[source]
Loads trained ensemble models
- Parameters:
meta_state_dict (str) – filepath to dictionary with trained weights and key information about model’s structure
- Return type:
Tuple
[Type
[Module
],Dict
[int
,Dict
[str
,Tensor
]]]- Returns:
Single model with averaged weights and dictionary with weights of all models