Stuck in Dimension problem in building QNN with pytorch

Hello, I’m trying to Build a Quantum Neural Network with pytorch and pennylane. The dimension error occurs.
The error occurs during defining the quantum Layer.

Colab Link: Google Colab

# Get data 
train = datasets.MNIST(root="data", download=True, train=True, transform=ToTensor())
dataset = DataLoader(train, 32)
n_qubits = 2
dev = qml.device("default.qubit", wires=n_qubits)

@qml.qnode(dev)
def qnode(inputs, weights_0, weight_1):
    print(inputs)
    qml.RX(inputs[0], wires=0)
    qml.RX(inputs[1], wires=1)
    qml.Rot(*weights_0, wires=0)
    qml.RY(weight_1, wires=1)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
weight_shapes = {"weights_0": 3, "weight_1": 1}
qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)
print(qlayer)
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(qlayer,
            nn.Conv2d(1, 32, (3, 3)),
            nn.ReLU(),
            nn.Conv2d(32, 64, (3, 3)),
            nn.ReLU(),
            nn.Conv2d(64, 64, (3, 3)),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(64 * (28 - 6) * (28 - 6), 10)
        )

    def forward(self, x):
        result = self.model(x)
        return result
# Instance of the neural network, loss, optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Instance of the neural network, loss, optimizer
clf = ImageClassifier().to('cpu')
opt = Adam(clf.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# Training flow 
if __name__ == "__main__":
    for epoch in range(1):  # train for 10 epochs
        for batch in dataset:
            X, y = batch
            X, y = X.to('cpu'), y.to(device)
            yhat = clf(X)
            loss = loss_fn(yhat, y)

            # Apply backprop 
            opt.zero_grad()
            loss.backward()
            opt.step()

        print(f"Epoch:{epoch} loss is {loss.item()}")

The Error Message I found:

RuntimeError                              Traceback (most recent call last)
<ipython-input-84-a98a57a9f607> in <cell line: 9>()
     12             X, y = batch
     13             X, y = X.to('cpu'), y.to(device)
---> 14             yhat = clf(X)
     15             loss = loss_fn(yhat, y)
     16 

10 frames
/usr/local/lib/python3.10/dist-packages/pennylane/qnn/torch.py in <listcomp>(.0)
    427 
    428         if len(x.shape) > 1:
--> 429             res = [torch.reshape(r, (x.shape[0], -1)) for r in res]
    430 
    431         return torch.hstack(res).type(x.dtype)

RuntimeError: shape '[896, -1]' is invalid for input of size 28

Hey @Hridoy_Chandra_Das!

At first glance, it’s difficult to see what’s causing the error. It looks like you’re doing image classification. Coincidentally, 896 / 28 = 32. Is your data being unintentionally flattened or something?

Hello, here is my colab link: Google Colab
Could you please check it?
I didn’t flatten the data. I don’t know what’s the problem, what should I do to solve it.

Hey @Hridoy_Chandra_Das,

I recommend this:

  1. Don’t do any parameter broadcasting / batching for now. Make sure your code works on one data point first:
for i, batch in enumerate(dataset):
  if i == 0:
    test_points = batch
    break

test_point = test_points[0][0] # one data point
print(test_point.size()) 
torch.Size([1, 28, 28])
  1. Go through each layer of your model, one by one, and try to identify which one gives you problems:
model = nn.Sequential(qlayer,
            #nn.Conv2d(1, 32, (3, 3)),
            #nn.ReLU(),
            #nn.Conv2d(32, 64, (3, 3)),
            #nn.ReLU(),
            #nn.Conv2d(64, 64, (3, 3)),
            #nn.ReLU(),
            #nn.Flatten(),
            #nn.Linear(64 * (28 - 6) * (28 - 6), 10)
        )

model(test_point)

Let me know if this helps!

Hello @isaacdevlugt , thanks for your rapid help. I followed your guidelines. It shows problems in the first layer. But I can’t fix it. Can you please help me build the layers? It would be beneficial for me.

model = nn.Sequential(qlayer,
            nn.Conv2d(1, 32, (3, 3)),
            # nn.ReLU(),
            # nn.Conv2d(32, 64, (2, 2)),
            # nn.ReLU(),
            # nn.Conv2d(64, 64, (3, 3)),
            # nn.ReLU(),
            # nn.Flatten(),
            # nn.Linear(64 * (28 - 6) * (28 - 6), 10)
        )
test_point = torch.randn(1, 28, 28)

model(test_point)
RuntimeError                              Traceback (most recent call last)
<ipython-input-23-0b1cd7fa4113> in <cell line: 12>()
     10         )
     11 
---> 12 model(test_point)

6 frames
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    454                             weight, bias, self.stride,
    455                             _pair(0), self.dilation, self.groups)
--> 456         return F.conv2d(input, weight, bias, self.stride,
    457                         self.padding, self.dilation, self.groups)
    458 

RuntimeError: Calculated padded input size per channel: (28 x 2). Kernel size: (3 x 3). Kernel size can't be greater than actual input size

Hey @Hridoy_Chandra_Das, apologies for the slow response.

The output after qlayer has size torch.Size([1, 28, 2]). Conv2d needs an input size (N,Cin,H,W) and outputs (N,Cout,Hout,Wout) (see here: Conv2d — PyTorch 2.2 documentation). Definitely check out the torch documentation for Conv2d and see if you can match the quantum layer’s output to the classical torch layer’s input!