Trainers and Predictors

Single Model Trainers

Base Trainer Class

class atomai.trainers.BaseTrainer

Base trainer class for training semantic segmentation and image-to-spectrum/spectrum-to-image deep learning models as well as custom PyTorch neural networks

Example:

>>> # Load 4 numpy arrays with training and test data
>>> dataset = np.load('training_data.npz')
>>> images = dataset['X_train']
>>> labels = dataset['y_train']
>>> images_test = dataset['X_test']
>>> labels_test = dataset['y_test']
>>> # Initialize a trainer
>>> t = BaseTrainer()
>>> # Set a model
>>> t.set_model(atomai.nets.Unet(), nb_classes=1)
>>> # Compile trainer
>>> t.compile_trainer(
>>>     (images, labels, images_test_1, labels_test_1),
>>>     loss="ce", full_epoch=True, training_cycles=25, swa=True)
>>> # Train and save model's weights
>>> t.fit()
>>> t.save_model("my_model")

set_data(X_train, y_train, X_test, y_test, **kwargs)

Sets training and test data by initializing PyTorch dataloaders or creating a list of PyTorch tensors from which it will randomly choose an element at each training iteration.

Parameters:

• X_train (Union[Tensor, ndarray]) – Training data

• y_train (Union[Tensor, ndarray]) – Training data labels/ground-truth

• X_test (Union[Tensor, ndarray]) – Test data

• y_test (Union[Tensor, ndarray]) – Test data labels/ground-truth

• memory_alloc – threshold (in GB) for holding all training data on GPU

Return type:

None

set_model(model, nb_classes=None)

Sets a neural network model and a number of classes (if any)

Parameters:

• model (Type[Module]) – initialized PyTorch model

• nb_classes (Optional[int]) – number of classes in classification scheme (if any)

Return type:

None

get_loss_fn(loss='mse', nb_classes=None)

Returns a loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems)

Return type:

None

train_step(feat, tar)

Propagates image(s) through a network to get model’s prediction and compares predicted value with ground truth; then performs backpropagation to compute gradients and optimizes weights.

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

Tuple[float]

test_step(feat, tar)

Forward pass for test data with deactivated autograd engine

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

float

step(e)

Single train-test step which passes a single mini-batch (for both training and testing), i.e. 1 “epoch” = 1 mini-batch

Return type:

None

step_full()

A standard PyTorch training loop where all available mini-batches are passed at a single step/epoch

Return type:

None

eval_model()

Evaluates model on the entire dataset

Return type:

None

dataloader(batch_num, mode='train')

Generates input training data with images/spectra and the associated labels (spectra/images)

Return type:

Tuple[Tensor]

save_model(*args)

Saves trained weights, optimizer and key information about model’s architecture (the latter works only for built-in AtomAI models)

Return type:

None

print_statistics(e, **kwargs)

Print loss and (optionally) accuracy score on train and test data, as well as GPU memory usage.

Return type:

None

accuracy_fn(*args)

Computes accuracy score

Return type:

None

weight_perturbation(e)

Time-dependent weights perturbation (role of time is played by “epoch” number)

Return type:

None

save_running_weights(e)

Saves running weights (for stochastic weights averaging)

Return type:

None

data_augmentation(augment_fn)

Set up data augmentation. To use it, pass a function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

Return type:

None

compile_trainer(train_data=None, loss='ce', optimizer=None, training_cycles=1000, batch_size=32, compute_accuracy=False, full_epoch=False, swa=False, perturb_weights=False, **kwargs)

Compile a trainer

Parameters:

• train_data (Union[Tuple[Tensor], Tuple[ndarray], None]) – 4-element tuple of ndarrays or torch tensors (train_data, train_labels, test_data, test_labels)

• loss (Union[str, Callable]) – loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems). One can also pass a custom loss function that takes prediction and ground truth and computes a loss score.

• 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 batch. Otherwise, each cycle corresponds to all mini-batches of data passing through a NN.

• batch_size (int) – Size of training and test batches

• compute_accuracy (bool) – Computes accuracy function at each training cycle

• full_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 training

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

• **batch_seed (int) – Random state for generating sequence of training and test batches

• **overwrite_train_data (bool) – Overwrites the exising training data using self.set_data() (Default: True)

• **memory_alloc (float) – threshold (in GB) for holding all training data on GPU

• **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 epochs at the end of training

select_lr(e)

Return type:

None

run()

Trains a neural network, prints the statistics, saves the final model weights. One can also pass kwargs for utils.datatransform class to perform the data augmentation “on-the-fly”

Return type:

Type[Module]

fit()

Return type:

None

Trainer for Semantic Segmentation

class atomai.trainers.SegTrainer(model='Unet', nb_classes=1, **kwargs)

Bases: BaseTrainer

Class for training a fully convolutional neural network for semantic segmentation of noisy experimental data

Parameters:

• model (Union[Type[Module], str]) – Type of model to train: ‘Unet’, ‘Uplusnet’ or ‘dilnet’ (Default: ‘Unet’). See atomai.nets for more details. One can also pass a custom fully convolutional neural network model.

• nb_classes (int) – Number of classes in the classification scheme adopted (must match the number of classes in training data)

• **seed (int) – Deterministic mode for model training (Default: 1)

• **batch_seed (int) – Separate seed for generating a sequence of batches for training/testing. Equal to ‘seed’ if set to None (default)

• **batch_norm (bool) – Apply batch normalization after each convolutional layer (Default: True)

• **dropout (bool) – Apply dropouts in the three inner blocks in the middle of a network (Default: False)

• **upsampling (str) – “bilinear” or “nearest” upsampling method (Default: “bilinear”)

• **nb_filters (int) – Number of convolutional filters 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 U-Net 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)

set_data(X_train, y_train, X_test=None, y_test=None, **kwargs)

Sets training and test data.

Parameters:

