aisquared.utils package


aisquared.utils.utils module

aisquared.utils.utils.get_model(model_type: str, input_shape: int | tuple, num_outputs: int, output_activation: str, size: str = 'small', vocab_size: None | int = None)[source]

Get a pre-configured model for different use cases

  • model_type (str) – Either ‘cv’, ‘nlp_embedding’, or ‘fc’, defining the model type

  • input_shape (int or tuple of int) – The input shape to the model

  • num_outputs (int) – The output shape of the model

  • output_activation (str or keras activation function) – The activation of the final layer of the model

  • size (str (default 'small')) – One of either ‘small’, ‘medium’, or ‘large’

  • vocab_size (str or None (default None)) – Size of the vocab, if model_type is ‘nlp_embedding’


model – The model

Return type:

TensorFlow Keras model

aisquared.utils.utils.mimic_model(trained_model: BaseEstimator, nnet: Model, training_data: ndarray, test_data: ndarray, test_labels: ndarray, problem_type: str, loss: str, metrics: str | list, optimizer: str, mimic_proba: bool = False, retention: float = 0.9, batch_size: int = 32, epochs: int = 100, starting_sparsification: int = 0, max_sparsification: int = 99, sparsification_rate: int = 5) Model[source]

Train a sparse neural network to mimic a scikit-learn model

  • trained_model (sklearn model) – The model that is already trained

  • nnet (TensorFlow keras Model) – The neural network to train to mimic the trained model

  • training_data (array or array-like) – The input data that was used to train the trained model

  • test_data (array or array-like) – The input data to be used for testing

  • test_labels (array or array-like) – The output data used in testing

  • problem_type (str) – The type of problem, either ‘classification’ or ‘regression’

  • loss (str or keras loss function) – The loss to use

  • metrics (str, function or list of str, function) – Metrics to measure

  • optimizer (str or keras optimizer) – The optimizer to use

  • mimic_proba (bool (default False)) – For classification, mimic the probability outputs

  • retention (float (default 0.9)) – The retention of performance to allow further pruning

  • batch_size (int (default 32)) – The batch size to use while training

  • epochs (int (default 100)) – The number of epochs (if early stopping is not met beforehand)

  • starting_sparsification (int (default 0)) – The starting model sparsification

  • max_sparsification (int (default 99)) – The maximum sparsification to allow

  • sparsification_rate (int (default 5)) – The sparsification rate when invoked


nnet – The trained model

Return type:

TensorFlow keras Model

Module contents

Additional utilities to use with the aisquared package. These utilities currently consist of two functions, the mimic_model and get_model functions. They utilize functionality that exists in our open source package BeyondML to train teacher-student models

To see in-depth examples of how to use these functions, please visit our GitHub repository at