add hybrid pipeline example all on gpu

Author: refraction-ray

.github/ISSUE_TEMPLATE/tc_enhancement_proposal.md

@@ -8,6 +8,8 @@ assignees: ""
 
 <!--Inspired from NEP: https://numpy.org/neps/nep-template.html-->
 
+<!-- If you have some small feature request or issue report, just open instead a plain issue -->
+
 # TEP - Title
 
 Author

CHANGELOG.md

@@ -12,7 +12,7 @@
 
 - Add `to_dlpack` and `from_dlpack` method on backends
 
-- Add dlpack path for interfaces
+- Add `enable_dlpack` option on interfaces and torchnn
 
 ### Changed
 
@@ -26,6 +26,8 @@
 
 - Fixed `numpy` method bug in pytorch backend when the input tensor requires grad (#24) and when the tensor is on GPU (#25)
 
+- Fixed `TorchLayer` parameter list auto registeration
+
 ## 0.2.1
 
 ### Added

examples/hybrid_gpu_pipeline.py

@@ -0,0 +1,108 @@
+"""
+quantum part in tensorflow or jax, neural part in torch, both on GPU,
+fantastic hybrid pipeline
+"""
+
+import os
+
+os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
+import time
+import numpy as np
+import tensorflow as tf
+import torch
+import tensorcircuit as tc
+
+K = tc.set_backend("tensorflow")
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+
+print(device)
+
+# dataset preparation
+
+(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+x_train = x_train[..., np.newaxis] / 255.0
+
+
+def filter_pair(x, y, a, b):
+    keep = (y == a) | (y == b)
+    x, y = x[keep], y[keep]
+    y = y == a
+    return x, y
+
+
+x_train, y_train = filter_pair(x_train, y_train, 1, 5)
+x_train_small = tf.image.resize(x_train, (3, 3)).numpy()
+x_train_bin = np.array(x_train_small > 0.5, dtype=np.float32)
+x_train_bin = np.squeeze(x_train_bin).reshape([-1, 9])
+y_train_torch = torch.tensor(y_train, dtype=torch.float32)
+x_train_torch = torch.tensor(x_train_bin)
+x_train_torch = x_train_torch.to(device=device)
+y_train_torch = y_train_torch.to(device=device)
+
+n = 9
+nlayers = 3
+
+# We define the quantum function,
+# note how this function is running on tensorflow
+
+
+def qpreds(x, weights):
+    c = tc.Circuit(n)
+    for i in range(n):
+        c.rx(i, theta=x[i])
+    for j in range(nlayers):
+        for i in range(n - 1):
+            c.cnot(i, i + 1)
+        for i in range(n):
+            c.rx(i, theta=weights[2 * j, i])
+            c.ry(i, theta=weights[2 * j + 1, i])
+
+    return K.stack([K.real(c.expectation_ps(z=[i])) for i in range(n)])
+
+
+# qpreds_vmap = K.vmap(qpreds, vectorized_argnums=0)
+# qpreds_batch = tc.interfaces.torch_interface(qpreds_vmap, jit=True, enable_dlpack=True)
+
+quantumnet = tc.TorchLayer(
+    qpreds,
+    weights_shape=[2 * nlayers, n],
+    use_vmap=True,
+    use_interface=True,
+    use_jit=True,
+    enable_dlpack=True,
+)
+
+
+model = torch.nn.Sequential(quantumnet, torch.nn.Linear(9, 1), torch.nn.Sigmoid())
+model = model.to(device=device)
+
+
+criterion = torch.nn.BCELoss()
+opt = torch.optim.Adam(model.parameters(), lr=1e-2)
+nepochs = 300
+nbatch = 32
+times = []
+for epoch in range(nepochs):
+    index = np.random.randint(low=0, high=100, size=nbatch)
+    # index = np.arange(nbatch)
+    inputs, labels = x_train_torch[index], y_train_torch[index]
+    opt.zero_grad()
+
+    with torch.set_grad_enabled(True):
+        time0 = time.time()
+        yps = model(inputs)
+        loss = criterion(
+            torch.reshape(yps, [nbatch, 1]), torch.reshape(labels, [nbatch, 1])
+        )
+        loss.backward()
+        if epoch % 100 == 0:
+            print(loss)
+        opt.step()
+        time1 = time.time()
+        times.append(time1 - time0)
+print("training time per step: ", np.mean(times[1:]))

tensorcircuit/interfaces/torch.py

@@ -70,12 +70,9 @@ def forward(ctx: Any, *x: Any) -> Any:  # type: ignore
             # (x, )
             if len(ctx.xdtype) == 1:
                 ctx.xdtype = ctx.xdtype[0]
+
             x = general_args_to_backend(x, enable_dlpack=enable_dlpack)
             y = fun(*x)
-            # if not is_sequence(y):
-            #     ctx.ydtype = [y.dtype]
-            # else:
-            #     ctx.ydtype = [yi.dtype for yi in y]
             ctx.ydtype = backend.tree_map(lambda s: s.dtype, y)
             if len(x) == 1:
                 x = x[0]

tensorcircuit/torchnn.py

@@ -22,6 +22,7 @@ def __init__(
         use_vmap: bool = True,
         use_interface: bool = True,
         use_jit: bool = True,
+        enable_dlpack: bool = False,
     ):
         """
         PyTorch nn Module wrapper on quantum function ``f``.
@@ -67,14 +68,16 @@ def qpred(x, weights):
         :type use_interface: bool, optional
         :param use_jit: whether jit ``f``, defaults to True
         :type use_jit: bool, optional
+        :param enable_dlpack: whether enbale dlpack in interfaces, defaults to False
+        :type enable_dlpack: bool, optional
         """
         super().__init__()
         if use_vmap:
             f = backend.vmap(f, vectorized_argnums=0)
         if use_interface:
-            f = torch_interface(f, jit=use_jit)
+            f = torch_interface(f, jit=use_jit, enable_dlpack=enable_dlpack)
         self.f = f
-        self.q_weights = []
+        self.q_weights = torch.nn.ParameterList()  # type: ignore
         if isinstance(weights_shape[0], int):
             weights_shape = [weights_shape]
         if not is_sequence(initializer):