<jupyter_start><jupyter_text>Stochastic Polyak Stepsizehttps://github.com/IssamLaradji/sps/  Setup<jupyter_code>!pip install git+https://github.com/IssamLaradji/sps.git



import sklearn
import scipy
import scipy.optimize
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns
import warnings

warnings.filterwarnings("ignore")
import itertools
import time
from functools import partial

import os

import numpy as np
from scipy.special import logsumexp

np.set_printoptions(precision=3)

np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)})
import torch
import torch.nn as nn
import torchvision

print("torch version {}".format(torch.__version__))
if torch.cuda.is_available():
    print(torch.cuda.get_device_name(0))
    print("current device {}".format(torch.cuda.current_device()))
else:
    print("Torch cannot find GPU")


def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
# torch.backends.cudnn.benchmark = True<jupyter_output>torch version 1.8.0+cu101
Tesla P100-PCIE-16GB
current device 0<jupyter_text>Binary logistic regression using Sklearn<jupyter_code># Fit the model usign sklearn

import sklearn.datasets
from sklearn.model_selection import train_test_split

iris = sklearn.datasets.load_iris()
X = iris["data"]
y = (iris["target"] == 2).astype(np.int)  # 1 if Iris-Virginica, else 0'
N, D = X.shape  # 150, 4

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

from sklearn.linear_model import LogisticRegression

# We set C to a large number to turn off regularization.
# We don't fit the bias term to simplify the comparison below.
log_reg = LogisticRegression(solver="lbfgs", C=1e5, fit_intercept=False)
log_reg.fit(X_train, y_train)

# extract probability of class 1
pred_sklearn_train = log_reg.predict_proba(X_train)[:, 1]
pred_sklearn_test = log_reg.predict_proba(X_test)[:, 1]

w_mle_sklearn = np.ravel(log_reg.coef_)
print(w_mle_sklearn)<jupyter_output>[-4.414 -9.111 6.539 12.686]<jupyter_text>PyTorch data and model<jupyter_code>from torch.utils.data import DataLoader, TensorDataset

# data. By default, numpy uses double but torch uses float
X_train_t = torch.tensor(X_train, dtype=torch.float)
y_train_t = torch.tensor(y_train, dtype=torch.float)

X_test_t = torch.tensor(X_test, dtype=torch.float)
y_test_t = torch.tensor(y_test, dtype=torch.float)

# To make things interesting, we pick a batchsize of B=33, which is not divisible by N=100
dataset = TensorDataset(X_train_t, y_train_t)
B = 33
dataloader = DataLoader(dataset, batch_size=B, shuffle=True)
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = torch.nn.Linear(D, 1, bias=False)

    def forward(self, x):
        y_pred = torch.sigmoid(self.linear(x))
        return y_pred[:, 0]  # (N,1) -> (N)
def criterion(ypred, ytrue, L2reg=0):
    loss = torch.nn.BCELoss(reduction="mean")(ypred, ytrue)
    w = 0.0
    for p in model.parameters():
        w += (p**2).sum()
    loss += L2reg * w
    return loss<jupyter_output><empty_output><jupyter_text>BFGS<jupyter_code>set_seed(0)
model = Model()
loss_trace = []
optimizer = torch.optim.LBFGS(model.parameters(), history_size=10)


def closure():
    optimizer.zero_grad()
    y_pred = model(X_train_t)
    loss = criterion(y_pred, y_train_t, L2reg=0)
    loss.backward()
    return loss.item()


max_iter = 10
for i in range(max_iter):
    loss = optimizer.step(closure)
    loss_trace.append(loss)
plt.figure()
plt.plot(loss_trace)

pred_sgd_train = model(X_train_t).detach().numpy()
pred_sgd_test = model(X_test_t).detach().numpy()


print("predicitons on test set using sklearn")
print(pred_sklearn_test)
print("predicitons on test set using sgd")
print(pred_sgd_test)
set_seed(0)
model = Model()
loss_trace = []
optimizer = torch.optim.LBFGS(model.parameters(), history_size=10)


def closure():
    optimizer.zero_grad()
    y_pred = model(X_train_t)
    loss = criterion(y_pred, y_train_t, L2reg=1e-4)
    loss.backward()
    return loss.item()


max_iter = 10
for i in range(max_iter):
    loss = optimizer.step(closure)
    loss_trace.append(loss)
plt.figure()
plt.plot(loss_trace)

pred_sgd_train = model(X_train_t).detach().numpy()
pred_sgd_test = model(X_test_t).detach().numpy()


print("predicitons on test set using sklearn")
print(pred_sklearn_test)
print("predicitons on test set using sgd")
print(pred_sgd_test)<jupyter_output>predicitons on test set using sklearn
[0.002 0.000 1.000 0.012 0.002 0.000 0.000 0.979 0.740 0.000 0.706 0.000
 0.000 0.000 0.000 0.001 1.000 0.000 0.009 1.000 0.000 0.650 0.000 1.000
 0.094 0.998 1.000 1.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.998
 0.000 0.000 0.000 0.000 0.999 0.000 0.000 0.000 0.000 0.000 0.281 0.909
 0.000 0.999]
