pytorch实践之线性模型使用框架

import torch


class LinearModel(torch.nn.Module):

    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)

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


# prepare dataset
x_data = torch.tensor([[1.0], [2.0], [3.0]])
y_data = torch.tensor([[2.0], [4.0], [6.0]])

model = LinearModel()
# 构造损失计算与优化器
criterion = torch.nn.MSELoss(size_average=False)
optimizer_sgd = torch.optim.SGD(model.parameters(), lr=0.01)
optimizer_adam = torch.optim.Adam(model.parameters(), lr=0.01)
optimizer_asgd = torch.optim.ASGD(model.parameters(), lr=0.01)
opt = []
opt.append(optimizer_sgd)
opt.append(optimizer_adam)
opt.append(optimizer_asgd)
pred_saver = {}
epoch_list = [i for i in range(1000)]
for optimizer in opt:
    loss_each = []
    for epoch in range(100):
        y_pred = model(x_data)
        loss = criterion(y_pred, y_data)
        print(epoch, loss.item())
        loss_each.append(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print("w", model.linear.weight.item(), "bias",model.linear.bias.item())
    # 测试模型在输入数据为4时，输出为何值
    pred_saver[optimizer] = model(torch.tensor([[4.0]])).item()

for k, v in pred_saver.items():
    print(k.__class__, v)

最后的测试情况