Merge pull request #78 from mit-han-lab/dev

[minor] fix VQE

Author: Hanrui-Wang

examples/simple_vqe/simple_vqe.py

@@ -1,14 +1,13 @@
 import torchquantum as tq
 import torch
-import torch.nn.functional as F
 from torchquantum.vqe_utils import parse_hamiltonian_file
-from torchquantum.datasets import VQE
 import random
 import numpy as np
 import argparse
 import torch.optim as optim
 
-from torch.optim.lr_scheduler import CosineAnnealingLR, ConstantLR
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torchquantum.measurement import expval_joint_analytical
 
 
 class QVQEModel(tq.QuantumModule):
@@ -32,55 +31,54 @@ def __init__(self, arch, hamil_info):
                                                    trainable=True,
                                                    circular=True
                                                    ))
-        self.measure = tq.MeasureMultipleTimes(
-            obs_list=hamil_info['hamil_list'])
 
-    def forward(self, q_device):
-        q_device.reset_states(bsz=1)
-        for k in range(self.n_blocks):
-            self.u3_layers[k](q_device)
-            self.cu3_layers[k](q_device)
-        x = self.measure(q_device)
-
-        hamil_coefficients = torch.tensor([hamil['coefficient'] for hamil in
-                                           self.hamil_info['hamil_list']],
-                                          device=x.device).double()
-
-        for k, hamil in enumerate(self.hamil_info['hamil_list']):
-            for wire, observable in zip(hamil['wires'], hamil['observables']):
-                if observable == 'i':
-                    x[k][wire] = 1
-            for wire in range(q_device.n_wires):
-                if wire not in hamil['wires']:
-                    x[k][wire] = 1
-
-        x = torch.cumprod(x, dim=-1)[:, -1].double()
-        x = torch.dot(x, hamil_coefficients)
+    def forward(self):
+        qdev = tq.QuantumDevice(n_wires=self.n_wires, bsz=1, device=next(self.parameters()).device)
 
-        if x.dim() == 0:
-            x = x.unsqueeze(0)
-
-        return x
+        for k in range(self.n_blocks):
+            self.u3_layers[k](qdev)
+            self.cu3_layers[k](qdev)
+        
+        expval = 0
+        for hamil in self.hamil_info['hamil_list']:
+            expval += expval_joint_analytical(qdev, observable=hamil["pauli_string"]) * hamil["coeff"]
 
+        return expval
 
-def train(dataflow, q_device, model, device, optimizer):
-    for _ in dataflow['train']:
-        outputs = model(q_device)
-        loss = outputs.mean()
 
+def train(model, optimizer, n_steps=1):
+    for _ in range(n_steps):
+        loss = model()
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
         print(f"Expectation of energy: {loss.item()}")
 
 
-def valid_test(dataflow, q_device, split, model, device):
+def valid_test(model):
     with torch.no_grad():
-        for _ in dataflow[split]:
-            outputs = model(q_device)
-    loss = outputs.mean()
+        loss = model()
+
+    print(f"validation: expectation of energy: {loss.item()}")
 
-    print(f"Expectation of energy: {loss}")
+
+def process_hamil_info(hamil_info):
+    hamil_list = hamil_info['hamil_list']
+    n_wires = hamil_info["n_wires"]
+    all_info = []
+
+    for hamil in hamil_list:
+        pauli_string = ""
+        for i in range(n_wires):
+            if i in hamil['wires']:
+                wire = hamil['wires'].index(i)
+                pauli_string += (hamil['observables'][wire].upper())
+            else:
+                pauli_string += "I"
+        all_info.append({"pauli_string": pauli_string,
+                            "coeff": hamil['coefficient']})
+    hamil_info['hamil_list'] = all_info
+    return hamil_info
 
 
 def main():
@@ -94,7 +92,7 @@ def main():
                         help='number of training epochs')
     parser.add_argument('--epochs', type=int, default=100,
                         help='number of training epochs')
-    parser.add_argument('--hamil_filename', type=str, default='./h2_new.txt',
+    parser.add_argument('--hamil_filename', type=str, default='./h2.txt',
                         help='number of training epochs')
 
     args = parser.parse_args()
@@ -108,49 +106,27 @@ def main():
     np.random.seed(seed)
     torch.manual_seed(seed)
 
-    dataset = VQE(steps_per_epoch=args.steps_per_epoch)
-
-    dataflow = dict()
-
-    for split in dataset:
-        if split == 'train':
-            sampler = torch.utils.data.RandomSampler(dataset[split])
-        else:
-            sampler = torch.utils.data.SequentialSampler(dataset[split])
-        dataflow[split] = torch.utils.data.DataLoader(
-            dataset[split],
-            batch_size=1,
-            sampler=sampler,
-            num_workers=1,
-            pin_memory=True)
-
-    hamil_info = parse_hamiltonian_file(args.hamil_filename)
+    hamil_info = process_hamil_info(parse_hamiltonian_file(args.hamil_filename))
 
     use_cuda = torch.cuda.is_available()
     device = torch.device("cuda" if use_cuda else "cpu")
-    model = QVQEModel(arch={"n_blocks": args.n_blocks},
-                       hamil_info=hamil_info)
+    model = QVQEModel(arch={"n_blocks": args.n_blocks}, hamil_info=hamil_info)
 
     model.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)
 
-    q_device = tq.QuantumDevice(n_wires=hamil_info['n_wires'])
-    q_device.reset_states(bsz=1)
-
     for epoch in range(1, n_epochs + 1):
         # train
         print(f"Epoch {epoch}, LR: {optimizer.param_groups[0]['lr']}")
-        train(dataflow, q_device, model, device, optimizer)
+        train(model, optimizer, n_steps=args.steps_per_epoch)
 
-        # valid
-        valid_test(dataflow, q_device, 'valid', model, device)
         scheduler.step()
 
     # final valid
-    valid_test(dataflow, q_device, 'valid', model, device)
+    valid_test(model)
 
 
 if __name__ == '__main__':