feat: separate inputs and weights for qnn

quafu/algorithms/ansatz.py

@@ -159,25 +159,34 @@ def __init__(
         self._transformer = InterfaceProvider.get(interface)
         self._layers = layers
 
-        # FIXME(zhaoyilun): don't use this default value
-        self._weights = np.empty((1, 1))
+        self._weights = None
 
         self._backend = backend
         super().__init__(num_qubits)
 
-    def __call__(self, features):
+    def __call__(self, inputs):
         """Compute outputs of QNN given input features"""
         from .estimator import Estimator
 
         estimator = Estimator(self, backend=self._backend)
-        return self._transformer.execute(self, features, estimator=estimator)
+        return self._transformer.execute(self, inputs, estimator=estimator)
 
     def _build(self):
         """Essentially initialize weights using transformer"""
         self.add_gates(self._layers)
 
-        self._weights = self._transformer.init_weights((1, self.num_parameters))
+        self._weights = self._transformer.init_weights((1, self.num_tunable_parameters))
 
     @property
     def weights(self):
         return self._weights
+
+    @property
+    def num_tunable_parameters(self):
+        num_tunable_params = 0
+        for g in self.gates:
+            paras = g.paras
+            for p in paras:
+                if hasattr(p, "tunable") and p.tunable:
+                    num_tunable_params += 1
+        return num_tunable_params

quafu/algorithms/interface/torch.py

@@ -17,7 +17,9 @@
 
 import numpy as np
 import torch
+from quafu.algorithms.ansatz import QuantumNeuralNetwork
 from quafu.algorithms.estimator import Estimator
+from torch import nn
 
 from quafu import QuantumCircuit
 
@@ -56,14 +58,7 @@ def execute(
             "estimator": estimator,
         }
 
-        if method == "external":
-            return ExecuteCircuits.apply(parameters, kwargs)
-        if method == "internal":
-            from ..ansatz import QuantumNeuralNetwork
-
-            assert isinstance(circ, QuantumNeuralNetwork)
-            return ExecuteCircuits.apply(circ.weights, kwargs)
-        raise NotImplementedError(f"Unsupported execution method: {method}")
+        return ExecuteCircuits.apply(parameters, kwargs)
 
 
 class ExecuteCircuits(torch.autograd.Function):
@@ -91,3 +86,42 @@ def backward(ctx, grad_out):
         vjp = compute_vjp(jac, grad_out.numpy())
         vjp = torch.from_numpy(vjp)
         return vjp, None
+
+
+class ModuleWrapper(nn.Module):
+    """
+    A wrapper class to transform quafu circuit to a torch module
+    """
+
+    def __init__(self, qnn: QuantumNeuralNetwork):
+        """
+        Initialization of quafu torch module
+
+        Args:
+            circ (QuantumCircuit): the original parameterized quantum circuit
+        """
+        super().__init__()
+        self._qnn = qnn
+        if qnn.weights is not None:
+            self.weights = nn.parameter.Parameter(qnn.weights)
+        else:
+            self.weights = None
+
+    def forward(self, inputs: torch.Tensor):
+        """
+        Args:
+            inputs (torch.Tensor): raw input data or output from previous
+                classical/quantum layers.
+        """
+        # if weights are not empty, it will be combined with inputs to form
+        # the complete parameter vector and feed to the quantum circuit
+        bsz, _ = inputs.shape  # FIXME: currently we assume 2-D inputs
+
+        # use the last dimension since it is currently initialized as (1, D)
+        if self.weights is not None:
+            weight_dim = self.weights.size(-1)
+            weights_expanded = self.weights.expand(bsz, weight_dim)
+            inputs_to_circ = torch.cat((inputs, weights_expanded), dim=1)
+        else:
+            inputs_to_circ = inputs
+        return self._qnn(inputs_to_circ)

quafu/algorithms/templates/angle.py

@@ -14,7 +14,7 @@
 """Angel Embedding in Quantum Data embedding"""
 import numpy as np
 import quafu.elements.element_gates as qeg
-from quafu.elements import QuantumGate
+from quafu.elements import Parameter, QuantumGate
 
 ROT = {"X": qeg.RXGate, "Y": qeg.RYGate, "Z": qeg.RZGate}
 
@@ -45,7 +45,9 @@ def _build(self):
         gate_list = []
         for j in range(self.batch_size):
             for i in range(self.num_qubits):
-                gate = self.op(i, self.features[j, i])
+                gate = self.op(
+                    i, Parameter(f"phi_{i}", self.features[j, i], tunable=False)
+                )
                 gate_list.append(gate)
         return gate_list
 

quafu/circuits/quantum_circuit.py

@@ -74,6 +74,7 @@ def __init__(self, qnum: int, cnum: Optional[int] = None, name="", *args, **kwar
     @property
     def parameterized_gates(self):
         """Return the list of gates which the parameters are tunable"""
+        # FIXME: if we add parameterized gates after calling this function it will not work
         if not self._parameterized_gates:
             self._parameterized_gates = [g for g in self.gates if len(g.paras) != 0]
         return self._parameterized_gates
@@ -272,7 +273,10 @@ def _update_params(self, values, order=[]):
             order: For transplied circuit that change the order of variables,
             need pass the order to match untranspiled circuit's variable.
         """
-
+        if len(values) != len(self.variables):
+            raise CircuitError(
+                "The size of input values must be the same to the parameters"
+            )
         for i in range(len(values)):
             val = values[order[i]] if order else values[i]
             self._variables[i].value = val

quafu/elements/parameters.py

@@ -233,12 +233,13 @@ def log(self):
 
 
 class Parameter(ParameterExpression):
-    def __init__(self, name, value: float = 0.0):
+    def __init__(self, name, value: float = 0.0, tunable: bool = True):
         self.name = name
         self.value = float(value)
         self.operands = []
         self.funcs = []
         self.latex = self.name
+        self.tunable = tunable
 
     @property
     def pivot(self):

tests/quafu/algorithms/qnn_test.py

@@ -16,7 +16,8 @@
 import torch
 from quafu.algorithms.ansatz import QuantumNeuralNetwork
 from quafu.algorithms.gradients import compute_vjp, jacobian
-from quafu.algorithms.interface.torch import TorchTransformer
+from quafu.algorithms.interface.torch import ModuleWrapper, TorchTransformer
+from quafu.algorithms.templates.angle import AngleEmbedding
 from quafu.algorithms.templates.basic_entangle import BasicEntangleLayers
 from quafu.circuits.quantum_circuit import QuantumCircuit
 from quafu.elements import Parameter
@@ -81,7 +82,7 @@ def __init__(self, circ: QuantumCircuit):
 
     def forward(self, features):
         out = self.linear(features)
-        out = TorchTransformer.execute(self.circ, out, method="external")
+        out = TorchTransformer.execute(self.circ, out)
         return out
 
 
@@ -124,7 +125,7 @@ def _model_grad(self, model, batch_size):
 
         # TODO(zhaoyilun): Make out dimension configurable
         features = torch.randn(
-            batch_size, 3, requires_grad=True, dtype=torch.double
+            batch_size, 2, requires_grad=True, dtype=torch.double
         )  # batch_size=4, num_params=3
         outputs = model(features)
         targets = torch.randn(batch_size, 2, dtype=torch.double)
@@ -157,8 +158,9 @@ def test_torch_layer_standard_circuit(self):
     def test_torch_layer_qnn(self):
         """Use QuantumNeuralNetwork ansatz"""
         weights = np.random.randn(2, 2)
-        entangle_layer = BasicEntangleLayers(weights, 2)
-        qnn = QuantumNeuralNetwork(2, entangle_layer)
+        # entangle_layer = BasicEntangleLayers(weights, 2)
+        encoder_layer = AngleEmbedding(np.random.random((2,)), 2)
+        qnn = QuantumNeuralNetwork(2, encoder_layer)
         batch_size = 1
 
         # Legacy invokation style
@@ -182,7 +184,85 @@ def test_torch_layer_qnn_real_machine(self):
         model = ModelQuantumNeuralNetworkNative(qnn)
         self._model_grad(model, batch_size)
 
-    def test_classification_on_random_dataset(self, num_epochs, batch_size):
+    def test_module_wrapper(self):
+        weights = np.random.randn(2, 2)
+        entangle_layer = BasicEntangleLayers(weights, 2)
+        qnn = QuantumNeuralNetwork(2, entangle_layer)
+        qnn.measure([0, 1], [0, 1])
+
+        qlayer = ModuleWrapper(qnn)
+        params = qlayer.parameters()
+
+        assert np.allclose(
+            qlayer.weights.detach().numpy(), params.__next__().detach().numpy()
+        )
+
+    def test_classify_random_dataset_quantum(self, num_epochs, batch_size):
+        """Test a pure quantum nn training using a synthetic dataset
+
+        Args:
+            num_epochs: number of epoches for training
+            batch_size: batch size for training
+
+        """
+        # Define the hyperparameters
+        num_inputs = 2
+        num_classes = 2
+        learning_rate = 0.01
+
+        # Generate the dataset
+        dataset = _generate_random_dataset(num_inputs, 100)
+
+        # Create QNN
+        num_qubits = num_classes
+        weights = np.random.randn(num_qubits, 2)
+        encoder_layer = AngleEmbedding(np.random.random((2,)), num_qubits=2)
+        entangle_layer = BasicEntangleLayers(weights, 2)
+        qnn = QuantumNeuralNetwork(num_qubits, encoder_layer + entangle_layer)
+
+        # Create hybrid model
+        model = ModuleWrapper(qnn)
+        # model = mlp
+
+        # Define the loss function and optimizer
+        criterion = nn.CrossEntropyLoss()
+        optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
+
+        # Create data loader
+        data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+        # Train the model
+        for epoch in range(num_epochs):
+            for inputs, labels in data_loader:
+                # Forward pass
+                outputs = model(inputs)
+
+                # Compute the loss
+                loss = criterion(outputs, labels)
+
+                # Backward pass
+                optimizer.zero_grad()
+                loss.backward()
+
+                # Update the parameters
+                optimizer.step()
+
+            # Print the loss
+            print(f"Epoch {epoch + 1}/{num_epochs}: Loss = {loss.item()}")
+
+        # Evaluate the model on the dataset
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for inputs, labels in data_loader:
+                outputs = model(inputs)
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                correct += (predicted == labels.argmax(dim=1)).sum().item()
+
+        print(f"Accuracy: {100 * correct / total:.2f}%")
+
+    def test_classify_random_dataset_hybrid(self, num_epochs, batch_size):
         """Test e2e hybrid quantum-classical nn training using a synthetic dataset
 
         Args: