pytorch入门 import torch.optim import torch.nn import torch x_data = torch.Tensor([[1.0], [2.0], [3.0]]) y_data = torch.Tensor([[2.0], [4.0], [6.0]]) 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 """首先实用类定义一个对象(网络结构+前向传播) 实例化对象,生成模型。实例化损失函数和优化器 开始训练for迭代,几个步骤: 1.输入数据、前向传播 2.计算损失函数 3.优化器梯度归零 4.反向传播 5.更新参数 """ model = LinearModel() # 实例化 mse_loss = torch.nn.MSELoss(size_average=False) # 实例化损失函数 optimizer = torch.optim.Adam(model.parameters(), lr=0.05,) # 实例化优化器 for epoch in range(1000): # 开始训练,前向传播,计算损失函数,优化器归零,反正传播,优化器更新参数 y_pred = model(x_data) loss = mse_loss(y_pred, y_data) print(epoch, loss) optimizer.zero_grad() loss.backward() optimizer.step() print("w", model.linear.weight.item()) print("b", model.linear.bias.item()) x_test = torch.Tensor([[4.0]]) y_test = model(x_test) print("y_pred",y_test.data)
2021-03-03
最新推荐文章于 2024-10-16 01:06:03 发布