pytorch 曲线拟合

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import torch
import matplotlib.pyplot as plt
from torch.autograd import Variable as var


# In[2]:


def get_data(x,w,b,d):
    c,r = x.shape
    y = (w * x * x + b*x + d)+ (0.1*(5*np.random.rand(c,r)-1))
    return(y)


# In[3]:


xs = np.arange(0,5,0.01).reshape(-1,1)
ys = get_data(xs,1,-2,3)

xs = var(torch.Tensor(xs))
ys = var(torch.Tensor(ys))


# In[4]:


class Fit_model(torch.nn.Module):
    def __init__(self):
        super(Fit_model,self).__init__()
        self.linear1 = torch.nn.Linear(1,32)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(32,32)
        self.linear3 = torch.nn.Linear(32,1)
        #torch.nn.

        self.criterion = torch.nn.MSELoss()
        self.opt = torch.optim.SGD(self.parameters(),lr=0.01)
        
    def forward(self, input):
        y = self.linear1(input)
        y = self.relu(y)
        y = self.linear2(y)
        y = self.relu(y)
        y = self.linear3(y)
        return y


# In[5]:


model = Fit_model()


# In[6]:


for e in range(40001):
    y_pre = model(xs)

    loss = model.criterion(y_pre,ys)
    if(e%200==0):
        print(e,loss.data)
    
    # Zero gradients
    model.opt.zero_grad()
    # perform backward pass
    loss.backward()
    # update weights
    model.opt.step()


# In[7]:


ys_pre = model(xs)


# In[8]:


plt.title("curve")
plt.plot(xs.data.numpy(),ys.data.numpy())
plt.plot(xs.data.numpy(),ys_pre.data.numpy())
plt.legend("ys","ys_pre")
plt.show()


# In[9]:


a = 1.5


# In[14]:


ys_preq = model(torch.Tensor([a]))


# In[21]:


print(ys_preq.item())
print(np.array(ys_preq.item()))
print(ys_preq.detach().numpy())


# In[ ]: