atomai

What's Changed

• Add toggle on-off output display by @arpanbiswas52 in https://github.com/pycroscopy/atomai/pull/62
• Update examples by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/63
• Update training data url by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/65
• Add regressor models by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/72
• Update installation instructions and URLs in the notebooks by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/73
• Add 'Classifier' class to atomai models by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/74
• Add a multi-task classifier by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/75
• minor fix by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/76
• Sparse image reconstructor based on the structured kernel interpolation framework by @ziatdinovmax in https://github.com/pycroscopy/atomai/pull/77

New Contributors

• @arpanbiswas52 made their first contribution in https://github.com/pycroscopy/atomai/pull/62

Full Changelog: https://github.com/pycroscopy/atomai/compare/v0.7.4...v0.7.8

Minor fixes and improvements

• Extend Thompson sampler to 'multi-task' DKL-GP models
• Fix bugs in stat.multivar.update_classes() that didn't allow specifying custom window size

• Ensemble deep kernel learning (DKL) as an 'approximation' to the fully Bayesian DKL
• Thompson sampler for active learning now comes as a built-in method in the DKL class
• Option to select between correlated and independent outputs for vector-valued function in DKL

Example of using an ensemble of DKL models:

# Initialize and train ensemble of models
dklgp = aoi.models.dklGPR(indim=X_train.shape[-1], embedim=2)
dklgp.fit_ensemble(X_train, y_train, n_models=5, training_cycles=1500, lr=0.01)
# Make a prediction
y_samples = dklgp.sample_from_posterior(X_test, num_samples=1000) # n_models x n_samples x n_data
y_pred = y_samples.mean(axis=(0,1)) # average over model and sample dimensions
y_var = y_samples.var(axis=(0,1))

Example of using a built-in Thompson sampler for active learning:

for e in range(exploration_steps):
    # obtain/update DKL-GP posterior
    dklgp = aoi.models.dklGPR(data_dim, embedim=2, precision="single")
    dklgp.fit(X_train, y_train, training_cycles=50)
    # Thompson sampling for selecting the next measurement/evaluation point
    obj, next_point = dklgp.thompson(X_cand)
    # Perform a 'measurement'
    y_measured = measure(next_point)
    # Update measured and candidate points, etc...

Bug fixes and additional utility functions for image and spectral data processing

Minor release to patch some numerical stability issues with DKL-GP

New functionalities

1) Deep kernel learning (DKL)-based Gaussian process (GP) regression.

The DKL-GP is based on this paper and can be used for predicting a functional property (or properties) from structural data such as images. Example of usage:

# Structural image data
n, d1, d2 = imgstack.shape
x_train = imgstack.reshape(n, d1*d2)
# Property
y_train = P[:, 0] # can be a scalar or vector variable

# Input data dims
data_dim = x_train.shape[-1]
# Initialize model
dklgp = aoi.models.dklGPR(data_dim)
# Train
dklgp.fit(
    x_train y_train, # inputs and outputs
    training_cycles=100, precision="single", lr=1e-2  # training parameters
)

# Make a prediction (with quantified uncertainty) with the trained model
mean, var = dklgp.predict(x_new)

For more details, see the example notebook

2) Pre-trained models

One can now load pre-trained models for atomic feature finding in graphene and BFO-like systems. Currently limited to STEM data. Example of usage:

# Load model for atom finding in graphene, which was trained on simulated data
model = aoi.models.load_pretrained_model("G_MD")
# Apply to your data
nn_out, coords = model.predict(new_data)

As with any machine learning model, there is a caveat that the performance of pre-trained models will likely degrade significantly on the out-of-distribution data (different feature size, presence of objects on the surface not accounted for in the simulations, etc.)

Bug fixes

• The extractor of image patches now checks for NaNs in the cropped data.

• Add atomai.utils.dataset module with experimental datasets for scientific machine learning
• Minor bug fixes and documentation improvement

New functionalities

• Utility functions for converting Segmentor output (coordinates and classes) to files readable by packages such as Atomic Simulation Environment, VESTA, etc.
• Optional time-dependent learning rate. For example,

# We are going to start with a constant learning rate, then after 600 iterations we begin
# linearly decreasing it over the next 200 iterations, and keep constant afterwards
lr_t = np.ones(800) * 1e-3
lr_t[600:800] = np.linspace(1e-3, 1e-4, 200)

