- import torch
- x_data = torch.Tensor([[1.0],[2.0],[3.0]])
- y_data = torch.Tensor([[2.0],[4.0],[6.0]])
- class MyLinear(torch.nn.Module):
- def __init__(self):
- super().__init__()
- self.linear = torch.nn.Linear(1,1)
-
- def forward(self, x):
- y_pred = self.linear(x)
- return y_pred
-
- model = MyLinear()
- criterion = torch.nn.MSELoss()
- optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
- for epoch in range(1000):
- y_pred = model(x_data)
- loss = criterion(y_pred,y_data)
- print('epoch==' + str(epoch), 'loss==' + str(loss.item()))
- 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_test=",y_test.data)
复制代码 输出结果:
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作! |