import torch
import torch.nn as nn
import torch.optim as optim
# Define the model
model = nn.Sequential(nn.Linear(1, 1))
# Define the optimizer and loss function
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()
# Prepare data
xs = torch.tensor([[-1.0], [0.0], [1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
ys = torch.tensor([[-3.0], [-1.0], [1.0], [3.0], [5.0], [7.0]], dtype=torch.float32)
# Training loop
for _ in range(500):
optimizer.zero_grad()
outputs = model(xs)
loss = criterion(outputs, ys)
loss.backward()
optimizer.step()
print("Loss = " + str(loss.item()))
# Predict
with torch.no_grad():
prediction = model(torch.tensor([[10.0]], dtype=torch.float32))
print(prediction)这段代码使用 PyTorch 框架实现了一个最基础的机器学习任务：**线性回归**。它的目标是让模型通过学习给定的数据，找出数据背后的线性规律，并对新的输入进行预测。
下面我将从**实现原理**、**用途**和**注意事项**三个方面为您详细解释这段代码。
---
### 一、 实现原理（代码逐行解析）
代码的核心流程遵循深度学习的标准范式：**定义模型 -> 定义优化器与损失函数 -> 准备数据 -> 循环训练 -> 预测**。
1. **定义模型**：
```python
model = nn.Sequential(nn.Linear(1, 1))
```
使用 nn.Sequential 构建了一个仅包含一个线性层（全连接层）的模型nn.Linear(1, 1) 表示输入特征维度为 1，输出特征维度为 1。该层实际执行的数学运算为 $y = wx + b$，模型会自动初始化权重 $w$ 和偏置 $b$。
2. **定义优化器和损失函数**：
```python
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()
```
- criterion = nn.MSELoss()：使用均方误差（Mean Squared Error）作为损失函数，用于衡量模型预测值与真实值之间的差距。
- optimizer = optim.SGD(...)：使用随机梯度下降法作为优化器lr=0.01 是学习率，决定了每次参数更新的步长model.parameters() 告诉优化器需要更新哪些参数（即上面的 $w$ 和 $b$）。
3. **准备数据**：
```python
xs = torch.tensor([[-1.0], [0.0], [1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
ys = torch.tensor([[-3.0], [-1.0], [1.0], [3.0], [5.0], [7.0]], dtype=torch.float32)
```
构造了 6 组训练数据。观察数据不难发现，$x$ 和 $y$ 之间满足严格的线性关系：$y = 2x - 1$。
4. **训练循环**：
```python
for _ in range(500):
optimizer.zero_grad() # 1. 梯度清零：防止上一步的梯度累积到这一步
outputs = model(xs) # 2. 前向传播：将输入数据喂给模型，得到预测值
loss = criterion(outputs, ys) # 3. 计算损失：衡量预测值与真实标签的差距
loss.backward() # 4. 反向传播：计算损失函数对各个参数的梯度
optimizer.step() # 5. 参数更新：根据梯度下降算法更新 w 和 b
print("Loss = " + str(loss.item())) # 打印当前损失，观察训练过程
```
这是 PyTorch 最经典的训练步骤，循环 500 次意味着模型会在这 6 个数据上反复学习 500 遍，逐渐将 $w$ 逼近 2，$b$ 逼近 -1。
5. **预测**：
```python
with torch.no_grad():
prediction = model(torch.tensor([[10.0]], dtype=torch.float32))
print(prediction)
```
使用 with torch.no_grad(): 关闭梯度计算，因为在推理/预测阶段不需要计算梯度，这样可以节省内存并提高速度。模型输入 $x=10.0$，预期输出应接近 $2 \times 10 - 1 = 19.0$。
---
### 二、 用途
1. **教学与入门**：这是深度学习最经典的 "Hello World" 级别代码，非常适合用来理解神经网络的工作原理（前向传播、损失计算、反向传播、参数更新）。
2. **简单线性关系拟合**：在现实场景中，如果怀疑两个变量之间存在简单的正比/线性关系（如面积与房价、学习时间与分数等），可以用此方法快速拟合出线性方程。
3. **框架验证**：在搭建复杂网络之前，程序员通常会用这种最简单的数据和模型跑通流程，以验证运行环境（PyTorch安装、GPU驱动等）是否正常。
---
### 三、 注意事项
1. **数据维度**：
PyTorch 中 nn.Linear 默认的输入格式是 (batch_size, input_features)。因此，即使只有一个特征，输入数据 xs 的维度也必须是 (6, 1) 而不是一维的 (6,)。代码中使用了 [[-1.0], ...] 的嵌套列表格式是正确的。
2. **梯度清零zero_grad）**：
PyTorch 中梯度是累加的grad += new_grad）。如果在每次前向传播前不调用 optimizer.zero_grad()，当前的梯度就会加上上一次的梯度，导致参数更新方向错误。这是新手最容易遗漏的一步。
3. **学习率lr）的选择**：
代码中 lr=0.01 是一个适中的值。如果学习率过大（如 1.0），可能导致损失震荡甚至发散；如果学习率过小（如 1e-6），则收敛极其缓慢，需要增加训练轮数。
4. **过拟合与泛化能力**：
本例中，数据是完全线性的，模型在训练集上能达到极低的损失。但在真实数据中，数据往往带有噪声。如果模型过于复杂或在噪声数据上训练过久，会导致过拟合，此时对未知数据的预测能力会下降。
5. **张量数据类型**：
代码中显式指定了 dtype=torch.float32。PyTorch 模型默认的权重类型是 32 位浮点数，如果输入数据是整型（如 torch.tensor([1, 2, 3])），会报类型不匹配的错误。
6. **打印频率**：
代码在每次循环中都打印了 Loss。在训练 500 次时这没问题，但在实际工程中训练可能是几万甚至上百万步，频繁打印会严重拖慢训练速度并刷屏，通常建议每 100 步或每个 Epoch 打印一次。