model.fit(images, labels, images_test, labels_test,  # training data
          training_cycles=1000, compute_accuracy=True,  # basic training parameters
          swa=True, lr_scheduler=lr_t  # advanced training parameters
)

Other changes

• Added new examples (Graph analysis and Im2Spec) and expanded explanations in the markdown parts for the old ones
• Improved (slightly) documentation

New functionality

• Added option for controlled information capacity increase to VAE and rVAE (jVAE and jrVAE have them by default). Based on Eq (8) in https://arxiv.org/pdf/1804.03599.pdf

Bug fixes

• Fixed bug that was preventing from loading (older) AtomAI models without saved optimizer in their meta-dictionary

Other

• Fixed some inconsistencies in classes/methods documentation
• Tested against the newest PyTorch release (1.8.0)

New functionalities:

• Add VAE that can learn (simultaneously) both discrete and continuous latent representations
  
• Add option for annealing of the KL terms associated with rotation and image content to rVAE

Bug fixes

• Fix bug that prevented rVAE from working in non-square images
• Fix bug that was causing VAE decoders to "forget" apply sigmoid in evaluation regime after training with BCE with logits

Other improvements

• Add option to set custom encoder and decoder modules in all VAEs
• Add a substantial amount of tests for VI trainer and VAE modules
• Update docs

New functionalities:

• ResHedNet model for advanced edge detection. This model is based on the holistically-nested edge detection paper. We improved the original model by replacing vanilla convolutional layers with ResNet-like blocks in each segment and by reducing the number of max-pooling operations to 2 (we found that 3 different scales are enough for learning the relevant features in typical microscopy images)
• SegResNet model for general semantic segmentation as an alternative to default UNet model. It has ResNet-like connections in each segment in addition to UNet-like skip connections between encoding and decoding paths.

Bug fixes:

• Fix bug that was preventing from saving/loading custom models
• Fix bug that was performing a zoom-in operation even when set to False during data augmentation
• Fix bug in the output_shape in BasePredictor, which required the output shape to be identical to the input shape

Improvements:

• Add option to pass a custom loss function to trainers for semantic segmentation and im2spec/spec2im
• Add option to store all training data on CPU when the size of the training data exceeds a certain limit (default limit is 4GB). In this case, only the individual batches are moved to a GPU device at training/test steps.
• Make computation of coordinates optional for SegPredictor
• Automatically save VAE models after each training cycle ("epoch") and not just at the end of training

New examples:

• New notebook on constructing and using (training+predicting) a custom image denoiser with AtomAI
• New notebook on applications of rotationally invariant VAE (rVAE) and class-conditioned rVAE to arbitrary rotated handwritten digits

Breaking changes

• out: atomnet, atomstat
• in: models, trainers, predictors

To install and use the old code run pip install git+https://github.com/ziatdinovmax/atomai@legacy --upgrade

The new version provides an easy, Keras-like interface for training and applying models for semantic segmentation, image-to-spectrum conversion, as well as different forms of variational autoencoders. For example, to train a model for semantic segmentation of data, (for atom/defect finding) simply run:

model = Segmentor()
model.fit(X, y, X_test, y_test, training_cycles=300)

To make a prediction with a trained model, run:

output, coords = model.predict(expdata)

See the updated documentation for more details.

New functionalities:

• ImSpec models for converting 2D images to 1D spectra and vice versa
• Graph analysis for identifying topologcial defects in the lattices
• Class-conditioned VAE and rVAE

Imrovements:

• AtomAI's trainers and predictors can now work with custom Pytorch models

New functionalities:

• Optional time-dependent weight perturbation w <- w + N(0, scale(t)) during NN model/ensemble training
• Vanilla Pytorch NN training loop (the customized on-the-fly data augmentation is not available in this mode)
• Basic utilities for working with graphs including:
  • Construction of graph from NN output
  • First-depth search for analyzing lattice topology
  • Using connected/disconnected subgraphs to clean the NN predictions
  • Plotting graphs

Bug fixes/improvements:

• fix bug that prevented NN training in a 'full_epoch' mode for multiclass case
• automatically load VAE's weights on cpu when cuda is not found
• return subimages together with VAE-encoded trajectories
• option to pass a custom latent variable range when plotting the VAE's manifold2d