1.准备数据集
2.设计模型所用的类别,inherit from nn.Module ()
3.设置loss和optimizer,使用PyTorch API
4.设置训练迭代 forward、backward、update。
import torch#loss 必须是一个标量才可以用backward。x_data = torch.Tensor([[1.0],[2.0],[3.0]])
y_data = torch.Tensor([[2.0],[3.0],[4.0]])#构造模型类class LinearModel(torch.nn.Module):#构造函数,初始化def __init__(self):super(LinearModel,self).__init__()self.linear = torch.nn.Linear(1,1) #(1,1)表示输入和输出的维度#前向传播函数#forward()相当于对父类_init_()进行重载def forward(self, x):y_pred = self.linear(x)return y_predmodel = LinearModel()
criterion = torch.nn.MSELoss(size_average = False) #size_average指在求MSE方差时是否求均值,样本数量一样可以不求,但是样本数量不一样时就需要求,默认True
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,loss)optimizer.zero_grad()loss.backward()optimizer.step() #进行更新update#输出权重w和偏置b
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)