shuffle-torch.py

import os

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

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

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


def random_nbinom(mean, disp, n=None):
    eps = 1e-10
    # translate these into gamma parameters
    gamma_shape = 1 / (disp+eps)
    gamma_scale = (mean / (gamma_shape+eps))+eps
    gamma_samples = np.random.gamma(gamma_shape, gamma_scale, n)
    return np.random.poisson(gamma_samples)


def generate_dummy(num_sample, num_feat, num_out):
    dispersion = np.random.uniform(0, 2, [1, num_out])

    X = np.random.normal(0, 0.5, (num_sample, num_feat)).astype(np.float32)
    W = np.random.normal(0, 0.5, (num_feat, num_out)).astype(np.float32)
    b = np.random.normal(0, 0.5, (1, num_out)).astype(np.float32)
    y_mean = np.exp(np.dot(X, W) + b)
    y_disp = np.zeros_like(y_mean) + dispersion

    Y = random_nbinom(mean=y_mean, disp=y_disp).astype(np.float32)
    return X, Y


def train(X, Y, model, loss, optimizer, epochs=100, val_split=0.1,
          batch_size=128, shuffle=True):

    X, Y = torch.from_numpy(X), torch.from_numpy(Y)
    dataset = TensorDataset(X, Y)

    if val_split > 0.:
        off = int(len(X) * (1.0 - val_split))
        train_data = TensorDataset(X[:off], Y[:off])
        val_data   = TensorDataset(X[off:], Y[off:])
        val_loader = DataLoader(val_data, batch_size=len(val_data), shuffle=False)
    else:
        train_data, val_data = dataset, None

    loader = DataLoader(train_data, batch_size=batch_size, shuffle=shuffle)

    train_hist = []
    val_hist = []

    for epoch in tqdm(range(epochs)):

        train_batch_losses = []

        for x, y in loader:
            x_var, y_var = Variable(x), Variable(y)
            pred = model(x_var)
            l = loss(pred, y_var)
            train_batch_losses.append(l.data[0])

            optimizer.zero_grad()
            l.backward()
            optimizer.step()

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

        if val_data:
            for x, y in val_loader:
                x_var, y_var = Variable(x), Variable(y)
                pred = model(x_var)
                l = loss(pred, y_var)
                val_hist.append(l.data[0])

    return {'model': model,
            'val_loss': val_hist,
            'loss': train_hist}

# generate random data via NB regression
np.random.seed(555)
torch.manual_seed(555)

num_sample = 1000
num_feat = 10
num_out = 10
X, Y = generate_dummy(num_sample, num_feat, num_out)

# train torch models with and without shuffling
model1 = torch.nn.Linear(num_feat, num_out)
opt1 = torch.optim.RMSprop(model1.parameters(), lr=0.01)
ret_shuf = train(X, Y, model1, torch.nn.MSELoss(), shuffle=True, optimizer=opt1)

model2 = torch.nn.Linear(num_feat, num_out)
opt2 = torch.optim.RMSprop(model2.parameters(), lr=0.01)
ret_noshuf = train(X, Y, model2, torch.nn.MSELoss(), shuffle=False, optimizer=opt2)

# train keras models with and without shuffling

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

opt = RMSprop(lr=0.01)
model.compile(optimizer=opt, loss='mse')
losses = model.fit(X, Y,
                   batch_size=128,
                   validation_split=0.1,
                   epochs=100, verbose=0, shuffle=True)

keras_shuf = losses.history

# now without shuffling
K.clear_session()

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

opt = RMSprop(lr=0.01)
model.compile(optimizer=opt, loss='mse')
losses = model.fit(X, Y,
                   batch_size=128,
                   validation_split=0.1,
                   epochs=100, verbose=0, shuffle=False)

keras_noshuf = losses.history

(p9.ggplot(pd.DataFrame({'torch_train_shuf': ret_shuf['loss'],
                         'torch_val_shuf':   ret_shuf['val_loss'],
                         'torch_train_noshuf': ret_noshuf['loss'],
                         'torch_val_noshuf': ret_noshuf['val_loss'],
                         'keras_train_shuf': keras_shuf['loss'],
                         'keras_val_shuf': keras_shuf['val_loss'],
                         'keras_train_noshuf': keras_noshuf['loss'],
                         'keras_val_noshuf': keras_noshuf['val_loss'],
                         'epochs': range(len(ret_shuf['loss']))}),
            p9.aes(x='epochs')) +
  p9.geom_path(p9.aes(y='torch_train_shuf', color='"train loss (Torch, shuffled)"')) +
  p9.geom_path(p9.aes(y='torch_val_shuf', color='"val loss (Torch, shuffled)"')) +
  p9.geom_path(p9.aes(y='torch_train_noshuf', color='"train loss (Torch, not shuffled)"')) +
  p9.geom_path(p9.aes(y='torch_val_noshuf', color='"val loss (Torch, not shuffled)"')) +
  p9.geom_path(p9.aes(y='keras_train_shuf', color='"train loss (Keras, shuffled)"')) +
  p9.geom_path(p9.aes(y='keras_val_shuf', color='"val loss (Keras, shuffled)"')) +
  p9.geom_path(p9.aes(y='keras_train_noshuf', color='"train loss (Keras, not shuffled)"')) +
  p9.geom_path(p9.aes(y='keras_val_noshuf', color='"val loss (Keras, not shuffled)"')) +
  p9.labs(color='Loss', y=' ') +
  p9.theme_minimal() +
  p9.scale_color_brewer(type='qualitative',
                        palette='Paired')).save('shuffle.png', width=8, height=7)