Pytorch实现简单线性回归Demo

Pytorch实现简单线性回归

import numpy as np
x_values = [i for i in range(11)]
x_train = np.array(x_values,dtype=np.float32)
x_train = x_train.reshape(-1,1)
print(x_train.shape)

y_values = [2*i+1 for i in x_values]
y_train = np.array(y_values,dtype=np.float32)
y_train = y_train.reshape(-1,1)
print(y_train.shape)

import  torch
import torch.nn as nn

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

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)
model.to(device)

losses = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.01)
epochs =  1000

for epoch in range(epochs):
    epoch += 1
    inputs = torch.from_numpy(x_train).to(device)
    outputs = torch.from_numpy(y_train).to(device)

    optimizer.zero_grad()

    out = model(inputs)

    loss = losses(out,outputs)

    loss.backward()
    optimizer.step()

    if epoch % 50 == 0:
        print('epoch {},loss {}'.format(epoch,loss))


#预测
predicted =model(torch.from_numpy(x_train).requires_grad_()).data.numpy()
print(predicted)

# #保存
# torch.save(model.state_dict(),'model.pkl') #保存模型的参数 w b
# #加载
# model.load_state_dict(torch.load('model.pkl')) #加载