AtomAI Models
Segmentor
- class atomai.models.Segmentor(model='Unet', nb_classes=1, **kwargs)[source]
Bases:
atomai.trainers.trainer.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:
atomai.trainers.trainer.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:
atomai.models.dgm.vae.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:
atomai.models.dgm.vae.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:
atomai.models.dgm.vae.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:
atomai.models.dgm.vae.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', **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”)
- 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:
atomai.trainers.gptrainer.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