SqueezeLight

By Jiaqi Gu, Chenghao Feng, Zheng Zhao, Zhoufeng Ying, Mingjie Liu, Ray T. Chen and David Z. Pan.

This repo is the official implementation of "SqueezeLight: Towards Scalable Optical Neural Networks with Multi-Operand Ring Resonators".

Introduction

SqueezeLight is a scalable optical neural network architecture based on multi-operand micro-ring resonators (MORRs), first experimentally demonstrated here. SqueezeLight is constructed with a nonlinear optical neuron based on multi-operand ring resonators to squeeze vector dot-product into a single device with low wavelength usage and built-in nonlinearity. A block-level squeezing technique with structured sparsity is exploited to support higher scalability via a fine-grained structured pruning technique. SqueezeLight achieves one-order-of-magnitude improvement in compactness and efficiency over previous MRR-based ONN designs with high fidelity, trainability, and robustness.

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/
    - morr_conv2d and morr_linear layers
    - activation.py: activation function
  - morr_*.py: SqueezeLight model definition
  - morr_base.py: base SqueezeLight model definition
- builder.py: build training utilities
script/: contains experiment scripts
train.py: training logic

Usage

Train SqueezeLight CNN model on MNIST dataset.
> python3 train.py config/mnist/cnn/train/train.yml or > nohup python3 script/mnist/cnn/launch_train.py &
Prune MORR in SqueezeLight CNN model on MNIST dataset.
> python3 train.py config/mnist/cnn/train/prune.yml or > nohup python3 script/mnist/cnn/launch_prune.py &

Citing SqueezeLight

@inproceedings{gu2021SqueezeLight,
  title={SqueezeLight: Towards Scalable Optical Neural Networks with Multi-Operand Ring Resonators},
  author={Jiaqi Gu and Chenghao Feng and Zheng Zhao and Zhoufeng Ying and Ray T. Chen and David Z. Pan},
  booktitle={Design, Automation and Test in Europe Conference (DATE)},
  year={2021}
}