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__() #继承Module的父类
self.linear = torch.nn.Linear(1,1)
def forward(self , x):
y_pred = self.linear(x)
return y_pred
model = LinearModel()
#构建损失函数
criterion = torch.nn.MSELoss(reduction='sum')
#构建优化器
optimizer = torch.optim.ASGD(model.parameters(),lr = 0.01)
#训练模型
for epoch in range(1000):
y_pred = model(X_data)
loss = criterion(y_pred, Y_data)
print(epoch , loss.item())
#梯度归零
optimizer.zero_grad()
loss.backward()
#更新
optimizer.step()
#打印权重与偏置
print ('w=',model.linear.weight.item())
print ('b=',model.linear.bias.item())
训练结果
990 7.009163027760223e-08
991 6.926188689249102e-08
992 6.860869916636148e-08
993 6.778782335459255e-08
994 6.710661182296462e-08
995 6.646365591223002e-08
996 6.565574039996136e-08
997 6.512215122711495e-08
998 6.435470822907519e-08
999 6.369333505062968e-08
w= 1.999827265739441
b= 0.0003789611510001123