TinyNeuralNetwork

TinyNeuralNetwork provides various pruning strategies for optimizing neural networks, specifically targeting MobileNetV1 architectures. The framework supports different pruning algorithms and methods to balance model size (channels/FLOPs) and accuracy.

Available Pruning Strategies

Based on the benchmark data, the following pruning components are available:

• OneShotPruner: A pruning method that can be combined with different importance metrics:
  • random: Random pruning.
  • l1_norm: Pruning based on L1 norm.
  • l2_norm: Pruning based on L2 norm.
  • fpgm: Pruning using FPGM.
  • hrank: Pruning using H-rank.
  • Distillation: Some OneShotPruner configurations support distillation (e.g., OneShotPruner, l2_norm, 蒸馏) to recover accuracy.
• BNSlimingPruner: A specialized pruning strategy.
• ADMMPruner: Pruning using the ADMM (Alternating Direction Method of Multipliers) approach, often used with l2_norm.
• GradualPruner: A strategy for gradual pruning, often used with l2_norm.
• RepPruner: A pruning strategy that can be combined with distillation.

Benchmark Results Summary

CIFAR10 Benchmarks

For MobileNetV1-0.5 (reducing to 50% channels and 25% FLOPs), OneShotPruner with l2_norm showed the highest accuracy improvement (+0.37%).

For MobileNetV1-0.25 (reducing to 25% channels and 6% FLOPs), GradualPruner with l2_norm provided the best accuracy retention (93.78%).

ImageNet Benchmarks

TinyNeuralNetwork's pruning methods (like OneShotPruner, l1_norm and GradualPruner, l2_norm) demonstrate competitive accuracy compared to original MobileNetV1 paper results and AMC (AutoML for Model Compression) benchmarks, often achieving lower accuracy loss than the baseline Google implementations when scaling down models.

TinyNeuralNetwork provides several key capabilities for deep learning model compression:

• Graph Tracer (Computation Graph Capture): Captures PyTorch operator connections to enable automatic pruning and quantization. It also supports reverse code generation (codegen) to convert PyTorch models into an equivalent model.py.
• Graph Modifier (Dependency Analysis): Manages dependencies within and between subgraphs. When a single operator is modified, the Graph Modifier automatically handles the resulting changes in related operators.
• Pruner: Implements various automated pruning algorithms including OneShot (L1, L2, FPGM), ADMM, NetAdapt, Gradual, and End2End (to be released incrementally).
• Quantization Training: Uses PyTorch's QAT (Quantization Aware Training) as a backend with extended support for BF16 training. It automates operator fusion and computation graph quantization, which typically requires significant manual effort in standard PyTorch implementations.
• Model Converter: Supports converting both floating-point and quantized PyTorch models into TFLite format for on-device deployment.

TinyNeuralNetwork simplifies the standard PyTorch quantization process by automating several manual steps.

Standard PyTorch Workflow (Manual):

1. Manually insert QuantStub and DeQuantStub into the forward function.
2. Replace arithmetic operations (e.g., add, mul, cat) with torch.nn.quantized.FloatFunctional operators.
3. Manually write a fuse_model function to fuse operators (e.g., Conv2D + BatchNorm2D $\rightarrow$ ConvBN2D).
4. Call prepare_qat to convert the graph to a quantized graph.
5. Deployment is difficult because quantized models are often limited to TorchScript and cannot easily convert to ONNX or TFLite.

TinyNeuralNetwork Workflow (Automated):

• Code Generation: Automatically generates new code equivalent to manual QuantStub insertion and FloatFunctional replacement.
• Intelligent Preparation: The prepare_qat function automatically handles operator fusion and mixed-precision analysis.
• Simplified Deployment: Provides conversion from TorchScript to TFLite (supporting both floating-point and quantized models), making edge deployment easier.

The DLContext object is used to provide the necessary training and validation components to the pruner. Depending on whether you are performing the main NetAdapt training or the final fine-tuning, different parameters are required.

For NetAdapt Training

You must provide:

• train_loader (Dataloader): Training data loader.
• val_loader (Dataloader): Validation data loader.
• criterion (lambda): A loss function with signature (output: Tensor, target: Tensor) -> Tensor.
• optimizer (Optimizer): The optimizer for training.

For Fine-tuning

You must provide:

• train_loader (Dataloader): Training data loader.
• val_loader (Dataloader): Validation data loader.
• criterion (lambda): A loss function with signature (output: Tensor, target: Tensor) -> Tensor.
• optimizer (Optimizer): The optimizer for training.
• scheduler (LR_Scheduler, optional): An optional learning rate scheduler for the optimizer.

Converting RNNs to TFLite can be done in two ways:

1. Unrolled RNN (unroll_rnn=True)

This translates the RNN into a series of standard operations. This is the easiest method but results in a more complex computation graph.

2. Single Operator (UnidirectionalLSTM)

To use a single UnidirectionalLSTM operator, you must manually manage the state tensors.

Important: When exporting from PyTorch, you must delete the LSTM state inputs and outputs from the model's inputs and outputs to simulate TensorFlow's behavior (where states are persistent Variables).

Usage Scenarios:

• Non-streaming: Call interpreter.reset_all_variables() in the TFLite Interpreter before each invoke() to reset states to zero.
• Streaming: You must manually read/write state variables. Use tinynn.converter.utils.tflite.parse_lstm_states(tflite_path) to find the indices of the state tensors, then use get_tensor and set_tensor to manage them. Note that state variables are 2D with shape [batch_size, hidden_size or input_size].