PyTorch实现可视化一元线性回归的Demo2（神经网络版）

一、环境

OS：Ubuntu 18.04
Environment: PyTorch&Anaconda3
Editor: Spyder

二、代码部分

比使用二乘法手工回归的版本：PyTorch实现可视化一元线性回归的Demo
有如下差别：
1、使用了神经网络
2、使用了自带的loss函数

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb 28 17:38:22 2020

@author: ftimes
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.functional as F

x = np.array([[3.3], [4.4], [5.5], [6.71], [6.93], [4.168],
                    [9.779], [6.182], [7.59], [2.167], [7.042],
                    [10.791], [5.313], [7.997], [3.1]], dtype=np.float32)

y = np.array([[1.7], [2.76], [2.09], [3.19], [1.694], [1.573],
                    [3.366], [2.596], [2.53], [1.221], [2.827],
                    [3.465], [1.65], [2.904], [1.3]], dtype=np.float32) 


plt.title('initial value')
plt.figure(1)    
plt.plot(x, y, 'bo')


#trans to tensor
x = torch.from_numpy(x)
y = torch.from_numpy(y)



class Net(torch.nn.Module):  
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()    
        self.hidden = torch.nn.Linear(n_feature, n_hidden)   
        self.predict = torch.nn.Linear(n_hidden, n_output)   
    def forward(self, x):   
        x = F.relu(self.hidden(x))      
        x = self.predict(x)             
        return x

net = Net(n_feature=1, n_hidden=10, n_output=1)

# optimizer 
learningRate=1e-2
optimizer = torch.optim.SGD(net.parameters(), learningRate)
loss_func = torch.nn.MSELoss()

print("Please input times of update: ")
times=int(input())

for t in range(times):
    prediction = net(x)
    loss = loss_func(prediction, y)   
    optimizer.zero_grad()  
    loss.backward()    
    optimizer.step()
    
plt.figure(2)
plt.title('after update')
plt.plot(x.data.numpy(), y.data.numpy(), 'bo', label='real')
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', label='estimated')
plt.legend()

三、示例

Please input times of update:

100