• X_train (Tuple[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 automatically

• y_train (Tuple[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 (Optional[Tuple[ndarray, Tensor]]) – 4D numpy array or pytorch tensor of test images (stacked along the first dimension)

• y_test (Optional[Tuple[ndarray, Tensor]]) – 4D (binary) / 3D (multiclass) numpy array or pytorch tensor of training masks (aka ground truth) stacked along the first dimension.

• kwargs (Union[float, int]) – Parameters for train_test_split (‘test_size’ and ‘seed’) when separate test set is not provided and ‘memory_alloc’, which sets a threshold (in GBs) for holding entire training data on GPU

Return type:

None

accuracy_fn(y, y_prob, *args)

Computes accuracy score

compile_trainer(train_data=None, loss='ce', optimizer=None, training_cycles=1000, batch_size=32, compute_accuracy=False, full_epoch=False, swa=False, perturb_weights=False, **kwargs)

Compile a trainer

Parameters:

• train_data (Union[Tuple[Tensor], Tuple[ndarray], None]) – 4-element tuple of ndarrays or torch tensors (train_data, train_labels, test_data, test_labels)

• loss (Union[str, Callable]) – loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems). One can also pass a custom loss function that takes prediction and ground truth and computes a loss score.

• 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 batch. Otherwise, each cycle corresponds to all mini-batches of data passing through a NN.

• batch_size (int) – Size of training and test batches

• compute_accuracy (bool) – Computes accuracy function at each training cycle

• full_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 training

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

• **batch_seed (int) – Random state for generating sequence of training and test batches

• **overwrite_train_data (bool) – Overwrites the exising training data using self.set_data() (Default: True)

• **memory_alloc (float) – threshold (in GB) for holding all training data on GPU

• **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 epochs at the end of training

data_augmentation(augment_fn)

Set up data augmentation. To use it, pass a function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

Return type:

None

dataloader(batch_num, mode='train')

Generates input training data with images/spectra and the associated labels (spectra/images)

Return type:

Tuple[Tensor]

eval_model()

Evaluates model on the entire dataset

Return type:

None

fit()

Return type:

None

get_loss_fn(loss='mse', nb_classes=None)

Returns a loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems)

Return type:

None

print_statistics(e, **kwargs)

Print loss and (optionally) accuracy score on train and test data, as well as GPU memory usage.

Return type:

None

run()

Trains a neural network, prints the statistics, saves the final model weights. One can also pass kwargs for utils.datatransform class to perform the data augmentation “on-the-fly”

Return type:

Type[Module]

save_model(*args)

Saves trained weights, optimizer and key information about model’s architecture (the latter works only for built-in AtomAI models)

Return type:

None

save_running_weights(e)

Saves running weights (for stochastic weights averaging)

Return type:

None

select_lr(e)

Return type:

None

set_model(model, nb_classes=None)

Sets a neural network model and a number of classes (if any)

Parameters:

• model (Type[Module]) – initialized PyTorch model

• nb_classes (Optional[int]) – number of classes in classification scheme (if any)

Return type:

None

step(e)

Single train-test step which passes a single mini-batch (for both training and testing), i.e. 1 “epoch” = 1 mini-batch

Return type:

None

step_full()

A standard PyTorch training loop where all available mini-batches are passed at a single step/epoch

Return type:

None

test_step(feat, tar)

Forward pass for test data with deactivated autograd engine

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

float

train_step(feat, tar)

Propagates image(s) through a network to get model’s prediction and compares predicted value with ground truth; then performs backpropagation to compute gradients and optimizes weights.

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

Tuple[float]

weight_perturbation(e)

Time-dependent weights perturbation (role of time is played by “epoch” number)

Return type:

None

ImSpec Trainer

class atomai.trainers.ImSpecTrainer(in_dim, out_dim, latent_dim=2, **kwargs)

Bases: BaseTrainer

Trainer of neural network for image-to-spectrum and spectrum-to-image transformations

Parameters:

• in_dim (Tuple[int]) – Input data dimensions. (height, width) for images or (length,) for spectra

• out_dim (Tuple[int]) – output dimensions. (length,) for spectra or (height, width) for images

• latent_dim (int) – dimensionality of the latent space (number of neurons in a fully connected bottleneck layer)

• **seed (int) – Deterministic mode for model training (Default: 1)

• **batch_seed (int) – Separate seed for generating a sequence of batches for training/testing. Equal to ‘seed’ if set to None (default)

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

set_data(X_train, y_train, X_test=None, y_test=None, **kwargs)

Sets training and test data.

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_train

• y_test (Union[ndarray, Tensor, None]) – Test data (targets) of the same dimesnionality (except for the number of samples) as y_train

• kwargs (Union[float, int]) – Parameters for train_test_split (‘test_size’ and ‘seed’) when separate test set is not provided and ‘memory_alloc’, which sets a threshold (in GBs) for holding entire training data on GPU

Return type:

None

accuracy_fn(*args)

Computes accuracy score

Return type:

None

compile_trainer(train_data=None, loss='ce', optimizer=None, training_cycles=1000, batch_size=32, compute_accuracy=False, full_epoch=False, swa=False, perturb_weights=False, **kwargs)

Compile a trainer

Parameters:

• train_data (Union[Tuple[Tensor], Tuple[ndarray], None]) – 4-element tuple of ndarrays or torch tensors (train_data, train_labels, test_data, test_labels)

• loss (Union[str, Callable]) – loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems). One can also pass a custom loss function that takes prediction and ground truth and computes a loss score.

• 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 batch. Otherwise, each cycle corresponds to all mini-batches of data passing through a NN.

• batch_size (int) – Size of training and test batches

• compute_accuracy (bool) – Computes accuracy function at each training cycle

• full_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 training

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

• **batch_seed (int) – Random state for generating sequence of training and test batches

• **overwrite_train_data (bool) – Overwrites the exising training data using self.set_data() (Default: True)

• **memory_alloc (float) – threshold (in GB) for holding all training data on GPU

• **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 epochs at the end of training

data_augmentation(augment_fn)

Set up data augmentation. To use it, pass a function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

Return type:

None

dataloader(batch_num, mode='train')

Generates input training data with images/spectra and the associated labels (spectra/images)

Return type:

Tuple[Tensor]

eval_model()

Evaluates model on the entire dataset

Return type:

None

fit()

Return type:

None

get_loss_fn(loss='mse', nb_classes=None)

Returns a loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems)

