/torchquantum

A PyTorch-based framework for Quantum ML and ML for Quantum, with support for easy deployments on real quantum computers.

Primary LanguagePythonMIT LicenseMIT

torchquantum Logo

@inproceedings{hanruiwang2022quantumnas,
    title     = {Quantumnas: Noise-adaptive search for robust quantum circuits},
    author    = {Wang, Hanrui and Ding, Yongshan and Gu, Jiaqi and Li, Zirui and Lin, Yujun and Pan, David Z and Chong, Frederic T and Han, Song},
    booktitle = {The 28th IEEE International Symposium on High-Performance Computer Architecture (HPCA-28)},
    year      = {2022}
}

TorchQuantum

A PyTorch-based hybrid classical-quantum dynamic neural networks framework.

MIT License Read the Docs Discourse status Website

News

  • Welcome to contribute! Please contact us or post in the forum if you want to have new examples implemented by TorchQuantum or any other questions.
  • Qmlsys website goes online: qmlsys.mit.edu

Installation

git clone https://github.com/mit-han-lab/torchquantum.git
cd torchquantum
pip install --editable .
python fix_qiskit_parameterization.py

Papers using TorchQuantum

Usage

Construct quantum NN models as simple as constructing a normal pytorch model.

import torch.nn as nn
import torch.nn.functional as F 
import torchquantum as tq
import torchquantum.functional as tqf

class QFCModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.n_wires = 4
    self.q_device = tq.QuantumDevice(n_wires=self.n_wires)
    self.measure = tq.MeasureAll(tq.PauliZ)
    
    self.encoder_gates = [tqf.rx] * 4 + [tqf.ry] * 4 + \
                         [tqf.rz] * 4 + [tqf.rx] * 4
    self.rx0 = tq.RX(has_params=True, trainable=True)
    self.ry0 = tq.RY(has_params=True, trainable=True)
    self.rz0 = tq.RZ(has_params=True, trainable=True)
    self.crx0 = tq.CRX(has_params=True, trainable=True)

  def forward(self, x):
    bsz = x.shape[0]
    # down-sample the image
    x = F.avg_pool2d(x, 6).view(bsz, 16)
    
    # reset qubit states
    self.q_device.reset_states(bsz)
    
    # encode the classical image to quantum domain
    for k, gate in enumerate(self.encoder_gates):
      gate(self.q_device, wires=k % self.n_wires, params=x[:, k])
    
    # add some trainable gates (need to instantiate ahead of time)
    self.rx0(self.q_device, wires=0)
    self.ry0(self.q_device, wires=1)
    self.rz0(self.q_device, wires=3)
    self.crx0(self.q_device, wires=[0, 2])
    
    # add some more non-parameterized gates (add on-the-fly)
    tqf.hadamard(self.q_device, wires=3)
    tqf.sx(self.q_device, wires=2)
    tqf.cnot(self.q_device, wires=[3, 0])
    tqf.qubitunitary(self.q_device0, wires=[1, 2], params=[[1, 0, 0, 0],
                                                           [0, 1, 0, 0],
                                                           [0, 0, 0, 1j],
                                                           [0, 0, -1j, 0]])
    
    # perform measurement to get expectations (back to classical domain)
    x = self.measure(self.q_device).reshape(bsz, 2, 2)
    
    # classification
    x = x.sum(-1).squeeze()
    x = F.log_softmax(x, dim=1)

    return x

Features

  • Easy construction of parameterized quantum circuits in PyTorch.
  • Support batch mode inference and training on CPU/GPU.
  • Support dynamic computation graph for easy debugging.
  • Support easy deployment on real quantum devices such as IBMQ.

TODOs

  • Support more gates
  • Support compile a unitary with descriptions to speedup training
  • Support other measurements other than analytic method
  • In einsum support multiple qubit sharing one letter. So that more than 26 qubit can be simulated.
  • Support bmm based implementation to solve scalability issue
  • Support conversion from torchquantum to qiskit

Dependencies

  • Python >= 3.7
  • PyTorch >= 1.8.0
  • configargparse >= 0.14
  • GPU model training requires NVIDIA GPUs

MNIST Example

Train a quantum circuit to perform MNIST task and deploy on the real IBM Quito quantum computer as in mnist_example.py script:

cd examples/simple_mnist
python mnist_example.py

Files

File Description
devices.py QuantumDevice class which stores the statevector
encoding.py Encoding layers to encode classical values to quantum domain
functional.py Quantum gate functions
operators.py Quantum gate classes
layers.py Layer templates such as RandomLayer
measure.py Measurement of quantum states to get classical values
graph.py Quantum gate graph used in static mode
super_layer.py Layer templates for SuperCircuits
plugins/qiskit* Convertors and processors for easy deployment on IBMQ
examples/ More examples for training QML and VQE models

Contact

TorchQuantum Forum

Hanrui Wang hanrui@mit.edu