PyTorch 神经图灵机 (NTM)

PyTorch 对 神经图灵机 (NTM) 的实现。

NTM 是一种带有外部记忆的增强型神经网络，其中与外部记忆的交互（寻址、读取、写入）通过可微分的变换来完成。总体而言，该网络是端到端可微的，因此可以使用基于梯度的优化器进行训练。

NTM 以序列的方式处理输入，类似于 LSTM，但具有额外的优势：(1) 外部记忆使网络更容易学习算法性任务；(2) 在不增加网络可训练参数的情况下，拥有更大的容量。

借助外部记忆，NTM 能够学习 LSTM 很难掌握的算法性任务，并且能够比传统 LSTM 更长时间地保持内部状态。

PyTorch 实现

本仓库以一种直观直接的方式实现了标准的 NTM。采用以下架构：

特性

• 支持批量学习
• 数值稳定
• 灵活的头部配置——可以使用 X 个读头部和 Y 个写头部，并指定操作顺序
• 复制 和 重复复制 实验结果与论文一致

---

复制任务

复制 任务用于测试 NTM 存储和回忆长序列任意信息的能力。网络的输入是一段随机的比特序列，以分隔符结尾。序列长度在 1 到 20 之间随机变化。

训练

使用 4 种不同随机种子进行的 复制 任务训练收敛情况（详情请参阅 notebook）

下图展示了训练过程中每种序列长度对应的损失。该网络使用 seed=10 进行训练，表现出快速收敛。其他种子的表现可能稍逊，但通常在 3 万次迭代内也能收敛。

评估

以下是一个动图，展示了模型的泛化能力。模型每隔 500 次训练样本进行一次评估，使用的目标序列显示在图像上方。下方则显示了网络在每个训练阶段的输出。

以下是相同任务，但序列长度为 80。请注意，该网络仅使用长度为 1 到 20 的序列进行训练。

---

重复复制任务

重复复制 任务用于测试 NTM 是否能够学习一个简单的嵌套函数，并通过学习执行一个 for 循环 来调用它。网络的输入是一段随机的比特序列，后接一个分隔符和一个标量值，表示要重复输出的次数。重复次数经过归一化处理，均值为 0，方差为 1（如论文所述）。序列长度和重复次数都在 1 到 10 之间随机化。

训练

使用 4 种不同随机种子进行的 重复复制 任务训练收敛情况（详情请参阅 notebook）

评估

下图展示了输入给网络的内容：一段比特序列 + 分隔符 + 重复次数标量。具体来说，这里的序列长度为 8，重复次数为 5。

而这是网络预测的输出：

以下是一个动图，展示了网络如何逐步学会预测目标。具体而言，网络在每次保存的检查点处都使用相同的输入序列进行了评估。

安装

NTM 目前尚未打包成可复用的模块，但可以直接作为模块使用。

1. 克隆仓库
2. 安装 PyTorch
3. 运行 pip install -r requirements.txt

使用

执行 ./train.py

用法: train.py [-h] [--seed SEED] [--task {copy,repeat-copy}] [-p PARAM]
                [--checkpoint-interval CHECKPOINT_INTERVAL]
                [--checkpoint-path CHECKPOINT_PATH]
                [--report-interval REPORT_INTERVAL]

可选参数:
  -h, --help            显示此帮助信息并退出
  --seed SEED           随机数生成器的种子值
  --task {copy,repeat-copy}
                        选择要训练的任务（默认：复制）
  -p PARAM, --param PARAM
                        覆盖模型参数。例如："--param batch_size=4
                        --param num_heads=2"
  --checkpoint-interval CHECKPOINT_INTERVAL
                        检查点间隔（默认：1000）。设置为 0 可禁用检查点保存
  --checkpoint-path CHECKPOINT_PATH
                        保存检查点数据的路径（默认：'.'）
  --report-interval REPORT_INTERVAL
                        报告间隔

PyTorch Neural Turing Machine (NTM) 快速上手指南

本指南帮助中国开发者快速部署并运行 PyTorch 版本的神经图灵机（NTM）。NTM 是一种带有外部记忆体的神经网络，擅长处理算法任务及长序列依赖问题。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：建议 Python 3.6+
• 核心依赖：
  • PyTorch (必须预先安装)
  • 其他基础库（如 numpy, matplotlib 等）

提示：国内用户安装 PyTorch 时，推荐使用清华或中科大镜像源以加速下载。 例如：pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目尚未打包发布，需通过源码克隆方式安装。

1. 克隆仓库

   git clone <repository_url>
   cd pytorch-ntm
   

   (注：请将 <repository_url> 替换为实际的仓库地址)

2. 安装依赖 建议使用国内镜像源安装 requirements.txt 中的依赖包：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

安装完成后，您可以直接运行 train.py 脚本来执行经典的 Copy（复制）或 Repeat Copy（重复复制）任务。

1. 运行默认任务 (Copy Task)

默认情况下，脚本将训练一个 NTM 模型来完成比特序列的复制任务（序列长度随机为 1 到 20）：

python train.py

2. 指定任务与参数

您可以通过命令行参数切换任务类型、设置随机种子或调整模型超参数。

示例：运行 Repeat Copy 任务

python train.py --task repeat-copy --seed 10

示例：自定义模型参数 使用 -p 或 --param 覆盖默认模型配置（例如设置批量大小为 4，读写头数量为 2）：

python train.py -pbatch_size=4 -pnum_heads=2

常用参数说明

参数	说明	默认值
--task	选择训练任务：copy 或 repeat-copy	copy
--seed	随机数生成器种子	无 (随机)
-p	覆盖模型参数，格式如 -pkey=value	无
--checkpoint-interval	保存检查点的间隔步数 (设为 0 禁用)	1000
--report-interval	打印训练报告的间隔步数	未指定

训练过程中，模型会自动保存检查点，并在控制台输出收敛情况。您可以参考项目中的 notebooks 目录查看详细的训练绘图分析。