PyTorch实现LinearModel

这里选择将模型和loss还有数据都在GPU上，方便后续能够将大型神经网络的计算图放在GPU上来加快训练速度

import torch  # 引入PyTorch模块

x = torch.tensor([[1.0], [2.0], [3.0]]).cuda()  # 创建张量x，y并使用cuda()将数据存放在GPU上使用cuda核心来进行并行计算
y = torch.tensor([[2.0], [4.0], [6.0]]).cuda()  # 在数据量和网络小时，使用CPU速度会更快


class LinearModel(torch.nn.Module):  # 创建线性模型类并继承类torch.nn.Module
    def __init__(self):  # 初始化类
        super(LinearModel, self).__init__()  # 继承父类的初始化
        self.linear = torch.nn.Linear(1, 1)  # 通过torch的线性神经网络层来创建类属性linear

    def forward(self, x):  # 自定义正向传播方式
        y_pred = self.linear(x)  # 输出经由线性层输出得到y_pred
        return y_pred  # 最后返回y_pred的值


model = LinearModel().cuda()  # 将模型的计算图放在GPU上
criterion = torch.nn.MSELoss(size_average=False).cuda()  # 为y_pred与y_true的损失定义损失函数，并不平均损失函数，同时将损失函数放在GPU上
optimizer = torch.optim.Adam(model.parameters(),
                             lr=0.01)  # 定义优化器Aadm来优化模型(model)的参数(parameters)并给定学习率learning_rate(lr)来优化线性层权重
for epoch in range(1000):  # 训练100个周期
    y_pred = model(x)  # y_pred（即预测值）由模型（model)正向传播获得
    loss = criterion(y_pred, y)  # 通过损失函数criterion来计算y_true和y_pred即真实值和预测值之间的误差$loss=\sqrt{(y-\hat{y})^2}$
    print(epoch + 1, loss)  # 打印出训练轮数和损失值
    optimizer.zero_grad()  # 将梯度初始化为0
    loss.backward()  # 通过反向传播来求的梯度
    optimizer.step()  # 优化器通过反向传播获得的梯度来更新参数
print('w=', model.linear.weight.item())
print('b=', model.linear.bias.item())
x_test = torch.tensor([[4.0]]).cuda()  # 模型的计算图和损失函数都放在GPU上，测试数据也需要放在GPU上
y_test = model(x_test)
print('y_pred=', y_test.data)

运行结果

一般来数据的device=('cuda:0')就说明数据放在了GPU的显存上
不使用cuda()函数时，默认数据都是放在内存中由CPU进行训练