Return type:

None

print_statistics(e, **kwargs)

Print loss and (optionally) accuracy score on train and test data, as well as GPU memory usage.

Return type:

None

run()

Trains a neural network, prints the statistics, saves the final model weights. One can also pass kwargs for utils.datatransform class to perform the data augmentation “on-the-fly”

Return type:

Type[Module]

save_model(*args)

Saves trained weights, optimizer and key information about model’s architecture (the latter works only for built-in AtomAI models)

Return type:

None

save_running_weights(e)

Saves running weights (for stochastic weights averaging)

Return type:

None

select_lr(e)

Return type:

None

set_model(model, nb_classes=None)

Sets a neural network model and a number of classes (if any)

Parameters:

• model (Type[Module]) – initialized PyTorch model

• nb_classes (Optional[int]) – number of classes in classification scheme (if any)

Return type:

None

step(e)

Single train-test step which passes a single mini-batch (for both training and testing), i.e. 1 “epoch” = 1 mini-batch

Return type:

None

step_full()

A standard PyTorch training loop where all available mini-batches are passed at a single step/epoch

Return type:

None

test_step(feat, tar)

Forward pass for test data with deactivated autograd engine

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

float

train_step(feat, tar)

Propagates image(s) through a network to get model’s prediction and compares predicted value with ground truth; then performs backpropagation to compute gradients and optimizes weights.

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

Tuple[float]

weight_perturbation(e)

Time-dependent weights perturbation (role of time is played by “epoch” number)

Return type:

None

Variational Inference Trainer

class atomai.trainers.viBaseTrainer

Bases: object

Initializes base trainer for VAE and VED models

set_model(encoder_net, decoder_net)

Sets encoder and decoder models

Return type:

None

set_encoder(encoder_net)

Sets an encoder network only

Return type:

None

set_decoder(decoder_net)

Sets a decoder network only

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

elbo_fn()

Computes ELBO

forward_compute_elbo()

Computes ELBO in “train” and “eval” modes. Specifically, it passes input data x through encoder, “compresses” it to latent variables z_mean and z_sd/z_logsd, performs reparametrization trick, passes the reparameterized latent vector through decoder to obtain y/x_reconstructed, and then computes the “loss” via self.elbo_fn, which usually takes as parameters x, y/x_reconstructed, z_mean, and z_sd/z_logsd.

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 labels

• train_data – Test data and (optionally) corresponding targets or labels

• optimizer (Optional[Type[Optimizer]]) – Weights optimizer. Defaults to Adam with learning rate 1e-4

• elbo_fn (Optional[Callable]) – function that calculates elbo loss

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

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

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 distribution

• p_param (Optional[Tuple[Tensor]]) – tuple with mean and SD of the 2nd distribution (optional)

Return type:

Tensor

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

train_epoch()

Trains a single epoch

evaluate_model()

Evaluates model on test data

print_statistics(e)

Prints training and (optionally) test loss after each training cycle

save_model(*args)

Saves trained weights and the key model parameters

Return type:

None

save_weights(*args)

Saves trained weights

Return type:

None

load_weights(filepath)

Loads saved weights

Return type:

None

Deep Ensemble Trainers

Ensemble Trainer

class atomai.trainers.EnsembleTrainer(model=None, nb_classes=1, **kwargs)

Bases: BaseEnsembleTrainer

Deep ensemble trainer

Parameters:

