File size: 2,660 Bytes
5cb9176 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import pennylane as qml
import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
def encode(n_qubits, inputs):
for wire in range(n_qubits):
qml.RX(inputs[wire], wires=wire)
def layer(n_qubits, y_weight, z_weight):
for wire, y_weight in enumerate(y_weight):
qml.RY(y_weight, wires=wire)
for wire, z_weight in enumerate(z_weight):
qml.RZ(z_weight, wires=wire)
for wire in range(n_qubits):
qml.CZ(wires=[wire, (wire + 1) % n_qubits])
def measure(n_qubits):
return [qml.expval(qml.PauliZ(wire)) for wire in range(n_qubits)]
def get_model(n_qubits, n_layers, data_reupload):
# NOTE: need to select an appropriate device
# dev = qml.device('lightning.gpu', wires=n_qubits)
dev = qml.device("default.qubit", wires=n_qubits)
shapes = {
"y_weights": (n_layers, n_qubits),
"z_weights": (n_layers, n_qubits)
}
@qml.qnode(dev, interface='torch')
def circuit(inputs, y_weights, z_weights):
for layer_idx in range(n_layers):
if (layer_idx == 0) or data_reupload:
encode(n_qubits, inputs)
layer(n_qubits, y_weights[layer_idx], z_weights[layer_idx])
return measure(n_qubits)
model = qml.qnn.TorchLayer(circuit, shapes)
return model
class QuantumNet(nn.Module):
def __init__(self, n_layers, w_input, w_output, data_reupload):
super(QuantumNet, self).__init__()
self.n_qubits = 2
self.n_actions = 3
self.data_reupload = data_reupload
self.q_layers = get_model(n_qubits=self.n_qubits, n_layers=n_layers, data_reupload=data_reupload)
# convert from 2 qubits to 3 actions
# not adding more complexity here because we want to learn through quantum circuit
self.layer1 = nn.Linear(2, 3)
if w_input:
self.w_input = Parameter(torch.Tensor(self.n_qubits))
nn.init.normal_(self.w_input)
else:
self.register_parameter("w_input", None)
if w_output:
self.w_output = Parameter(torch.Tensor(self.n_actions))
nn.init.normal_(self.w_output, mean=90.0)
else:
self.register_parameter("w_output", None)
def forward(self, inputs):
if self.w_input is not None:
inputs = inputs * self.w_input
inputs = torch.atan(inputs)
q_outputs = self.q_layers(inputs)
q_outputs = (1 + q_outputs) / 2
outputs = self.layer1(q_outputs)
if self.w_output is not None:
outputs = outputs * self.w_output
else:
outputs = 90 * outputs
return outputs
|