pytorch 1d regression

regression1d.py

import numpy as np
import torch
from torch import nn
from itertools import count
import time

#%%
# make some data
N = 100

def func(x):
    y = 1.789 * x + 5.3
    return y

def make_data(numsamp):
    xdom = np.random.rand(numsamp).astype(np.float32)
    offset = (np.random.rand(numsamp).astype(np.float32) * 2 -1) * 0.15
    y = func(xdom) + offset
    return xdom.reshape(-1,1), y.reshape(-1,1)

train = make_data(N)
truth = func(train[0])

#%%

# linear regression for 1 variable
input_dim = 1
output_dim = 1

device = torch.device('cpu')
#device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu' )

x = torch.from_numpy(train[0]).to(device)
#x = x.view(10,-1,1)
#y = torch.from_numpy(train[1]).to(device) # adding noise to data
y = torch.from_numpy(truth).to(device)
#y = y.view(10,-1,1)

#%%
# Structured Model
class LinSolve(nn.Module):
    def __init__(self,_inD,_outD):
        super(LinSolve,self).__init__()
        
        self.fc1 = nn.Linear(_inD,_outD)
        
    def forward(self,x):
        x = self.fc1(x)
        return x

def trainLin(npX,npY,batchsize=10,maxEpoch=5000):
    model = LinSolve(1,1)
    
    betas = (0.9, 0.999)
    lr_decay = 1.0
    lr = 0.001
    optimizer = optim.Adam(model.parameters(), lr=lr, betas=betas)
    
    X = torch.from_numpy(npX)
    Y = torch.from_numpy(npY)
    
    print(X.shape)
    
    dataset = torch.utils.data.TensorDataset(X,Y)
    loader  = torch.utils.data.DataLoader(dataset, batch_size=batchsize, shuffle=True)
    
    loss_func = nn.MSELoss()
    
    t0 = time.time()
    for epoch in range(maxEpoch):
        for batch_idx, (x,y) in enumerate(loader):
            optimizer.zero_grad()
            
            pred = model(x)
            loss = loss_func(pred,y)
            
            loss.backward()
            optimizer.step()
            
        with torch.no_grad():
            out = model(X)
            full_loss = loss_func(out,Y).item()
            if full_loss < 1e-6:
                break
    t1 = time.time()
    
    print('{} epoch in {} sec\nLoss: {}'.format( epoch, t1-t0, full_loss ))
    return model,out.data.numpy()
 

#%%

#model = nn.Sequential( nn.Linear(input_dim, output_dim) ).to(device)
model = nn.Linear(input_dim, output_dim).to(device)
loss_func = nn.MSELoss()

learning_rate = 1e-4
optimizer = torch.optim.Adam( model.parameters(), lr=learning_rate )

#%%
t0 = time.time()
for epoch in count(1):#range(20000):

    # forward prediction pass
    y_pred = model(x)

    # compute loss
    loss = loss_func(y_pred, y)

    # back prop
    optimizer.zero_grad()
    loss.backward()

    if loss.item() < 1e-6 or epoch > 100000:
        break

    # update
    optimizer.step()
t1 = time.time()

#%%
print('{} epoch in {} sec'.format( epoch, t1-t0 ))
x = x.view(-1,1)
pred = model(x).data.cpu().numpy()

plt.plot( train[0], pred, 'r+')
plt.plot( train[0], truth, 'g+')

print(model.state_dict())