• model (Union[str, Type[Module], None]) – Built-in AtomAI model (passed as string) or initialized custom PyTorch model

• nb_classes (int) – Number of classes (if any) in the model’s output

• **kwargs – Number of input, output, and latent dimensions for imspec models (in_dim, out_dim, latent_dim)

Example

>>> # Train an ensemble of Unet-s
>>> etrainer = aoi.trainers.EnsembleTrainer(
>>>    "Unet", batch_norm=True, nb_classes=3, with_dilation=False)
>>> etrainer.compile_ensemble_trainer(training_cycles=500)
>>> # Train 10 different models from scratch
>>> smodel, ensemble = etrainer.train_ensemble_from_scratch(
>>>    images, labels, images_test, labels_test, n_models=10)

accuracy_fn(*args)

Computes accuracy score

Return type:

None

compile_ensemble_trainer(**kwargs)

Compile ensemble trainer.

Parameters:

kwargs (Union[Type[Optimizer], str, int]) – Keyword arguments to be passed to BaseTrainer.compile_trainer (loss, optimizer, compute_accuracy, full_epoch, swa, perturb_weights, batch_size, training_cycles, accuracy_metrics, filename, print_loss, plot_training_history)

Return type:

None

train_baseline(X_train, y_train, X_test=None, y_test=None, seed=1, augment_fn=None)

Trains baseline weights

Parameters:

• X_train (ndarray) – Training features

• y_train (ndarray) – Training targets

• X_test (Optional[ndarray]) – Test features

• y_test (Optional[ndarray]) – Test targets

• seed (int) – seed to be used for pytorch and numpy random numbers generator

• augment_fn (Optional[Callable[[Tensor, Tensor, int], Tuple[Tensor, Tensor]]]) – function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

Return type:

Type[Module]

Returns:

Trained baseline weights

preprocess_train_data(*args)

Training and test data preprocessing

Return type:

Tuple[Tensor]

compile_trainer(train_data=None, loss='ce', optimizer=None, training_cycles=1000, batch_size=32, compute_accuracy=False, full_epoch=False, swa=False, perturb_weights=False, **kwargs)

Compile a trainer

Parameters:

• train_data (Union[Tuple[Tensor], Tuple[ndarray], None]) – 4-element tuple of ndarrays or torch tensors (train_data, train_labels, test_data, test_labels)

• loss (Union[str, Callable]) – loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems). One can also pass a custom loss function that takes prediction and ground truth and computes a loss score.

• 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 batch. Otherwise, each cycle corresponds to all mini-batches of data passing through a NN.

• batch_size (int) – Size of training and test batches

• compute_accuracy (bool) – Computes accuracy function at each training cycle

• full_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 training

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

• **batch_seed (int) – Random state for generating sequence of training and test batches

• **overwrite_train_data (bool) – Overwrites the exising training data using self.set_data() (Default: True)

• **memory_alloc (float) – threshold (in GB) for holding all training data on GPU

• **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 epochs at the end of training

data_augmentation(augment_fn)

Set up data augmentation. To use it, pass a function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

Return type:

None

dataloader(batch_num, mode='train')

Generates input training data with images/spectra and the associated labels (spectra/images)

Return type:

Tuple[Tensor]

eval_model()

Evaluates model on the entire dataset

Return type:

None

fit()

Return type:

None

get_loss_fn(loss='mse', nb_classes=None)

Returns a loss function. Available loss functions are: ‘mse’ (MSE), ‘ce’ (cross-entropy), ‘focal’ (focal loss; single class only), and ‘dice’ (dice loss; for semantic segmentation problems)

Return type:

None

print_statistics(e, **kwargs)

Print loss and (optionally) accuracy score on train and test data, as well as GPU memory usage.

Return type:

None

run()

Trains a neural network, prints the statistics, saves the final model weights. One can also pass kwargs for utils.datatransform class to perform the data augmentation “on-the-fly”

Return type:

Type[Module]

