Размер партии выше 1 капли производительности Pytorch

Question

Я пытаюсь использовать некоторые очень простые модели для регрессии с помощью кераса или pytorch.Поэтому я моделирую данные X=np.random.random((10000, 10)) и y = np.matmul(X, W), где W - случайная матрица размера (10, 1).

Я хочу изучать W с помощью модели с одним слоем и без функции активации.

Я использую SGD с lr=0.01, потерей MSE, 20 эпохами и хотел проверить влияние на размер партии.

Когда я использую Keras, использование размера партии, отличного от 1, не влияет на конечную производительность, и модель предсказывает отличные результаты с учетом этой проблемы с игрушкой.

Однако на pytorch он отлично работает, когда размер моего пакета равен 1, но как только я его изменяю, производительность сильно падает.

Я пытался изменить скорость обучения, но это не сильно помогло.

Кто-нибудь знает, почему он работает так плохо?

У меня есть небольшой код, чтобы опробовать его:

import numpy as np
import torch
import keras
from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F

# Data
X = np.random.random((10000, 10))
W = np.random.random((10, 1))
y = np.matmul(X, W).ravel()

batch_size = 32
n_epochs = 20

# Keras 
model = keras.models.Sequential()
model.add(keras.layers.Dense(units=1, input_dim=10))

sgd = keras.optimizers.SGD(lr=0.01)
model.compile(loss='MSE', optimizer=sgd)

model.fit(X, y, epochs=n_epochs, batch_size=batch_size, verbose=0)

# pytorch
train_dataset = TensorDataset(torch.Tensor(X), torch.Tensor(y))
params = {'batch_size': batch_size}
trainloader = DataLoader(train_dataset, **params)

net = torch.nn.Linear(10, 1)

optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
criterion = F.mse_loss

for epoch in range(n_epochs):
    for X_batch, y_batch in trainloader:
        optimizer.zero_grad()
        outputs = net(X_batch)
        loss = criterion(outputs, y_batch)
        loss.backward()
        optimizer.step()

# eval
X_test = np.random.random((100, 10))
y_test = np.matmul(X_test, W)

y_keras = model.predict(X_test)
y_torch = net(torch.Tensor(X_test)).detach().numpy()

print('error keras {}'.format(np.mean((y_test-y_keras)**2)))
print('error torch {}'.format(np.mean((y_test-y_torch)**2)))