Hi,

I am currently testing out some hybrid neural networks with the MNIST dataset. The problem is that it is very slow.

5 epochs takes around 4/5 hours.

I have tried Forest pyQVM because it was shown to be the fastest in Speeding up grad computation but it was a lot slower.

Any suggestions would be appreciated.

Thanks for your help!

```
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import pennylane as qml
from pennylane import numpy as np
# Params
q_depth = 2
n_qubits = 4
q_delta = 1
input_size = 784
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001
dev = qml.device("default.qubit", wires=4)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def RY_layer(w):
"""Layer of parametrized qubit rotations around the y axis.
"""
for idx, element in enumerate(w):
qml.RY(element, wires=idx)
@qml.qnode(dev, interface="torch")
def q_net(q_in, q_weights_flat):
# Reshape weights
q_weights = q_weights_flat.reshape(q_depth, n_qubits)
# Embed features in the quantum node
RY_layer(q_in)
for i in range(n_qubits - 1):
qml.CNOT(wires=[i, i+1])
for i in range(n_qubits):
qml.RX(q_weights[0][i], wires = i)
for i in range(n_qubits):
qml.RY(q_weights[1][i], wires = i)
return tuple(qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits))
# MNIST dataset
train_dataset = torchvision.datasets.MNIST(root='../data',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = torchvision.datasets.MNIST(root='../data',
train=False,
transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
class NeuralNet(nn.Module):
def __init__(self, input_size, num_classes):
super(NeuralNet, self).__init__()
self.relu = nn.ReLU()
self.pre_net = nn.Linear(input_size, n_qubits)
self.q_params = nn.Parameter(q_delta * torch.randn(q_depth * n_qubits))
self.post_net = nn.Linear(n_qubits, 10)
def forward(self, x):
q_in = self.pre_net(x)
q_in = self.relu(q_in)
q_out = torch.Tensor(0, n_qubits).to(device)
for elem in q_in:
q_out_elem = q_net(elem, self.q_params).float().unsqueeze(0)
q_out = torch.cat((q_out, q_out_elem))
q_out = self.relu(q_out)
out = self.post_net(q_out)
return out
model = NeuralNet(input_size, num_classes).to(device)
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# Train the model
total_step = len(train_loader)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
# Move tensors to the configured device
images = images.reshape(-1, 28*28).to(device)
labels = labels.to(device)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(i)
if (i+1) % 100 == 0:
print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, num_epochs, i+1, total_step, loss.item()))
# Test
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.reshape(-1, 28*28).to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: {} %'.format(100 * correct / total))
# Save
torch.save(model.state_dict(), 'model1.ckpt')
```