save_ensemble_metadict(filename=None)

Saves meta dictionary with ensemble weights and key information about model’s structure (needed to load it back) to disk

Return type:

None

save_model(*args)

Saves trained weights, optimizer and key information about model’s architecture (the latter works only for built-in AtomAI models)

Return type:

None

save_running_weights(e)

Saves running weights (for stochastic weights averaging)

Return type:

None

select_lr(e)

Return type:

None

set_data(X_train, y_train, X_test, y_test, **kwargs)

Sets training and test data by initializing PyTorch dataloaders or creating a list of PyTorch tensors from which it will randomly choose an element at each training iteration.

Parameters:

• X_train (Union[Tensor, ndarray]) – Training data

• y_train (Union[Tensor, ndarray]) – Training data labels/ground-truth

• X_test (Union[Tensor, ndarray]) – Test data

• y_test (Union[Tensor, ndarray]) – Test data labels/ground-truth

• memory_alloc – threshold (in GB) for holding all training data on GPU

Return type:

None

set_model(model, nb_classes=None)

Sets a neural network model and a number of classes (if any)

Parameters:

• model (Type[Module]) – initialized PyTorch model

• nb_classes (Optional[int]) – number of classes in classification scheme (if any)

Return type:

None

step(e)

Single train-test step which passes a single mini-batch (for both training and testing), i.e. 1 “epoch” = 1 mini-batch

Return type:

None

step_full()

A standard PyTorch training loop where all available mini-batches are passed at a single step/epoch

Return type:

None

test_step(feat, tar)

Forward pass for test data with deactivated autograd engine

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

float

train_ensemble_from_baseline(X_train, y_train, X_test=None, y_test=None, basemodel=None, n_models=10, training_cycles_base=1000, training_cycles_ensemble=100, augment_fn=None, **kwargs)

Trains ensemble of models starting each time from baseline model. Each ensemble model is trained each with different random shuffling of batches (and different seed for data augmentation if any). If a baseline model is not provided, the baseline weights are trained for N epochs and then used as a baseline to train multiple ensemble models for n epochs (n << N),

Parameters:

• X_train (ndarray) – Training features

• y_train (ndarray) – Training targets

• X_test (Optional[ndarray]) – Test features

• y_test (Optional[ndarray]) – Test targets

• basemodel (Optional[Type[Module]]) – Provide a baseline model (Optional)

• n_models (int) – number of models in ensemble

• training_cycles_base (int) – Number of training iterations for baseline model

• training_cycles_ensemble (int) – Number of training iterations for every ensemble model

• augment_fn (Optional[Callable[[Tensor, Tensor, int], Tuple[Tensor, Tensor]]]) – function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

• **kwargs – Updates kwargs from initial compilation (can be useful for iterative training)

Return type:

Tuple[Type[Module], Dict[int, Dict[str, Tensor]]]

Returns:

Model with averaged weights and dictionary with ensemble weights

train_ensemble_from_scratch(X_train, y_train, X_test=None, y_test=None, n_models=10, augment_fn=None, **kwargs)

Trains ensemble of models starting every time from scratch with different initialization

Parameters:

• X_train (ndarray) – Training features

• y_train (ndarray) – Training targets

• X_test (Optional[ndarray]) – Test features

• y_test (Optional[ndarray]) – Test targets

• n_models (int) – number of models to be trained

• augment_fn (Optional[Callable[[Tensor, Tensor, int], Tuple[Tensor, Tensor]]]) – function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

• **kwargs – Updates kwargs from initial compilation, which can be useful for iterative training.

Return type:

Tuple[Type[Module], Dict[int, Dict[str, Tensor]]]

Returns:

The last trained model and dictionary with ensemble weights

train_step(feat, tar)

Propagates image(s) through a network to get model’s prediction and compares predicted value with ground truth; then performs backpropagation to compute gradients and optimizes weights.

Parameters:

• feat (Tensor) – input features

• tar (Tensor) – targets

Return type:

