上次手搓一个简单的线性回归,这次该上PyTorch了。

import torch  
from torch import nn

数据

还用上一次的数据:

先转成Tensor:

X = torch.tensor(df['population'].values, dtype=torch.float)  
y = torch.tensor(df['profit'].values, dtype=torch.float)  
X = X.view(len(X), 1)  
y = y.view(len(y), 1)

定义模型

class LinearRegressionModel(nn.Module):  
    def __init__(self):  
        super().__init__()  
        self.linear_layer = nn.Linear(in_features=1, out_features=1)  
      
    def forward(self, x: torch.Tensor) -> torch.Tensor:  
        return self.linear_layer(x)

然后我们需要初始化一个模型对象,指定用什么loss function和optimizer:

model = LinearRegressionModel()  
  
loss_fn = nn.L1Loss()  
  
optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01)

训练

无非是搭积木:

epochs = 5000  
  
for epoch in range(epochs):  
    model.train()  
      
    y_pred = model(X)  
    loss = loss_fn(y_pred, y)  
      
    optimizer.zero_grad()  
    loss.backward()  
    optimizer.step()

结果