# 数据加载 · get_batch 怎么喂饭

> 逐行读 get_batch——用 np.memmap 把几百 MB 的 train.bin 当数组用而不撑爆内存，随机取一批窗口，输入 x 与目标 y 只差一个位置（这就是「预测下一个」的全部监督信号），再用 pin_memory 异步搬上 GPU。

- 作者：David（道雾轩）
- 专栏：nanoGPT 源码解读（https://daiw.org/manual/nanogpt.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/nanogpt/get-batch
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

# 数据加载 · get_batch 怎么喂饭

训练循环每一步都要一批新鲜数据。`get_batch` 就是「食堂打饭窗口」：从[数据那站](https://daiw.org/manual/nanogpt/data-prepare-char)存的 `train.bin` 里，随机舀一批喂给模型。别看它只有十来行，里头有三个漂亮的工程点。

## 全代码

`train.py` 约 116–131 行：

```python
data_dir = os.path.join('data', dataset)
def get_batch(split):
    # We recreate np.memmap every batch to avoid a memory leak, as per
    # https://stackoverflow.com/questions/45132940/...
    if split == 'train':
        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
    else:
        data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
    if device_type == 'cuda':
        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y
```

## 亮点一：memmap——不把大文件读进内存

```python
data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
```

`train.bin` 可能几百 MB 甚至几十 GB(OpenWebText)，**全读进内存不现实**。`np.memmap` 的绝活是：**把磁盘上的文件「映射」成一个数组，但并不真的把它加载进内存**——你访问哪一小段，操作系统才把哪一小段**按需换页**进来。于是你能像操作普通数组一样 `data[i:i+block_size]`，内存占用却只有实际碰到的那点。**这就是几十 GB 数据集能在有限内存里训练的关键。**

那句注释也值得一提：**每次 `get_batch` 都重新创建 `memmap`**，是为了绕开一个已知的内存泄漏（注释里附了 StackOverflow 链接）。这是踩过坑才会写的一行——真实工程的味道。

## 亮点二：随机取窗口，x 与 y 差一个位置

```python
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
```

- `ix = randint(len(data) - block_size, (batch_size,))`：随机挑 `batch_size` 个**起始位置**。注意上界是 `len(data) - block_size`，保证从任何一个起点都能取够一整窗 `block_size`。**训练数据不是顺序喂的，而是每批从整条数据里随机开窗**——这样批与批之间更独立，训练更稳。
- `x`：每个起点 `i` 取 `data[i : i+block_size]`，一段长 `block_size` 的 token，`batch_size` 段叠成一批。
- `y`：每个起点取 `data[i+1 : i+1+block_size]`——**整体往后挪一位**。

**这个「挪一位」，是整个语言模型监督信号的全部秘密。** 把 x 和 y 对齐看：

```
位置:      0    1    2    3    4
x (输入):  To   be   or   not   to
y (目标):  be   or   not   to    be     ← 每一位都是 x 的下一个
```

模型在位置 0 看到 `To`，要预测 `be`（y 的第 0 位）；在位置 1 看到 `To be`，要预测 `or`……**「预测下一个 token」这句口号，落到数据上就是「让 y = x 右移一位」**。回忆[GPT.forward 里的 cross_entropy](https://daiw.org/manual/nanogpt/gpt-assembly)，比的正是「模型对每个位置的预测」和「这里的 y」。前后完全咬合。

`.astype(np.int64)`：`train.bin` 里是 `uint16`（省空间），但 PyTorch 索引/嵌入要 `int64`，所以取出来时转一下。

## 亮点三：pin_memory + 异步搬运

```python
x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
```

数据在 CPU 上备好，要送上 GPU。这里用了两个加速小技巧：

- `pin_memory()`：把张量放进**锁页内存**（不会被操作系统换出），这样才能……
- `.to(device, non_blocking=True)`：**异步**拷贝到 GPU——不必干等拷贝完成，CPU 可以继续干别的。

配合[训练循环](https://daiw.org/manual/nanogpt/training-loop)里「一边算当前批、一边预取下一批」的写法，**数据搬运和 GPU 计算就能重叠起来**，GPU 不必停下来等饭。CPU 模式则退回普通的 `.to(device)`。

<Callout type="info">
  回头看，数据这条链特别干净：[prepare.py](https://daiw.org/manual/nanogpt/data-prepare-char) 把文本压成一条 `uint16` 长龙存盘 → `get_batch` 用 `memmap` 随机开窗、右移一位造出 (x， y) → [forward](https://daiw.org/manual/nanogpt/gpt-assembly) 拿 x 出 logits、和 y 算 loss。**没有花哨的 DataLoader，十来行搞定，却把「大文件、随机采样、异步上卡」全办了。**
</Callout>

## 饭有了，开火

配置、优化器、数据加载——训练的三大件齐了。下一章，我们终于走进 `train.py` 的**主循环**：学习率怎么随步数起落、梯度累积怎么用小显存模拟大 batch、混合精度和梯度裁剪各自在防什么，以及一步「前向—反向—更新」到底怎么写。这是第三站的高潮。👉 [训练主循环：一步到底怎么走](https://daiw.org/manual/nanogpt/training-loop)

## 参考文献

- 源码：Andrej Karpathy, nanoGPT， `train.py` 的 `get_batch`. https://github.com/karpathy/nanoGPT
- NumPy `memmap` 文档：https://numpy.org/doc/stable/reference/generated/numpy.memmap.html