Tuple[float]

train_swag(X_train, y_train, X_test=None, y_test=None, n_models=10, augment_fn=None, **kwargs)

Performs SWAG-like weights sampling at the end of single model training

Parameters:

• X_train (ndarray) – Training features

• y_train (ndarray) – Training targets

• X_test (Optional[ndarray]) – Test features

• y_test (Optional[ndarray]) – Test targets

• n_models (int) – number fo samples to be drawn

• augment_fn (Optional[Callable[[Tensor, Tensor, int], Tuple[Tensor, Tensor]]]) – function that takes two torch tensors (features and targets), peforms some transforms, and returns the transformed tensors. The dimensions of the transformed tensors must be the same as the dimensions of the original ones.

• **kwargs – Updates kwargs from initial compilation (can be useful for iterative training)

Return type:

Tuple[Type[Module], Dict[int, Dict[str, Tensor]]]

Returns:

Baseline model and dictionary with sampled weights

update_training_parameters(kwargs)

weight_perturbation(e)

Time-dependent weights perturbation (role of time is played by “epoch” number)

Return type:

None

Single Model Predictors

BasePredictor

class atomai.predictors.BasePredictor(model=None, use_gpu=False, **kwargs)

Base predictor class

preprocess(data)

Preprocess input data

Return type:

None

forward_(xnew)

Pass data through a trained neural network

Return type:

Tensor

batch_predict(data, out_shape, num_batches)

Make a prediction batch-by-batch (for larger datasets)

Return type:

Tensor

predict(data, out_shape=None, num_batches=1)

Make a prediction on the new data with a trained model

Return type:

Tensor

SegPredictor

class atomai.predictors.SegPredictor(trained_model, refine=False, resize=None, use_gpu=False, logits=True, **kwargs)

Bases: BasePredictor

Prediction with a trained fully convolutional neural network

Parameters:

• trained_model (Type[Module]) – Trained pytorch model (skeleton+weights)

• refine (bool) – Atomic positions refinement with 2d Gaussian peak fitting

• resize (Union[Tuple, List, None]) – Target dimensions for optional image(s) resizing

• use_gpu (bool) – Use gpu device for inference

• logits (bool) – Indicates that the image data is passed through a softmax/sigmoid layer when set to False (logits=True for AtomAI models)

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

• **nb_classes (int) – Number of classes in the model

• **downsampling (int or float) – Downsampling factor (equal to \(2^n\) where n is a number of pooling operations)

Example

>>> # Here we load new experimental data (as 2D or 3D numpy array)
>>> expdata = np.load('expdata-test.npy')
>>> # Get prediction from a trained model
>>> pseg = atomnet.SegPredictor(trained_model)
>>> nn_output, coords = pseg.run(expdata)

preprocess(image_data, norm=True)

Prepares an input for a neural network

Return type:

Tensor

forward_(images)

Returns ‘probability’ of each pixel in image(s) belonging to an atom/defect

Return type:

ndarray

predict(image_data, return_image=False, **kwargs)

Make prediction

Parameters:

• image_data (ndarray) – 3D image stack or a single 2D image (all greyscale)

• return_image (bool) – Returns images used as input into NN

• **num_batches (int) – number of batches

• **norm (bool) – Normalize data to (0, 1) during pre-processing

Return type:

Tuple[ndarray]

run(image_data, compute_coords=True, **kwargs)

Make prediction with a trained model and calculate coordinates

Parameters:

• image_data (ndarray) – Image stack or a single image (all greyscale)

• compute_coords – Computes centers of the mass of individual blobs in the segmented images (Default: True)

• **num_batches (int) – number of batches for batch-by-batch prediction which ensures that one doesn’t run out of memory (Default: 10)

• **norm (bool) – Normalize data to (0, 1) during pre-processing

Return type:

Tuple[ndarray, Dict[int, ndarray]]

batch_predict(data, out_shape, num_batches)

Make a prediction batch-by-batch (for larger datasets)

