风的彷徨的学习笔记

LeRobot入门教程 04：PyTorch 入门

目录：00_LeRobot入门教程_目录.md

1. 什么是 PyTorch？

PyTorch 是一个深度学习框架。它帮你完成三件重要的事情：

  1. 张量计算：张量可以理解成多维数组。图像、文本、动作序列最终都会变成数字数组。
  2. 自动求导：神经网络训练需要根据误差调整参数。PyTorch 能自动算出梯度。
  3. GPU 加速：GPU 特别擅长并行矩阵运算，PyTorch 可以把张量和模型放到 GPU 上。

张量的维度直觉

  • 一个数字：标量，0 维张量。
  • [1, 2, 3]：向量，1 维张量。
  • 一张灰度图：通常是 2 维张量，形状类似 (高, 宽)。
  • 一批 RGB 图像：通常是 4 维张量，形状类似 (batch_size, channels, height, width)。
  • 一段机器人轨迹：可能是 (batch_size, sequence_length, action_dim)。
import torch

print("PyTorch 版本:", torch.__version__)
print("CUDA 是否可用:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("GPU 名称:", torch.cuda.get_device_name(0))
    print("GPU 数量:", torch.cuda.device_count())

import torch

# 如果有 GPU 就用 GPU，否则用 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("当前设备:", device)

# 创建一个 3 行 4 列的张量
x = torch.arange(12, dtype=torch.float32).reshape(3, 4)
y = torch.ones(4, 2, device=device)

x = x.to(device)

print("x 形状:", x.shape)
print("x 内容:\n", x)
print("矩阵乘法结果形状:", (x @ y).shape)
print("矩阵乘法结果:\n", x @ y)

2. 自动求导：PyTorch 训练模型的核心

深度学习的训练过程可以简化成下面几步：

  1. 模型根据输入算出预测；
  2. 用损失函数衡量预测和真实答案差多少；
  3. 根据损失计算每个参数的梯度；
  4. 优化器根据梯度更新参数；
  5. 重复很多次。

PyTorch 的 autograd 自动完成第 3 步。你只需要定义前向计算和损失函数。

import torch

# requires_grad=True 表示希望 PyTorch 记录这个张量的梯度
w = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# 定义一个简单的损失：loss = w 的平方和
loss = (w ** 2).sum()

# 反向传播，自动计算 d(loss)/d(w)
loss.backward()

print("w:", w)
print("loss:", loss.item())
print("梯度 w.grad:", w.grad)

3. 一个最小的 PyTorch 训练例子

下面我们造一个很简单的任务：学习下面这个关系。

y = 2.0 * x1 - 1.5 * x2 + noise

模型一开始不知道这些系数，只能根据数据慢慢调整自己的参数。

这个例子很小，但结构和真实深度学习任务一样：

数据 -> 模型 -> 预测 -> 损失 -> 反向传播 -> 参数更新

import torch
from torch import nn

torch.manual_seed(0)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. 造数据
num_samples = 256
X = torch.randn(num_samples, 2, device=device)
true_w = torch.tensor([2.0, -1.5], device=device)
y = X @ true_w + 0.3 * torch.randn(num_samples, device=device)
y = y.unsqueeze(1)

# 2. 定义一个线性模型
model = nn.Linear(in_features=2, out_features=1).to(device)

# 3. 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

# 4. 训练循环
for step in range(200):
    model.train()

    prediction = model(X)
    loss = criterion(prediction, y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 20 == 0:
        print(f"step {step:03d} | loss {loss.item():.4f}")

print("训练完成")
print("模型学到的权重:", model.weight.detach().cpu().numpy())
print("模型学到的偏置:", model.bias.detach().cpu().numpy())
print("真实权重:", true_w.cpu().numpy())

4. 保存和加载模型

训练可能需要几小时甚至几天，所以一定要保存检查点。

常见做法是保存模型参数，而不是保存整个 Python 对象。这样代码结构更清晰，也更容易跨环境加载。

torch.save(model.state_dict(), "model.pth")
model.load_state_dict(torch.load("model.pth", map_location=device))

真实项目里通常还会保存优化器状态、当前 epoch、损失曲线和配置文件，这样中断后可以从检查点继续训练。

import os
import torch

save_dir = os.path.expanduser("~/LeRobot入门教程_demo")
os.makedirs(save_dir, exist_ok=True)

save_path = os.path.join(save_dir, "small_linear_model.pth")
torch.save(model.state_dict(), save_path)
print("模型已保存到:", save_path)
print("文件大小:", os.path.getsize(save_path), "bytes")

# 重新加载到同一个模型结构中
state_dict = torch.load(save_path, map_location=device)
model.load_state_dict(state_dict)
print("模型加载成功")
具身智能