数据准备

  1. def ready_data():
  2.     """准备数据"""
  3.     x_value = [i for i in range(11)]
  5.     x_train = np.array(x_value, dtype=np.float32)
  6.     x_train = x_train.reshape(-1, 1)
  8.     y_value = [(2 * i + 2) for i in x_value]
  10.     y_train = np.array(y_value, dtype=np.float32)
  11.     y_train = y_train.reshape(-1, 1)
  12.     return x_train, y_train

建立模型

  1. class LinearRegressionModel(nn.Module):
  2.     """建模"""
  4.     def __init__(self, input_dim, output_dim):
  5.         super().__init__()
  6.         self.linear = nn.Linear(input_dim, output_dim)
  8.     def forward(self, x):
  9.         out = self.linear(x)
  10.         return out

优化策略

  1. input_dim = 1
  2. output_dim = 1
  4. model = LinearRegressionModel(input_dim=input_dim, output_dim=output_dim)
  6. print(model)
  8. epochs = 1000
  10. learning_rate = 0.01
  12. optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
  13. criterion = nn.MSELoss()
  14. x_train, y_train = ready_data()
  16. for epoch in range(epochs):
  17.     inputs = torch.from_numpy(x_train)
  18.     labels = torch.from_numpy(y_train)
  19.     # 将梯度清零
  20.     optimizer.zero_grad()
  21.     # 前向传播
  22.     outputs = model(inputs)
  23.     # 计算损失
  24.     loss = criterion(outputs, labels)
  25.     # 反向传播
  26.     loss.backward()
  27.     # 更新参数
  28.     optimizer.step()
  30.     if epoch % 50 == 0:
  31.         print('epoch {}, loss {}'.format(epoch, loss.item()))

梯度清零

  1. optimizer.zero_grad()

前向传播

  1. outputs = model(inputs)

计算损失

  1. loss = criterion(outputs, labels)

反向传播

  1. loss.backward()

更新参数

  1.  optimizer.step()

模型保存

  1. torch.save(model.state_dict(), 'model.pkl')

模型读取

  1. model.load_state_dict(torch.load('model.pkl'))