Optical CNNs

Knowledge Distillation Circumvents Nonlinearity for OCNNs

Keywords: Python, Pytorch, Deep Learning, CNN, ONN

Introduction

This repository contains the code for the paper “Knowledge Distillation Circumvents Nonlinearity for Optical Convolutional Neural Networks”, which is available here, published in Applied Optics. Our open source code is here

Figure 1: Nonlinear CNN (top) and SCLC (botom)
Figure 2: Illustration of KD training.

Abstract

In recent years, Convolutional Neural Networks (CNNs) have enabled ubiquitous image processing applications. As such, CNNs require fast runtime (forward propagation) to process high-resolution visual streams in real time. This is still a challenging task even with state-of-the-art graphics and tensor processing units. The bottleneck in computational efficiency primarily occurs in the convolutional layers. Performing operations in the Fourier domain is a promising way to accelerate forward propagation since it transforms convolutions into elementwise multiplications, which are considerably faster to compute for large kernels. Furthermore, such computation could be implemented using an optical 4f system with orders of magnitude faster operation. However, a major challenge in using this spectral approach, as well as in an optical implementation of CNNs, is the inclusion of a nonlinearity between each convolutional layer, without which CNN performance drops dramatically. Here, we propose a Spectral CNN Linear Counterpart (SCLC) network architecture and develop a Knowledge Distillation (KD) approach to circumvent the need for a nonlinearity and successfully train such networks. While the KD approach is known in machine learning as an effective process for network pruning, we adapt the approach to transfer the knowledge from a nonlinear network (teacher) to a linear counterpart (student). We show that the KD approach can achieve performance that easily surpasses the standard linear version of a CNN and could approach the performance of the nonlinear network. Our simulations show that the possibility of increasing the resolution of the input image allows our proposed 4f optical linear network to perform more efficiently than a nonlinear network with the same accuracy on two fundamental image processing tasks: (i) object classification and (ii) semantic segmentation.

Example

Our training example are showed in jupyter notebook.

Requirements

pyTorch 1.4

Python 3.8

scikit-learn 0.21.2

matplotlib 3.1.0

numpy 1.16.4

Dataset

Object Classfication:

Cats Vs. Dogs

Cifar-10

HIGH-10, available here: train and test

Oject Segementation:

VOC2012

Car Segementation

Face Recognition

Citation

Please cite the Knowledge Distillation Circumvents Nonlinearity for Optical Convolutional Neural Networks when you use this code.

@article{xiang2022knowledge,
  title={Knowledge distillation circumvents nonlinearity for optical convolutional neural networks},
  author={Xiang, Jinlin and Colburn, Shane and Majumdar, Arka and Shlizerman, Eli},
  journal={Applied Optics},
  volume={61},
  number={9},
  pages={2173--2183},
  year={2022},
  publisher={Optica Publishing Group}