torch实现简单的线性回归模型

描述

拟合y=2*x+1（其中w =2,b =1）
最后模型求得的参数w近似等于2，b近似等于1

代码

import numpy as np
import torch
import torch.nn as nn

#拟合y=x*2+1
#创造训练数据集x 和 y
x = [i for i in range(11)]
x_train = np.array(x,dtype=np.float32)
x_train = x_train.reshape(-1,1)
# print(x_train)
# print(x_train.shape)

y = [i*2+1 for i in range(11)]
y_train = np.array(y,dtype=np.float32)
y_train = y_train.reshape(-1,1)
# print(y_train)
# print(y_train.shape)

#构建模型
class LinearRegressionModel(nn.Module):
    def __init__(self,input_dim,output_dim):
        super(LinearRegressionModel,self).__init__()
        self.linear = nn.Linear(input_dim,output_dim)
    #前向传播 全连接层
    def forward(self,x):
        out = self.linear(x)
        return out

input_dim = 1
output_dim = 1
model = LinearRegressionModel(input_dim,output_dim)
# print(model)

#设置训练模型的设备gpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

#设置迭代次数和学习率 优化器 和 损失函数
epochs = 1000
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate)
criterion = nn.MSELoss()

for epoch in range(epochs):
    epoch += 1
    #数据
    inputs = torch.from_numpy(x_train).to(device)
    labels = torch.from_numpy(y_train).to(device)
    #梯度清零
    optimizer.zero_grad()
    #前向传播
    outputs = model(inputs)
    #计算损失
    loss = criterion(outputs,labels)
    #反向传播
    loss.backward()
    #更新权重参数
    optimizer.step()
    if epoch % 50 == 0:
        print('epoch {},loss {}'.format(epoch,loss.item()))
#预测数据 numpy只支持cpu
predicted = model(torch.from_numpy(x_train).requires_grad_().to("cpu")).data.numpy()
print(predicted)
#保存模型参数
torch.save(model.state_dict,'model.pkl')
#model.load_state_dict(torch.load('model.pkl'))

参考文献

https://www.bilibili.com/video/BV1s3411F7c5?p=8&spm_id_from=pageDriver&vd_source=7cd34e51498fa7a003b2f6546bcf0ecc