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}
}

Related Papers

• Jiaqi Gu, Chenghao Feng, Zheng Zhao, Zhoufeng Ying, Mingjie Liu, Ray T. Chen, David Z. Pan, "SqueezeLight: SqueezeLight: Towards Scalable Optical Neural Networks with Multi-Operand Ring Resonators," IEEE Design, Automation & Test in Europe Conference & Exhibition (DATE), 2021. [paper]
• Yunfei Sun, "Design and optimization of multi-operand ring resonator based efficient optical recurrent neural network", Master Thesis, 2021. [paper]
• Zhoufeng Ying, Chenghao Feng, Zheng Zhao, Richard Soref, David Z. Pan, Ray T. Chen, "Integrated multi-operand electro-optic logic gates for optical computing," Applied Physics Letters (APL), 2019. [paper]
• Zhoufeng Ying, Chenghao Feng, Zheng Zhao, David Z. Pan, Ray T. Chen, "Multi-Operand Directed Logic-Based Electro-Optic Gates for Ultracompact Optical Computing," IEEE Photonics Conference (IPC), 2019. [paper]