M³ICRO-MOMMI

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

This repo is the official implementation of "M³ICRO: Machine Learning-Enabled Compact Photonic Tensor Core based on Programmable Multi-Operand Multimode Interference".

Introduction

M³ICRO is an ultra-compact photonic tensor core design based on programmable multi-operand multimode interference (MOMMI).

It consists of a multi-path tensor core topology with interleaved MOMMIs and modulators for coherent matrix multiplication.

A machine learning-assisted training flow enables differentiable optimization of the control signals of MOMMIs. A novel block unfolding method is proposed to allow efficient full-range real-to-real linear transform using coherent photonic tensor cores.

Dependencies

• Python >= 3.10
• pyutils >= 0.0.1. See pyutils for installation.
• pytorch-onn >= 0.0.5. See pytorch-onn for installation.
• Python libraries listed in requirements.txt
• NVIDIA GPUs and CUDA >= 11.7

Structures

• configs/: configuration files
• core/
  • models/
    • layers/
      • dpe_conv2d.py: Convolution layer definition
      • dpe_linear.py: Linear layer definition
      • dpe_layer_base.py: Base layer module
      • utils.py: Utility function
      • activation.py: Customized activation
    • dpe_nn_base.py: DPE neural network base module
    • dpe_resnet.py: DPE ResNet
    • dpe_mobilenet_v3.py: DPE MobileNetV3
    • resnet.py: Standard ResNet
    • mobilenet_v3.py: Standard MobileNetV3
    • dpe_base.py: Base module for NN-based device predictor
    • dpe.py: Module definition of NN-based device predictor
    • butterfly_utils.py: Customized modules for butterfly transform
    • utils.py: Utility functions
  • builder.py: Build training utilities
  • utils.py: Customized register hooks
• scripts/: contains experiment scripts
• data/: MMI simulation dataset
• train.py: training logic
• train_dpe.py: training logic for NN-based device predictor

Usage

• Pretrain an ideal software digital ResNet20 on CIFAR10 as the teacher model:
  python3 scripts/cifar10/resnet20/pretrain.py The pretrained checkpoint is available at ./checkpoint/cifar10/resnet20/ResNet20_CIFAR10_pretrain.pt
• Simulated (Lumerical EME) transfer matrices of 4x4, 5x5, and 10x10 programmable MOMMIs are under ./data/mmi/raw/
• Train NN-based differentiable device behavior estimator (DPE) on 4x4 MMI,
  python3 scripts/mmi/dpe/train.py The trained checkpoint is under ./checkpoint/mmi/dpe/pretrain. The pretrained checkpoints are available under the above directory.
• Train ResNet20 on 4x4 MOMMI-based PTC based on the trained DPE,
  > python3 scripts/cifar10/resnet20/train_m3icro.py The training log is under ./log and the trained checkpoint is under ./checkpoint/cifar10/resnet20/train
• For comparison, butterfly and FFT-based PTCs can be trained as
  python3 scripts/cifar10/resnet20/train_butterfly.py

Citing M³ICRO

@article{gu2023M3ICRO,
  title={M3ICRO: Machine Learning-Enabled Compact Photonic Tensor Core based on PRogrammable Multi-Operand Multimode Interference},
  author={Jiaqi Gu and Hanqing Zhu and Chenghao Feng and Zixuan Jiang and Ray T. Chen and David Z. Pan},
  journal={APL Machine Learning},
  year={2024}
}