Question of parameter-shift method in broadcasting

PennyLane Help
6

Hi @Noah_niu !

Sorry it took me some time to look into this.
It doesn’t seem like there’s any obvious solution on how to make this work with TorchLayer.
Using finite-diff also doesn’t work.

You need to avoid using TorchLayer completely and make a class that inherits from nn.Module.

I made a working example based on the TorchLayers demo and the Transfer Learning demo. See the code below:

import torch
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_moons

# Set random seeds
torch.manual_seed(42)
np.random.seed(42)

X, y = make_moons(n_samples=200, noise=0.1)
y_ = torch.unsqueeze(torch.tensor(y), 1)  # used for one-hot encoded labels
y_hot = torch.scatter(torch.zeros((200, 2)), 1, y_, 1)

c = ["#1f77b4" if y_ == 0 else "#ff7f0e" for y_ in y]  # colours for each class
plt.axis("off")
plt.scatter(X[:, 0], X[:, 1], c=c)
plt.show()

import pennylane as qml
import torch.nn as nn

n_qubits = 2
n_layers = 6
dev = qml.device("default.qubit", wires=n_qubits)

@qml.qnode(dev,diff_method='parameter-shift')
def qnode(inputs, weights):
    qml.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

class HybridModel(nn.Module):

    def __init__(self):

        super().__init__()
        self.clayer_1 = nn.Linear(2, n_qubits)
        self.q_params = nn.Parameter(torch.randn(n_layers , n_qubits))
        self.clayer_2 = nn.Linear(n_qubits, 2)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, inputs):
        q_in = self.clayer_1(inputs)

        # Apply the quantum circuit to each element of the batch and append to q_out
        q_out = torch.Tensor(0, n_qubits)

        for elem in q_in:
            q_out_elem = torch.hstack(qnode(elem, self.q_params)).float().unsqueeze(0)
            q_out = torch.cat((q_out, q_out_elem))

        # return the two-dimensional prediction from the postprocessing layer
        x = self.clayer_2(q_out)
        return self.softmax(x)

model = HybridModel()

# Data and batches
X = torch.tensor(X, requires_grad=True).float()
y_hot = y_hot.float()

batch_size = 5
batches = 200 // batch_size

data_loader = torch.utils.data.DataLoader(
    list(zip(X, y_hot)), batch_size=batch_size, shuffle=True, drop_last=True
)

# Optimization
epochs = 6
loss = nn.L1Loss()
opt = torch.optim.SGD(model.parameters(), lr=0.2)

for epoch in range(epochs):

    running_loss = 0

    for xs, ys in data_loader:
        opt.zero_grad()

        loss_evaluated = loss(model(xs), ys)
        loss_evaluated.backward()

        opt.step()

        running_loss += loss_evaluated

    avg_loss = running_loss / batches
    print("Average loss over epoch {}: {:.4f}".format(epoch + 1, avg_loss))

y_pred = model(X)
predictions = torch.argmax(y_pred, axis=1).detach().numpy()

correct = [1 if p == p_true else 0 for p, p_true in zip(predictions, y)]
accuracy = sum(correct) / len(correct)
print(f"Accuracy: {accuracy * 100}%")

It took me a while to make it but I hope it can help you to see how to update your code by using nn.Parameter instead of TorchLayer and breaking up the batches with for elem in q_in:.
Let me know if you have any questions on the code specifics!