pytorch模型组件构建简单线性回归

# 导入工具包
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torch.utils.data import random_split

# 数据集构造
x=torch.randn([10000,1])
y=x*2+0.5

#如果设备上有GPU使用GPU 否则使用CPU
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu" )
# 构建特征值 目标值
x,y=x.to(device),y.to(device)


# 构建模型
class Linearmodel(nn.Module):
    def __init__(self,input):
        super().__init__()
        self.fc=nn.Linear(1,1)
           
    def forward(self,input):
        output=self.fc(input)
        return output


# 实例化模型  损失函数选择，优化器选择 学习率选择
model=Linearmodel(x).to(device)
criterion=nn.MSELoss()
optimizer=torch.optim.SGD(model.parameters(),lr=0.05)
# 迭代训练模型
epotchs=10
for epotch in range(epotchs):
    optimizer.zero_grad()
    output=model(x)
    loss=criterion(y,output)
    loss.backward()
    optimizer.step()
    print(loss)
    
    if (epotch+1)%5==0:
        print("第{}次的损失为{}".format(epotch,loss.data))


# 模型保存与加载
torch.save(model,'./model01')
model01=torch.load('./model01')
# 模型评估
model01.eval() 
out_predict=model01(x)
out_predict=out_predict.cpu().detach().numpy()
# 真实的
plt.scatter(x.cpu().data.numpy(),y.cpu().data.numpy(),c='r')
# 预测
plt.plot(x.cpu().data.numpy(),out_predict)
plt.show()

迭代10个批次：

迭代100个批次：