Trainers and Predictors
Single Model Trainers
Base Trainer Class
- class atomai.trainers.BaseTrainer[source]
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)[source]
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 datay_train (
Union
[Tensor
,ndarray
]) – Training data labels/ground-truthX_test (
Union
[Tensor
,ndarray
]) – Test datay_test (
Union
[Tensor
,ndarray
]) – Test data labels/ground-truthmemory_alloc – threshold (in GB) for holding all training data on GPU
- Return type
None
- set_model(model, nb_classes=None)[source]
Sets a neural network model and a number of classes (if any)
- Parameters
model (
Type
[Module
]) – initialized PyTorch modelnb_classes (
Optional
[int
]) – number of classes in classification scheme (if any)
- Return type
None
- get_loss_fn(loss='mse', nb_classes=None)[source]
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)[source]
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 featurestar (
Tensor
) – targets
- Return type
Tuple
[float
]
- test_step(feat, tar)[source]
Forward pass for test data with deactivated autograd engine
- Parameters
feat (
Tensor
) – input featurestar (
Tensor
) – targets
- Return type
float
- step(e)[source]
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()[source]
A standard PyTorch training loop where all available mini-batches are passed at a single step/epoch
- Return type
None
- dataloader(batch_num, mode='train')[source]
Generates input training data with images/spectra and the associated labels (spectra/images)
- Return type
Tuple
[Tensor
]
- save_model(*args)[source]
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)[source]
Print loss and (optionally) IoU score on train and test data, as well as GPU memory usage.
- Return type
None
- weight_perturbation(e)[source]
Time-dependent weights perturbation (role of time is played by “epoch” number)
- Return type
None
- save_running_weights(e)[source]
Saves running weights (for stochastic weights averaging)
- Return type
None
- data_augmentation(augment_fn)[source]
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)[source]
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 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.
**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
Trainer for Semantic Segmentation
- class atomai.trainers.SegTrainer(model='Unet', nb_classes=1, **kwargs)[source]
Bases:
atomai.trainers.trainer.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)[source]
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 automaticallyy_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
- 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 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.
**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) IoU 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 modelnb_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 featurestar (
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 featurestar (
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)[source]
Bases:
atomai.trainers.trainer.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 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) – 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)[source]
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_trainy_test (
Union
[ndarray
,Tensor
,None
]) – Test data (targets) of the same dimesnionality (except for the number of samples) as y_trainkwargs (
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 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.
**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) IoU 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 modelnb_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 featurestar (
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 featurestar (
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[source]
Bases:
object
Initializes base trainer for VAE and VED models
- set_data(X_train, y_train=None, X_test=None, y_test=None, memory_alloc=4)[source]
Initializes train and (optionally) test data loaders
- Return type
None
- forward_compute_elbo()[source]
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)[source]
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
- classmethod reparameterize(z_mean, z_sd)[source]
Reparameterization trick for continuous distributions
- Return type
Tensor
- classmethod reparameterize_discrete(alpha, tau)[source]
Reparameterization trick for discrete gumbel-softmax distributions
- kld_normal(z, q_param, p_param=None)[source]
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
- classmethod log_normal(x, mu, log_sd)[source]
Computes log-pdf for a normal distribution
- Return type
Tensor
Deep Ensemble Trainers
Ensemble Trainer
- class atomai.trainers.EnsembleTrainer(model=None, nb_classes=1, **kwargs)[source]
Bases:
atomai.trainers.etrainer.BaseEnsembleTrainer
Deep ensemble trainer
- Parameters
model (
Union
[str
,Type
[Module
],None
]) – Built-in AtomAI model (passed as string) or initialized custom PyTorch modelnb_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)[source]
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)[source]
Trains baseline weights
- Parameters
X_train (
ndarray
) – Training featuresy_train (
ndarray
) – Training targetsX_test (
Optional
[ndarray
]) – Test featuresy_test (
Optional
[ndarray
]) – Test targetsseed (
int
) – seed to be used for pytorch and numpy random numbers generatoraugment_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)[source]
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 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.
**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) IoU 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 datay_train (
Union
[Tensor
,ndarray
]) – Training data labels/ground-truthX_test (
Union
[Tensor
,ndarray
]) – Test datay_test (
Union
[Tensor
,ndarray
]) – Test data labels/ground-truthmemory_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 modelnb_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 featurestar (
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 featuresy_train (
ndarray
) – Training targetsX_test (
Optional
[ndarray
]) – Test featuresy_test (
Optional
[ndarray
]) – Test targetsbasemodel (
Optional
[Type
[Module
]]) – Provide a baseline model (Optional)n_models (
int
) – number of models in ensembletraining_cycles_base (
int
) – Number of training iterations for baseline modeltraining_cycles_ensemble (
int
) – Number of training iterations for every ensemble modelaugment_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 featuresy_train (
ndarray
) – Training targetsX_test (
Optional
[ndarray
]) – Test featuresy_test (
Optional
[ndarray
]) – Test targetsn_models (
int
) – number of models to be trainedaugment_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 featurestar (
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 featuresy_train (
ndarray
) – Training targetsX_test (
Optional
[ndarray
]) – Test featuresy_test (
Optional
[ndarray
]) – Test targetsn_models (
int
) – number fo samples to be drawnaugment_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
SegPredictor
- class atomai.predictors.SegPredictor(trained_model, refine=False, resize=None, use_gpu=False, logits=True, **kwargs)[source]
Bases:
atomai.predictors.predictor.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 fittingresize (
Union
[Tuple
,List
,None
]) – Target dimensions for optional image(s) resizinguse_gpu (
bool
) – Use gpu device for inferencelogits (
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)[source]
Prepares an input for a neural network
- Return type
Tensor
- forward_(images)[source]
Returns ‘probability’ of each pixel in image(s) belonging to an atom/defect
- Return type
ndarray
- predict(image_data, return_image=False, **kwargs)[source]
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)[source]
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)[source]
Bases:
atomai.predictors.predictor.BasePredictor
Prediction with a trained im2spec or spec2im model
- Parameters
trained_model (
Type
[Module
]) – Pre-trained neural networkoutput_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 predictionverbose – 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)[source]
Preprocess input signal (images or spectra)
- Return type
Tensor
- predict(signal, **kwargs)[source]
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)[source]
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)[source]
Bases:
atomai.predictors.predictor.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 weightsdata_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)[source]
Preprocesses data depending on type (image or spectra)
- Return type
Tensor
- ensemble_forward_(data, out_shape)[source]
Computes mean and variance of prediction with ensemble models
- Return type
Tuple
[ndarray
]
- ensemble_forward(data, out_shape, num_batches=1)[source]
Computes prediction with ensemble models. Returns ALL calculated predictions (n_models * n_samples).
- Return type
ndarray
- ensemble_batch_predict(data, num_batches=10)[source]
Batch-by-batch prediction with ensemble models
- Return type
Tuple
[ndarray
]
- predict(data, num_batches=10, format_out='channel_last', norm=True)[source]
Predicts mean and variance for all the data points with ensemble of models
- Parameters
data (
ndarray
) – input datanum_batches (
int
) – number of batches for batch-by-batch prediction (Default: 10)format_out (
str
) – ‘channel_last’ of ‘channel_first’ dimension order in outputnorm (
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)[source]
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