coding=utf-8

“”"
author:lei
function:
“”"

from torch.optim import SGD
import torch
from torch import nn

定义一个device对象 如果要转换为gpu上运行

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

准备数据

x = torch.rand([500, 1]).to(device)
y_true = 3 * x + 0.8

1、定义模型

class MyLinear(nn.Module):
def init(self):
# 继承父类的init
super(MyLinear, self).init()
self.linear = nn.Linear(1, 1)

def forward(self, x):
    out = self.linear(x)
    return out

2、实例化模型，优化器实例化，loss实例化 将模型进行gpu实例化

my_linear = MyLinear().to(device)
optimizer = SGD(my_linear.parameters(), 0.001)
loss_fn = nn.MSELoss()

3、循环，进行梯度下降，参数更新

for i in range(120000):
# 得到预测值
y_predict = my_linear(x)

loss = loss_fn(y_predict, y_true)
# 梯度置为一
optimizer.zero_grad()

# 反向传播
loss.backward()

# 参数更新
optimizer.step()

if i % 500 == 0:
    params = list(my_linear.parameters())
    print(loss.item(), params[0].item(), params[1].item())

4、模型评估 设置模型评估模式，即预测模式

my_linear.eval()