神经网络实现线性回归

线性回归

import torch
import matplotlib.pyplot as plt

#生成一组数据
torch.manual_seed(100)
def get_fake_data(batch_size=8):
    x=torch.rand(batch_size,1)*10
    y=x*2+3
    return x,y


#网络模型
class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)
    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred


model = LinearModel()
#损失函数
criterion = torch.nn.MSELoss(size_average=False)
#优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

#训练神经网络
def train(x_data):
    for epoch in range(1000):
        y_pred = model(x_data)
        #计算损失
        loss = criterion(y_pred, y_data)
        #梯度清0
        optimizer.zero_grad()
        #反向传播
        loss.backward()
        #参数更新
        optimizer.step()


if __name__ == '__main__':
    x_data, y_data = get_fake_data(batch_size=20)
    train(x_data)
    print('w = ', model.linear.weight.item())
    print('b = ', model.linear.bias.item())
    plt.scatter(x_data.numpy(), y_data.numpy())
    plt.show()
    plt.pause(0.5)