Pytorch学习笔记—线性回归模型的实现

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

#输入数据的构架
x_value = [i for i in range(11)]
x_train = np.array(x_value, dtype=np.float32)
print(x_value)
print(x_train)
x_train = x_train.reshape(-1,1) #将数组reshape成1列
print(x_train)

#label的构建
y_value = [2 * i + 1 for i in x_value] #根据训练集生成正确的标签
y_train = np.array(y_value, dtype=np.float32)
y_train = y_train.reshape(-1,1)

#构建网络模型
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)

#定义超参数和损失函数
epoch = 500
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate, momentum=0)
criterion = nn.MSELoss()

#训练模型
for epoch in range(epoch):
inputs = torch.from_numpy(x_train)
labels = torch.from_numpy(y_train)
optimizer.zero_grad() #梯度清零
#前向传播
outputs = model(inputs)
#计算损失
loss = criterion(outputs,labels)
#反向传播
loss.backward()
#更新权重参数
optimizer.step()
#每50个epoch打印训练数据
if epoch%50 == 0:
print(‘epoch:{}, loss:{}’.format(epoch, loss.item()))
‘’’
若使用GPU进行训练，需进行两点改动
device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)
1.将模型导入GPU
model.to(device)
2.将数据导入GPU
inputs = torch.from_numpy(x_train).to(device)
labels = torch.from_numpy(y_train).to(device)
‘’’