使用PyTorch与PennyLane构建QNN时遭遇维度不匹配错误

问题描述 #

使用PyTorch和PennyLane搭建量子神经网络（QNN），环境配置完成，但定义量子层时出现维度错误，怀疑输入数据与量子层预期形状不匹配。

原代码 #

# 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()}")

报错信息 #

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

错误原因 #

1. 输入维度不匹配：MNIST输入是(batch_size, 1, 28, 28)的4维张量，但量子层qnode仅接受2个输入特征（对应2个量子比特），直接传入高维图像会导致维度解析错误。
2. 网络顺序错误：量子层被放在卷积层之前，而卷积层需要4维输入，量子层输出的是2维张量，两者形状不兼容。
3. 设备不一致：数据和模型分别放在不同设备（CPU/GPU），可能引发隐性维度问题。

解决方案 #

调整网络结构，先用经典模块将图像特征压缩到量子层所需的维度，再接入量子层，最后完成分类任务。修改后的代码如下：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import pennylane as qml

# 数据加载
transform = transforms.ToTensor()
train = datasets.MNIST(root="data", download=True, train=True, transform=transform)
dataset = DataLoader(train, batch_size=32)

# 量子设备与量子节点定义
n_qubits = 2
dev = qml.device("default.qubit", wires=n_qubits)

@qml.qnode(dev)
def qnode(inputs, weights_0, weight_1):
    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)

# 修改后的分类器结构：经典特征压缩 → 量子层 → 分类头
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        # 经典编码器：将28×28图像压缩为2维特征，匹配量子层输入
        self.classical_encoder = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, 2)
        )
        # 量子层+分类输出层
        self.qnn_classifier = nn.Sequential(
            qlayer,
            nn.Linear(2, 10)
        )

    def forward(self, x):
        # 先提取并压缩图像特征
        features = self.classical_encoder(x)
        # 送入量子层完成分类
        return self.qnn_classifier(features)

# 训练配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
clf = ImageClassifier().to(device)
opt = torch.optim.Adam(clf.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# 训练流程
if __name__ == "__main__":
    for epoch in range(1):
        total_loss = 0.0
        for batch in dataset:
            X, y = batch
            X, y = X.to(device), y.to(device)
            yhat = clf(X)
            loss = loss_fn(yhat, y)

            opt.zero_grad()
            loss.backward()
            opt.step()

            total_loss += loss.item()
        
        avg_loss = total_loss / len(dataset)
        print(f"Epoch:{epoch} average loss: {avg_loss:.4f}")

关键修改点 #

• 新增classical_encoder模块：通过卷积、池化和全连接层，将高维图像压缩为量子层所需的2维特征。
• 调整网络顺序：经典特征提取在前，量子层接在特征之后，最后用全连接层完成10分类，解决形状不兼容问题。
• 统一设备：模型和数据都部署到同一设备（CPU/GPU），避免设备不一致引发的隐性错误。
• 优化训练日志：计算每个epoch的平均损失，提升训练过程的可读性。

内容的提问来源于stack exchange，提问作者Hrridoy V2