PyTorch auf M4 Mac auf MPS läuft im Vergleich zur CPU deutlich langsamer ⇐ Python

[Anonymous]

Ich verwende ein M4 MacBook Pro und versuche, ein einfaches NN für MNIST-Daten auszuführen. Die Leistung auf MPS soll besser sein als die auf CPU. Aber es ist dramatisch langsamer. Selbst für ein einfaches NN wie das folgende dauert es auf der CPU etwa 1 Sekunde, auf MPS jedoch etwa 8 Sekunden. Vermisse ich etwas?

Code: Select all

def fit(X, Y, epochs, model, optimizer):
for epoch in range(epochs):
y_pred = model.forward(X)

loss = F.binary_cross_entropy(y_pred, Y)

optimizer.zero_grad() # zero the gradients
loss.backward() # Compute new gradients
optimizer.step() # update the parameters (weights)

if (epoch % 2000 == 0):
print(f'Epoch: {epoch} | Loss: {loss.item()}')

class NeuralNet(nn.Module):
def __init__(self):
super().__init__()

self.fc1 = nn.Linear(X.shape[1], 3)
self.fc2 = nn.Linear(3, 1)

def forward(self, x):
x = F.sigmoid(self.fc1(x))
x = F.sigmoid(self.fc2(x))
return x

def predict(self, x):
output = self.forward(x)
return (output > 0.5).int()

model = NeuralNet().to(device=device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)