torchtrainloss.py

import numpy as np
import pandas as pd
import plotnine as p9

import torch
from torch.autograd import Variable
from torch.utils.data import TensorDataset, DataLoader, Dataset

np.random.seed(555)
torch.manual_seed(555)

num_sample = 300
num_feat = 10
num_out = 10

batch_size = 32
lr = 1e-4
epochs = 300

X = np.random.normal(0, 0.5, (num_sample, num_feat)).astype(np.float64)
W = np.random.normal(0, 0.5, (num_feat, num_out)).astype(np.float64)
b = np.random.normal(0, 0.5, (1, num_out)).astype(np.float64)
Y = np.dot(X, W) + b
X -= X.mean(0)

ds = TensorDataset(torch.from_numpy(X).double(), torch.from_numpy(Y).double())

# train torch models with and without shuffling
model1 = torch.nn.Linear(num_feat, num_out).double()
opt1 = torch.optim.SGD(model1.parameters(), lr=lr)
loss1 = torch.nn.MSELoss()

train_hist_torch = []

for epoch in range(epochs):

    train_batch_losses = []
    for x, y in DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=True):
        x_var, y_var = Variable(x, requires_grad=False), Variable(y, requires_grad=False)
        pred = model1(x_var)
        l = loss1(pred, y_var)
        train_batch_losses.append(l.data[0])

        opt1.zero_grad()
        l.backward()
        opt1.step()

    # save mean of all batch errors within the epoch
    train_hist_torch.append(np.array(train_batch_losses).mean())


from keras.models import Model
from keras.layers import Input, Dense
from keras.optimizers import SGD
from keras import backend as K


inputs = Input(shape=(num_feat,))
predictions = Dense(num_out, activation='linear')(inputs)
model = Model(inputs=inputs, outputs=predictions)

opt = SGD(lr=lr)
model.compile(optimizer=opt, loss='mse')
losses = model.fit(X, Y,
                   batch_size=batch_size,
                   epochs=epochs, verbose=0, shuffle=True)
        
train_hist_keras = losses.history['loss']
    
(p9.ggplot(pd.DataFrame({'torch_train_torch': train_hist_torch,
                         'torch_train_keras': train_hist_keras,
                         'epochs': range(len(train_hist_torch))}),
            p9.aes(x='epochs')) +
  p9.geom_path(p9.aes(y='torch_train_torch', color='"training loss (sgd, torch)"')) +
  p9.geom_path(p9.aes(y='torch_train_keras', color='"training loss (sgd, keras)"')) +
  p9.labs(color='Loss', y=' ') +
  p9.theme_minimal())