上次手搓一个简单的线性回归,这次该上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()结果
