We release TorchONN, an open-source PyTorch-centric optical neural network library. We support fast and scalable development, training, and optimization of customized ONN layers and models. The tool runs on both CPU and GPU. CUDA-accelerated batched operators achieves 10-50X speedup over CPU-based matrix decomposition and unitary group parametrization.
PhD Student at UTDA, the University of Texas at Austin