To compress a model using a Python script it is possible to import aiminify.minify
and call the main(...)
function. As shown in the example below:
from aiminify.minify import main
from your_project import model
compressed_model, _ = main(model)
Depending on the type of model you are compressing it can be stored using the appropriate methods. As aiminify
supports PyTorch and Tensorflow as backends, the different storing methods are implemented in the save_model()
function.
from aiminify.minify import main as minify
from aiminify import save_model
from your_project import model
compressed_model, _ = minify(model)
save_model(compressed_model, "./compressed-pytorch-model.pt", input_shape)
compressed_model, _ = main(model, quantization=false)
save_model(compressed_model, "./compressed-pytorch-model.pt")
from aiminify.minify import main as minify
from aiminify import save_model
from your_project import model
compressed_model, _ = minify(model)
save_model(compressed_model, "./compressed-tensorflow-model.keras")
aiminify.minify.main
arguments
Compression strength to use for defining the compression strategy
Training set used for fine tuning
Validation set used for fine tuning
Loss function used for fine tuning
Fine tune the model after compression
Use a strategy for determining the amount of filters to prune instead of flat pruning x % of all filters.
The return value of main
is (compressed_model, feedback_dictionary)
. compressed_model
is, as the name would suggest, your compressed model using the same backend as the input (Tensorflow, PyTorch). feedback_dictionary
contains logs and miscellaneous messages from the compression algorithm.
training_generator
, validation_generator
and loss_function
can be specific to the backend you’re using. For example when using pytorch the training_generator
and validation_generator
need to be a subclass of torch.utils.data.Dataset
. For Tensorflow these can be a subclass of tensorflow.data.Dataset
. Similar with loss_function
for PyTorch this can be any member of torch.nn.modules.loss
and for Tensorflow this can be any implementation of tf.keras.losses.Loss
.