Memory-Efficient Multi-Level In-Situ Generation (MLG)

By Jiaqi Gu, Hanqing Zhu, Chenghao Feng, Mingjie Liu, Zixuan Jiang, Ray T. Chen and David Z. Pan.

This repo is the official implementation of “Towards Memory-Efficient Neural Networks via Multi-Level in situ Generation“.

Introduction

MLG is a general and unified framework to trade expensive memory transactions with ultra-fast on-chip computations, directly translating to performance improvement.
MLG explores the intrinsic correlations and bit-level redundancy within DNN kernels and propose a multi-level in situ generation mechanism with mixed-precision bases to achieve on-the-fly recovery of high-resolution parameters with minimum hardware overhead.
MLG can boost the memory efficiency by 10-20× with comparable accuracy over four state-of-theart designs, when benchmarked on ResNet-18/DenseNet121/MobileNetV2/V3 with various tasks

flow

We explore intra-kernel and cross-kernel correlation in the accuracy (blue curve) and memory compression ratio (black curve) space with ResNet18/CIFAR-10.
Our method generalizes prior DSConv and Blueprint Conv with better efficiency-performance trade-off.
teaser

On CIFAR-10/100 and ResNet-18/DenseNet-121, we surpass prior low-rank methods with 10-20x less weight storage cost.
exp

Dependencies

  • Python >= 3.6
  • pyutils >= 0.0.1. See pyutils for installation.
  • pytorch-onn >= 0.0.2. See pytorch-onn for installation.
  • Python libraries listed in requirements.txt
  • NVIDIA GPUs and CUDA >= 10.2

Structures

  • core/
    • models/
      • layers/
        • mlg_conv2d and mlg_linear: MLG layer definition
      • resnet.py: MLG-based ResNet definition
      • model_base.py: base model definition with all model utilities
    • builder.py: build training utilities
  • configs: YAML-based config files
  • scripts/: contains experiment scripts
  • train.py: training logic

Usage

  • Pretrain teacher model.
    > python3 train.py configs/cifar10/resnet18/train/pretrain.yml

  • Train MLG-based student model with L2-norm-based projection, knowledge distillation, multi-level orthonormality regularization, (Bi, Bo, qb, qu, qv) = (2, 44, 3, 6, 3).
    > python3 train.py configs/cifar10/resnet18/train/train.yml --teacher.checkpoint=path-to-teacher-ckpt --mlg.projection_alg=train --mlg.kd=1 --mlg.base_in=2 --mlg.base_out=44 --mlg.basis_bit=3 --mlg.coeff_in_bit=6 --mlg.coeff_out_bit=3 --criterion.ortho_weight_loss=0.05

  • Scripts for experiments are in ./scripts. For example, to run teacher model pretraining, you can write proper task setting in SCRIPT=scripts/cifar10/resnet18/pretrain.py and run
    > python3 SCRIPT

  • To train ML-based student model with KD and projection, you can write proper task setting in SCRIPT=scripts/cifar10/resnet18/train.py (need to provide the pretrained teacher checkpoint) and run
    > python3 SCRIPT

Citing Memory-Efficient Multi-Level In-Situ Generation (MLG)

@inproceedings{gu2021MLG,
  title={Towards Memory-Efficient Neural Networks via Multi-Level in situ Generation},
  author={Jiaqi Gu and Hanqing Zhu and Chenghao Feng and Mingjie Liu and Zixuan Jiang and Ray T. Chen and David Z. Pan},
  journal={International Conference on Computer Vision (ICCV)},
  year={2021}
}

Related Papers

  • Jiaqi Gu, Hanqing Zhu, Chenghao Feng, Mingjie Liu, Zixuan Jiang, Ray T. Chen, David Z. Pan, “Towards Memory-Efficient Neural Networks via Multi-Level in situ Generation,” ICCV, 2021. [paper | slides]