import torch
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]]) #x与y必须是矩阵
class LinearModel(torch.nn.Module): #继承自Module,Module自动实现基本图(实现backward过程)
def __init__(self): #初始化对象默认调用的函数
super(LinearModel,self).__init__() #调用父类的构造函数
self.linear = torch.nn.Linear(1,1) #构造一个对象，包含权重和偏置,同时继承自Module
def forward(self,x): #进行前馈时所要执行的计算
y_pred = self.linear(x) #实现一个可调用的对象，实现 w * x + b
return y_pred
model = LinearModel() #可以直接被调用
criterion = torch.nn.MSELoss(size_average = False) #构造损失函数，继承自Module
optimizer = torch.optim.SGD(model.parameters(),lr = 0.01) #优化器，来自优化模块的SGD类
#训练过程
for epoch in range(1000):
y_pred = model(x_data) #前馈计算
loss = criterion(y_pred,y_data) #计算损失
print(epoch,loss.item()) #loss在print时自动调用__str__()，不会产生计算图
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)
