google-research/morph-net
Fast & Simple Resource-Constrained Learning of Deep Network Structure
What it solves
MorphNet is designed to shrink deep neural networks to meet specific resource constraints, such as memory, latency, or computational cost (FLOPs), without manually redesigning the network architecture. It allows a model to learn its own optimal structure by adjusting the number of output channels in each convolutional layer during training.
How it works
MorphNet uses continuous relaxation to treat network structure learning as a constrained optimization problem. It adds a regularizer to the training loss that targets a specific resource (like FLOPs or model size). As the model trains, this regularizer pushes the influence of certain filters down; once they fall below a certain threshold, those channels are marked for removal.
The project provides several regularization algorithms:
- LogisticSigmoid (FiGS): A probabilistic approach to channel regularization that requires adding gating layers to the model. This is the recommended method.
- Gamma: Used for models with BatchNorm layers.
- GroupLasso: Used for models without BatchNorm.
Who it’s for
Machine learning engineers and researchers who have an existing "seed network" and need to optimize it for deployment on resource-constrained hardware where latency and memory are critical.
Highlights
- Resource-Aware Optimization: Specifically targets FLOPs, model size, or hardware-specific latency.
- Differentiable Architecture Search: Supports both pruning as a pruning algorithm or as a full architecture search method.
- Automated Structure Learning: Automatically determines which channels to remove based on a target cost.
- Structure Exporter: Includes tools to the save the learned architecture in JSON format for final retraining.