tinyml-compress 1.0.1

A library for converting, optimizing, and saving/loading machine learning models.

Model Conversion and Optimization Library (tinyml_compress)

This library provides functions to convert, save, load, distill, prune, and quantize machine learning models for both TensorFlow and PyTorch. It supports model conversion to TFLite and ONNX formats, as well as optimization techniques like distillation, pruning, and quantization.

Features

• Model Conversion:

  • Convert TensorFlow models to TFLite format.
  • Convert PyTorch models to ONNX format.

• Model Optimization:

  • Distillation:
    • Perform knowledge distillation between a teacher and a student model for both TensorFlow and PyTorch.
  • Pruning:
    • Apply pruning to TensorFlow and PyTorch models to reduce the number of parameters.
  • Quantization:
    • Apply quantization to TensorFlow models to reduce model size and improve inference speed.
    • Apply dynamic quantization to PyTorch models.

• Model Serialization:

  • Save and load TensorFlow and PyTorch models to and from disk.

Installation

To install this library, you can use the following pip command:

pip install tinyML

Install required dependencies:

pip install -r requirements.txt

Usage

Model Conversion

Convert TensorFlow Model to TFLite

import tensorflow as tf from model_conversion import convert_to_tflite

Create a simple model model = tf.keras.Sequential([tf.keras.layers.Dense(10, input_shape=(5,))]) tflite_model = convert_to_tflite(model)

Convert PyTorch Model to ONNX

import torch from model_conversion import convert_to_onnx

model = torch.nn.Linear(5, 10) dummy_input = torch.randn(1, 5) onnx_path = convert_to_onnx(model, dummy_input, "model.onnx")

Model Optimization

Distillation in TensorFlow

from model_conversion import distill_model_tflite

teacher_model = tf.keras.models.load_model("teacher_model.h5") student_model = tf.keras.models.Sequential([tf.keras.layers.Dense(10, input_shape=(5,))]) train_data, train_labels = ... # Your training data student_model = distill_model_tflite(teacher_model, student_model, train_data, train_labels)

Distillation in PyTorch

from model_conversion import distill_model_pytorch

teacher_model = torch.load("teacher_model.pt") student_model = torch.nn.Linear(5, 10) train_loader = ... # Your training data loader optimizer = torch.optim.Adam(student_model.parameters()) student_model = distill_model_pytorch(teacher_model, student_model, train_loader, optimizer)

Pruning in TensorFlow

from model_conversion import prune_model_tflite

model = tf.keras.Sequential([tf.keras.layers.Dense(10, input_shape=(5,))]) pruned_model = prune_model_tflite(model, sparsity=0.5)

Pruning in PyTorch

from model_conversion import prune_model_pytorch model = torch.nn.Linear(5, 10) pruned_model = prune_model_pytorch(model, amount=0.3)

Quantization in TensorFlow

from model_conversion import quantize_model_tflite

model = tf.keras.Sequential([tf.keras.layers.Dense(10, input_shape=(5,))]) quantized_model = quantize_model_tflite(model)

Quantization in PyTorch

from model_conversion import quantize_model_pytorch

model = torch.nn.Linear(5, 10) quantized_model = quantize_model_pytorch(model)

Model Serialization

Save a Model

from model_conversion import save_model

model = tf.keras.Sequential([tf.keras.layers.Dense(10, input_shape=(5,))]) save_model(model, "model.h5")

PyTorch model

import torch model = torch.nn.Linear(5, 10) save_model(model, "model.pt")

Load a Model

from model_conversion import load_model

Load TensorFlow model

tensorflow_model = load_model("model.h5", model_type="tensorflow")

Load PyTorch model

pytorch_model = load_model("model.pt", model_type="pytorch", input_shape=(5,))