predicitons on test set using sgd
[0.012 0.000 1.000 0.039 0.009 0.000 0.000 0.907 0.617 0.001 0.634 0.000
 0.000 0.000 0.000 0.008 0.999 0.003 0.040 1.000 0.000 0.619 0.000 0.999
 0.205 0.985 0.998 0.998 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.986
 0.001 0.000 0.000 0.000 0.995 0.004 0.003 0.000 0.000 0.001 0.387 0.877
 0.001 0.994]<jupyter_text>SGD<jupyter_code>nepochs = 100
learning_rate = 1e-1
loss_trace = []

set_seed(0)
model = Model()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0)

for epoch in range(nepochs):
    for X, y in dataloader:
        y_pred = model(X)
        loss = criterion(y_pred, y, L2reg=0)
        loss_trace.append(loss)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
plt.figure()
plt.plot(loss_trace)
plt.ylim([0, 2])

pred_sgd_train = model(X_train_t).detach().numpy()
pred_sgd_test = model(X_test_t).detach().numpy()


print("predicitons on test set using sklearn")
print(pred_sklearn_test)
print("predicitons on test set using sgd")
print(pred_sgd_test)<jupyter_output><empty_output><jupyter_text>Momentum<jupyter_code>nepochs = 100
loss_trace = []

set_seed(0)
model = Model()

optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

for epoch in range(nepochs):
    for X, y in dataloader:
        y_pred = model(X)
        loss = criterion(y_pred, y, L2reg=0)
        loss_trace.append(loss)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
plt.figure()
plt.plot(loss_trace)
plt.ylim([0, 2])

pred_sgd_train = model(X_train_t).detach().numpy()
pred_sgd_test = model(X_test_t).detach().numpy()


print("predicitons on test set using sklearn")
print(pred_sklearn_test)
print("predicitons on test set using sgd")
print(pred_sgd_test)<jupyter_output>predicitons on test set using sklearn
[0.002 0.000 1.000 0.012 0.002 0.000 0.000 0.979 0.740 0.000 0.706 0.000
 0.000 0.000 0.000 0.001 1.000 0.000 0.009 1.000 0.000 0.650 0.000 1.000
 0.094 0.998 1.000 1.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.998
 0.000 0.000 0.000 0.000 0.999 0.000 0.000 0.000 0.000 0.000 0.281 0.909
 0.000 0.999]
predicitons on test set using sgd
[0.364 0.000 0.988 0.374 0.258 0.000 0.086 0.643 0.565 0.138 0.630 0.000
 0.000 0.000 0.000 0.279 0.945 0.200 0.439 0.944 0.000 0.654 0.000 0.934
 0.626 0.783 0.931 0.922 0.000 0.000 0.000 0.000 0.090 0.000 0.000 0.819
 0.169 0.000 0.000 0.000 0.901 0.247 0.193 0.000 0.000 0.152 0.644 0.808
 0.119 0.886]<jupyter_text>SPS<jupyter_code>import sps

set_seed(0)
model = Model()
score_list = []

opt = sps.Sps(model.parameters(), c=0.5, eta_max=1, adapt_flag="constant")
# , fstar_flag=True)
# c=0.2 blows up

nepochs = 100
for epoch in range(nepochs):
    for X, y in dataloader:

        def closure():
            loss = criterion(model(X), y, L2reg=1e-4)
            loss.backward()
            return loss

        opt.zero_grad()
        loss = opt.step(closure=closure)
        loss_trace.append(loss)

        # Record metrics
        score_dict = {"epoch": epoch}
        score_dict["step_size"] = opt.state.get("step_size", {})
        score_dict["step_size_avg"] = opt.state.get("step_size_avg", {})
        score_dict["train_loss"] = loss

        score_list += [score_dict]
import pandas as pd

df = pd.DataFrame(score_list)
df.head()
plt.figure()
plt.plot(df["train_loss"])
plt.ylim([0, 2])

pred_sgd_train = model(X_train_t).detach().numpy()
pred_sgd_test = model(X_test_t).detach().numpy()

print("predicitons on test set using sklearn")
print(pred_sklearn_test)
print("predicitons on test set using sgd")
print(pred_sgd_test)
plt.figure()
plt.plot(df["step_size"])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13154858066005076, 'fraction_numerical': 0.1065081929379183, 'mean_word_length': 4.065458796025716, 'pattern_counts': {'":': 2, '<': 25, '<?xml version=': 0, '>': 26, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '297395', 'license': 'MIT', 'n_tokens_mistral': 4125, 'n_tokens_neox': 3436, 'n_words': 783, 'path': 'notebooks/misc/sps_logreg_torch.ipynb', 'repo_name': 'karm-patel/pyprobml', 'size': '8666'}