搭建环境
开头总是装库。PyTorch 用 pip 装就行,顺便带上 numpy 和 matplotlib:
pip install torch torchvision
pip install numpy matplotlib
线性回归:最简模型
第一个程序用线性回归,拟合带噪声的直线。生成数据、定义模型、训练、画图,一条龙。代码直给:
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
# 生成训练数据
x = torch.linspace(0, 10, 100).unsqueeze(1)
y = 2 * x + 1 + torch.randn(100, 1) * 0.5
# 定义模型
class LinearModel(nn.Module):
def __init__(self):
super(LinearModel, self).__init__()
self.linear = nn.Linear(1, 1)
def forward(self, x):
return self.linear(x)
# 创建模型实例
model = LinearModel()
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 训练模型
epochs = 100
for epoch in range(epochs):
# 前向传播
outputs = model(x)
loss = criterion(outputs, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
(epoch + ) % == :
()
torch.no_grad():
predicted = model(x)
plt.scatter(x.numpy(), y.numpy(), label=)
plt.plot(x.numpy(), predicted.numpy(), , label=)
plt.legend()
plt.show()
()