Return type:

Tensor

ImSpecPredictor

class atomai.predictors.ImSpecPredictor(trained_model, output_dim, use_gpu=False, **kwargs)

Bases: BasePredictor

Prediction with a trained im2spec or spec2im model

Parameters:

• trained_model (Type[Module]) – Pre-trained neural network

• output_dim (Tuple[int]) – Output dimensions. For im2spec, the output_dim is (length,). For spec2im, the output_dim is (height, width)

• use_gpu (bool) – Use GPU accelration for prediction

• verbose – Verbosity

Example

>>> # Predict spectra from images with pretrained im2spec model
>>> out_dim = (16,)  # spectra length
>>> prediction = ImSpecPredictor(trained_model, out_dim).run(data)

preprocess(signal, norm=True)

Preprocess input signal (images or spectra)

Return type:

Tensor

predict(signal, **kwargs)

Predict spectra from images or vice versa

Parameters:

• signal (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

Return type:

ndarray

run(signal, **kwargs)

Make prediction with a trained model

Parameters:

• signal (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

Return type:

ndarray

batch_predict(data, out_shape, num_batches)

Make a prediction batch-by-batch (for larger datasets)

Return type:

Tensor

forward_(xnew)

Pass data through a trained neural network

Return type:

Tensor

Deep Ensemble Predictors

Ensemble Predictor

class atomai.predictors.EnsemblePredictor(skeleton, ensemble, data_type='image', output_type='image', nb_classes=None, in_dim=None, out_dim=None, **kwargs)

Bases: BasePredictor

Prediction with ensemble of models

Parameters:

• skeleton (Type[Module]) – Model skeleton (cam be with randomly initialized weights)

• ensemble (Dict[int, Dict[str, Tensor]]) – Ensemble of trained weights

• data_type (str) – Input data type (image or spectra)

• output_type (str) – Output data type (image or spectra)

• nb_classes (Optional[int]) – Number of classes (e.g. for semantic segmentation)

• in_dim (Optional[Tuple[int]]) – Input data size (for models with fully-connected layers)

• out_dim (Optional[Tuple[int]]) – Output data size (for models with fully-connected layers)

• **output_shape – Optionally one may specify the exact output shape

• **verbose – verbosity

Example

>>> p = aoi.predictors.EnsemblePredictor(skeleton, ensemble, nb_classes=3)
>>> nn_out_mean, nn_out_var = p.predict(expdata)

preprocess(data, norm=True)

Preprocesses data depending on type (image or spectra)

Return type:

Tensor

ensemble_forward_(data, out_shape)

Computes mean and variance of prediction with ensemble models

Return type:

Tuple[ndarray]

ensemble_forward(data, out_shape, num_batches=1)

Computes prediction with ensemble models. Returns ALL calculated predictions (n_models * n_samples).

Return type:

ndarray

ensemble_batch_predict(data, num_batches=10)

Batch-by-batch prediction with ensemble models

Return type:

Tuple[ndarray]

predict(data, num_batches=10, format_out='channel_last', norm=True)

Predicts mean and variance for all the data points with ensemble of models

Parameters:

• data (ndarray) – input data

• num_batches (int) – number of batches for batch-by-batch prediction (Default: 10)

• format_out (str) – ‘channel_last’ of ‘channel_first’ dimension order in output

• norm (bool) – Normalize input data to (0, 1)

Return type:

Tuple[ndarray]

Returns:

Tuple of numpy arrays with predicted mean and variance

atomai.predictors.ensemble_locate(nn_output_ensemble, **kwargs)

Finds coordinates for each ensemble predictions

Parameters:

• nn_output_ensembles (numpy array) – 5D numpy array with ensemble predictions

• **eps (float) – DBSCAN epsilon for clustering coordinates

• **threshold (float) – threshold value for atomnet.locator

Return type:

Tuple[ndarray, ndarray]

Returns:

Mean and variance for every detected atom/defect/particle coordinate