docs/source/tutorials/torch_qml_cn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用 PyTorch 进行量子机器学习"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "在本教程中，我们展示了具有与 [mnist_qml](mnist_qml.ipynb) 相同设置的 MNIST 二元分类 QML 示例。\n",
    "这一次，我们使用 PyTorch 机器学习管道来构建 QML 模型。\n",
    "同样，本教程不是关于最佳 QML 实践或最佳 PyTorch 管道实践，而是关于 PyTorch 和 TensorCircuit 之间集成的演示。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import torch\n",
    "\n",
    "import tensorcircuit as tc\n",
    "\n",
    "K = tc.set_backend(\"tensorflow\")\n",
    "\n",
    "# 使用 TensorFlow 作为后端，同时将量子函数封装在 PyTorch 接口中"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([12163, 9]), torch.Size([12163]))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 我们使用和之前 notebook 相同的数据集与预处理\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
    "x_train = x_train[..., np.newaxis] / 255.0\n",
    "\n",
    "\n",
    "def filter_pair(x, y, a, b):\n",
    "    keep = (y == a) | (y == b)\n",
    "    x, y = x[keep], y[keep]\n",
    "    y = y == a\n",
    "    return x, y\n",
    "\n",
    "\n",
    "x_train, y_train = filter_pair(x_train, y_train, 1, 5)\n",
    "x_train_small = tf.image.resize(x_train, (3, 3)).numpy()\n",
    "x_train_bin = np.array(x_train_small > 0.5, dtype=np.float32)\n",
    "x_train_bin = np.squeeze(x_train_bin).reshape([-1, 9])\n",
    "y_train_torch = torch.tensor(y_train, dtype=torch.float32)\n",
    "x_train_torch = torch.tensor(x_train_bin)\n",
    "x_train_torch.shape, y_train_torch.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 ``torch_interface`` 包装量子函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 9\n",
    "nlayers = 3\n",
    "\n",
    "# 我们定义量子函数，\n",
    "# 注意这个函数是在 tensorflow 上运行的\n",
    "\n",
    "\n",
    "def qpred(x, weights):\n",
    "    c = tc.Circuit(n)\n",
    "    for i in range(n):\n",
    "        c.rx(i, theta=x[i])\n",
    "    for j in range(nlayers):\n",
    "        for i in range(n - 1):\n",
    "            c.cnot(i, i + 1)\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=weights[2 * j, i])\n",
    "            c.ry(i, theta=weights[2 * j + 1, i])\n",
    "    ypred = c.expectation_ps(z=[n // 2])\n",
    "    ypred = K.real(ypred)\n",
    "    return K.sigmoid(ypred)\n",
    "\n",
    "\n",
    "# 将函数包装成 pytorch 形式，但具有 tensorflow 速度！\n",
    "qpred_torch = tc.interfaces.torch_interface(qpred, jit=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在我们拥有 PyTorch 格式的可微功能之后，我们可以进一步将其包装为一个 Torch 模块（网络层）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QuantumNet(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.q_weights = torch.nn.Parameter(torch.randn([2 * nlayers, n]))\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        ypred = qpred_torch(inputs, self.q_weights)\n",
    "        return ypred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.4539, grad_fn=<FunBackward>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = QuantumNet()\n",
    "net(x_train_torch[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.7287, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.5947, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.5804, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6358, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6503, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "training time per step:  0.12587213516235352\n"
     ]
    }
   ],
   "source": [
    "criterion = torch.nn.BCELoss()\n",
    "opt = torch.optim.Adam(net.parameters(), lr=1e-2)\n",
    "nepochs = 500\n",
    "nbatch = 32\n",
    "times = []\n",
    "\n",
    "for epoch in range(nepochs):\n",
    "    index = np.random.randint(low=0, high=100, size=nbatch)\n",
    "    # index = np.arange(nbatch)\n",
    "    inputs, labels = x_train_torch[index], y_train_torch[index]\n",
    "    opt.zero_grad()\n",
    "\n",
    "    with torch.set_grad_enabled(True):\n",
    "        time0 = time.time()\n",
    "        yps = []\n",
    "        for i in range(nbatch):\n",
    "            yp = net(inputs[i])\n",
    "            yps.append(yp)\n",
    "        yps = torch.stack(yps)\n",
    "        loss = criterion(\n",
    "            torch.reshape(yps, [nbatch, 1]), torch.reshape(labels, [nbatch, 1])\n",
    "        )\n",
    "        loss.backward()\n",
    "        if epoch % 100 == 0:\n",
    "            print(loss)\n",
    "        opt.step()\n",
    "        time1 = time.time()\n",
    "        times.append(time1 - time0)\n",
    "\n",
    "print(\"training time per step: \", np.mean(time1 - time0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 批处理版本\n",
    "\n",
    "现在让我们尝试使用向量化版本来加快批量输入处理。请注意，本质上是 ``tf.vectorized_map`` 帮助批处理管道的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "qpred_vmap = K.vmap(qpred, vectorized_argnums=0)\n",
    "\n",
    "# `qpred_vmap` 是一个具有向量化能力的 TensorFlow 函数\n",
    "\n",
    "qpred_batch = tc.interfaces.torch_interface(qpred_vmap, jit=True)\n",
    "\n",
    "# 我们进一步将函数包装为 PyTorch 函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-6.2068e-03, -3.0100e-05, -1.0997e-02, -1.8381e-02, -9.1800e-02,\n",
      "          1.2481e-01, -6.5200e-02,  1.1176e-08,  7.4506e-09],\n",
      "        [-3.2353e-03,  3.4989e-03, -1.1344e-02, -1.6136e-02,  1.9075e-02,\n",
      "          2.1119e-02,  2.6881e-02, -1.1176e-08,  0.0000e+00],\n",
      "        [-1.1777e-02, -1.1572e-03, -5.0570e-03,  6.4838e-03, -5.5077e-02,\n",
      "         -3.4250e-02, -7.4506e-09, -1.1176e-08,  3.7253e-09],\n",
      "        [-1.4748e-02, -2.3818e-02, -4.3567e-02, -4.7879e-02,  1.2331e-01,\n",
      "          1.4314e-01,  3.7253e-09,  1.1176e-08,  3.7253e-09],\n",
      "        [-3.7253e-09,  3.7253e-09,  0.0000e+00,  0.0000e+00, -2.7574e-02,\n",
      "          7.4506e-09,  7.4506e-09, -1.1176e-08,  0.0000e+00],\n",
      "        [ 3.7253e-09,  3.7253e-09,  1.4901e-08, -7.4506e-09,  7.1655e-02,\n",
      "         -7.4506e-09,  3.7253e-09,  1.4901e-08,  0.0000e+00]])\n"
     ]
    }
   ],
   "source": [
    "# 测试 PyTorch 函数的 AD 支持\n",
    "\n",
    "w = torch.ones([2 * nlayers, n])\n",
    "w.requires_grad_()\n",
    "with torch.set_grad_enabled(True):\n",
    "    yps = qpred_batch(x_train_torch[:3], w)\n",
    "    loss = torch.sum(yps)\n",
    "    loss.backward()\n",
    "print(w.grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QuantumNetV2(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.q_weights = torch.nn.Parameter(torch.randn([2 * nlayers, n]))\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        ypred = qpred_batch(inputs, self.q_weights)\n",
    "        return ypred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.4706, 0.4706, 0.4767], grad_fn=<FunBackward>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net2 = QuantumNetV2()\n",
    "net2(x_train_torch[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "借助借用自 TensorFlow 的 vmap 基础设施，训练的性能大大提升！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6973, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6421, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6419, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6498, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6466, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "training time per step:  0.009107916531916371\n"
     ]
    }
   ],
   "source": [
    "criterion = torch.nn.BCELoss()\n",
    "opt = torch.optim.Adam(net2.parameters(), lr=1e-2)\n",
    "nepochs = 500\n",
    "nbatch = 32\n",
    "times = []\n",
    "for epoch in range(nepochs):\n",
    "    index = np.random.randint(low=0, high=100, size=nbatch)\n",
    "    # index = np.arange(nbatch)\n",
    "    inputs, labels = x_train_torch[index], y_train_torch[index]\n",
    "    opt.zero_grad()\n",
    "\n",
    "    with torch.set_grad_enabled(True):\n",
    "        time0 = time.time()\n",
    "        yps = net2(inputs)\n",
    "        loss = criterion(\n",
    "            torch.reshape(yps, [nbatch, 1]), torch.reshape(labels, [nbatch, 1])\n",
    "        )\n",
    "        loss.backward()\n",
    "        if epoch % 100 == 0:\n",
    "            print(loss)\n",
    "        opt.step()\n",
    "        time1 = time.time()\n",
    "        times.append(time1 - time0)\n",
    "print(\"training time per step: \", np.mean(times[1:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 具有经典后处理的混合模型\n",
    "\n",
    "我们现在构建了一个量子经典混合机器学习模型管道，其中输出测量结果被进一步馈送到经典的全连接层。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.2839,  0.3786,  0.0158,  0.1512,  0.1395,  0.1364,  0.1403,  0.1423,\n",
       "         -0.1285],\n",
       "        [ 0.2839,  0.3786,  0.0158,  0.1512,  0.1395,  0.1364,  0.1403,  0.1423,\n",
       "         -0.1285]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def qpreds(x, weights):\n",
    "    c = tc.Circuit(n)\n",
    "    for i in range(n):\n",
    "        c.rx(i, theta=x[i])\n",
    "    for j in range(nlayers):\n",
    "        for i in range(n - 1):\n",
    "            c.cnot(i, i + 1)\n",
    "        for i in range(n):\n",
    "            c.rx(i, theta=weights[2 * j, i])\n",
    "            c.ry(i, theta=weights[2 * j + 1, i])\n",
    "\n",
    "    return K.stack([K.real(c.expectation_ps(z=[i])) for i in range(n)])\n",
    "\n",
    "\n",
    "qpreds_vmap = K.vmap(qpreds, vectorized_argnums=0)\n",
    "qpreds_batch = tc.interfaces.torch_interface(qpreds_vmap, jit=True)\n",
    "\n",
    "qpreds_batch(x_train_torch[:2], torch.ones([2 * nlayers, n]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QuantumNetV3(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.q_weights = torch.nn.Parameter(torch.randn([2 * nlayers, n]))\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        ypred = qpreds_batch(inputs, self.q_weights)\n",
    "        return ypred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.2931,  0.5393, -0.0369, -0.0450,  0.0511, -0.0121,  0.0156, -0.0406,\n",
       "         -0.1330],\n",
       "        [ 0.2931,  0.5393, -0.0369, -0.0450,  0.0511, -0.0121,  0.0156, -0.0406,\n",
       "         -0.1330]], grad_fn=<FunBackward>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net3 = QuantumNetV3()\n",
    "net3(x_train_bin[:2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们现在用量子层 ``net3``构建一个混合模型，并在后面附加一个线性层。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.5500],\n",
       "        [0.5500]], grad_fn=<SigmoidBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = torch.nn.Sequential(QuantumNetV3(), torch.nn.Linear(9, 1), torch.nn.Sigmoid())\n",
    "\n",
    "model(x_train_bin[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6460, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.6086, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.5199, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.5697, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "tensor(0.5248, grad_fn=<BinaryCrossEntropyBackward0>)\n",
      "training time per step:  0.020270218113381304\n"
     ]
    }
   ],
   "source": [
    "criterion = torch.nn.BCELoss()\n",
    "opt = torch.optim.Adam(model.parameters(), lr=1e-2)\n",
    "nepochs = 500\n",
    "nbatch = 32\n",
    "times = []\n",
    "for epoch in range(nepochs):\n",
    "    index = np.random.randint(low=0, high=100, size=nbatch)\n",
    "    # index = np.arange(nbatch)\n",
    "    inputs, labels = x_train_torch[index], y_train_torch[index]\n",
    "    opt.zero_grad()\n",
    "\n",
    "    with torch.set_grad_enabled(True):\n",
    "        time0 = time.time()\n",
    "        yps = model(inputs)\n",
    "        loss = criterion(\n",
    "            torch.reshape(yps, [nbatch, 1]), torch.reshape(labels, [nbatch, 1])\n",
    "        )\n",
    "        loss.backward()\n",
    "        if epoch % 100 == 0:\n",
    "            print(loss)\n",
    "        opt.step()\n",
    "        time1 = time.time()\n",
    "        times.append(time1 - time0)\n",
    "print(\"training time per step: \", np.mean(times[1:]))"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "2d121f26c836fa50684100f1e49bcafe1050f1ef27bfc6ab5222e4b5b1a90afa"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}