[Fix] Encoding for density matrix. An example in mnist_new_noise.py

Author: Zhuoyang Ye

examples/PauliSumOp/pauli_sum_op_noise.py

@@ -100,10 +100,23 @@ def forward(self, x, use_qiskit=False):
         bsz = x.shape[0]
         x = F.avg_pool2d(x, 6).view(bsz, 16)
 
+
+        print("Shape 1:")
+        print(self.q_device.states.shape)
         self.encoder(self.q_device, x)
         self.q_layer(self.q_device)
+
+
+
+        print("X shape before measurement")
+        print(x.shape)
+
         x = self.measure(self.q_device)
 
+
+        print("X shape after measurement")
+        print(x.shape)
+
         x = x.reshape(bsz, 2, 2).sum(-1).squeeze()
         x = F.log_softmax(x, dim=1)
 

examples/amplitude_encoding_mnist/mnist_example.py

@@ -0,0 +1,223 @@
+"""
+MIT License
+
+Copyright (c) 2020-present TorchQuantum Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+import argparse
+
+import torchquantum as tq
+import torchquantum.functional as tqf
+
+from torchquantum.dataset import MNIST
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+import random
+import numpy as np
+
+
+class QFCModel(tq.QuantumModule):
+    class QLayer(tq.QuantumModule):
+        def __init__(self):
+            super().__init__()
+            self.n_wires = 4
+            self.random_layer = tq.RandomLayer(
+                n_ops=50, wires=list(range(self.n_wires))
+            )
+
+            # gates with trainable parameters
+            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)
+
+        @tq.static_support
+        def forward(self, q_device: tq.NoiseDevice):
+            """
+            1. To convert tq QuantumModule to qiskit or run in the static
+            model, need to:
+                (1) add @tq.static_support before the forward
+                (2) make sure to add
+                    static=self.static_mode and
+                    parent_graph=self.graph
+                    to all the tqf functions, such as tqf.hadamard below
+            """
+            self.q_device = q_device
+
+            self.random_layer(self.q_device)
+
+            # some trainable gates (instantiated 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, static=self.static_mode, parent_graph=self.graph
+            )
+            tqf.sx(
+                self.q_device, wires=2, static=self.static_mode, parent_graph=self.graph
+            )
+            tqf.cnot(
+                self.q_device,
+                wires=[3, 0],
+                static=self.static_mode,
+                parent_graph=self.graph,
+            )
+
+    def __init__(self):
+        super().__init__()
+        self.n_wires = 4
+        self.q_device = tq.NoiseDevice(n_wires=self.n_wires,
+                                       noise_model=tq.NoiseModel(kraus_dict={"Bitflip": 0.08, "Phaseflip": 0.08})
+                                       )
+        self.encoder = tq.AmplitudeEncoder()
+
+        self.q_layer = self.QLayer()
+        self.measure = tq.MeasureAll_density(tq.PauliZ)
+
+    def forward(self, x, use_qiskit=False):
+        bsz = x.shape[0]
+        x = F.avg_pool2d(x, 6).view(bsz, 16)
+        self.encoder(self.q_device, x)
+        self.q_layer(self.q_device)
+        x = self.measure(self.q_device)
+        x = x.reshape(bsz, 2, 2).sum(-1).squeeze()
+        x = F.log_softmax(x, dim=1)
+        return x
+
+
+def train(dataflow, model, device, optimizer):
+    for feed_dict in dataflow["train"]:
+        inputs = feed_dict["image"].to(device)
+        targets = feed_dict["digit"].to(device)
+
+        outputs = model(inputs)
+        loss = F.nll_loss(outputs, targets)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        print(f"loss: {loss.item()}", end="\r")
+
+
+def valid_test(dataflow, split, model, device, qiskit=False):
+    target_all = []
+    output_all = []
+    with torch.no_grad():
+        for feed_dict in dataflow[split]:
+            inputs = feed_dict["image"].to(device)
+            targets = feed_dict["digit"].to(device)
+
+            outputs = model(inputs, use_qiskit=qiskit)
+
+            target_all.append(targets)
+            output_all.append(outputs)
+        target_all = torch.cat(target_all, dim=0)
+        output_all = torch.cat(output_all, dim=0)
+
+    _, indices = output_all.topk(1, dim=1)
+    masks = indices.eq(target_all.view(-1, 1).expand_as(indices))
+    size = target_all.shape[0]
+    corrects = masks.sum().item()
+    accuracy = corrects / size
+    loss = F.nll_loss(output_all, target_all).item()
+
+    print(f"{split} set accuracy: {accuracy}")
+    print(f"{split} set loss: {loss}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--static", action="store_true", help="compute with " "static mode"
+    )
+    parser.add_argument("--pdb", action="store_true", help="debug with pdb")
+    parser.add_argument(
+        "--wires-per-block", type=int, default=2, help="wires per block int static mode"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=5, help="number of training epochs"
+    )
+
+    args = parser.parse_args()
+
+    if args.pdb:
+        import pdb
+
+        pdb.set_trace()
+
+    seed = 0
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    dataset = MNIST(
+        root="./mnist_data",
+        train_valid_split_ratio=[0.9, 0.1],
+        digits_of_interest=[3, 6],
+        n_test_samples=75,
+    )
+    dataflow = dict()
+
+    for split in dataset:
+        sampler = torch.utils.data.RandomSampler(dataset[split])
+        dataflow[split] = torch.utils.data.DataLoader(
+            dataset[split],
+            batch_size=256,
+            sampler=sampler,
+            num_workers=8,
+            pin_memory=True,
+        )
+
+    use_cuda = torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    model = QFCModel().to(device)
+
+    n_epochs = args.epochs
+    optimizer = optim.Adam(model.parameters(), lr=5e-3, weight_decay=1e-4)
+    scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs)
+
+    if args.static:
+        # optionally to switch to the static mode, which can bring speedup
+        # on training
+        model.q_layer.static_on(wires_per_block=args.wires_per_block)
+
+    for epoch in range(1, n_epochs + 1):
+        # train
+        print(f"Epoch {epoch}:")
+        train(dataflow, model, device, optimizer)
+        print(optimizer.param_groups[0]["lr"])
+
+        # valid
+        valid_test(dataflow, "valid", model, device)
+        scheduler.step()
+
+    # test
+    valid_test(dataflow, "test", model, device, qiskit=False)
+
+
+if __name__ == "__main__":
+    main()

examples/amplitude_encoding_mnist/mnist_example_noise.py

@@ -171,3 +171,4 @@ def train_tq(model, device, train_dl, epochs, loss_fn, optimizer):
 
 print("--Training--")
 train_losses = train_tq(model, device, train_dl, 1, loss_fn, optimizer)
+