用于 PyTorch 的稀疏门控专家混合层

本仓库包含了论文《极其庞大的神经网络》（arXiv:1701.06538）中描述的稀疏门控 MoE 层的 PyTorch 重实现。

from moe import MoE
import torch

# 实例化 MoE 层
model = MoE(input_size=1000, output_size=20, num_experts=10,hidden_size=66, k= 4, noisy_gating=True)

X = torch.rand(32, 1000)

# 训练模式
model.train()
# 前向传播
y_hat, aux_loss = model(X)

# 评估模式
model.eval()
y_hat, aux_loss = model(X)

环境要求

安装依赖项请运行：

pip install -r requirements.py

示例

文件 example.py 包含一个最小可运行示例，展示了如何使用虚拟输入和目标来训练和评估 MoE 层。运行示例命令如下：

python example.py

CIFAR-10 示例

文件 cifar10_example.py 包含一个基于 CIFAR-10 数据集的最小可运行示例。在未完全收敛且采用任意超参数的情况下，该示例达到了 39% 的准确率。运行示例命令如下：

python cifar10_example.py

使用单位

FastMoE：一种快速的专家混合训练系统 本实现被用作单 GPU 训练的参考 PyTorch 实现。

致谢

该代码基于 TensorFlow 实现，TensorFlow 实现可在 这里 找到。

引用

@misc{rau2019moe,
    title={稀疏门控专家混合层的 PyTorch 实现},
    author={Rau, David},
    journal={https://github.com/davidmrau/mixture-of-experts},
    year={2019}
}

Mixture-of-Experts (MoE) 快速上手指南

本指南基于 David Rau 的 PyTorch 实现，帮助开发者快速部署稀疏门控混合专家（Sparsely Gated MoE）层。该实现参考了论文《Outrageously Large Neural Networks》。

环境准备

• 系统要求：支持 Python 3.6+ 的操作系统（Linux/macOS/Windows）。
• 核心依赖：
  • PyTorch
  • 其他依赖项详见项目根目录下的 requirements.py 文件。

提示：国内开发者建议使用清华或阿里镜像源加速 PyTorch 及相关库的安装。

安装步骤

1. 克隆或下载本项目代码到本地。
2. 进入项目目录，运行以下命令安装依赖：

pip install -r requirements.py

国内加速建议：若安装缓慢，可添加国内镜像源参数：

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

基本使用

以下是最小化的代码示例，展示如何实例化 MoE 层并进行前向传播（训练与评估模式）。

from moe import MoE
import torch

# 实例化 MoE 层
# 参数说明：
# input_size: 输入维度
# output_size: 输出维度
# num_experts: 专家数量
# hidden_size: 隐藏层维度
# k: 每次激活的专家数量 (top-k)
# noisy_gating: 是否启用噪声门控
model = MoE(input_size=1000, output_size=20, num_experts=10, hidden_size=66, k=4, noisy_gating=True)

# 构造随机输入数据 (Batch size=32, Input dim=1000)
X = torch.rand(32, 1000)

# --- 训练模式 ---
model.train()
# 前向传播，返回预测值和辅助损失 (aux_loss 用于负载均衡正则化)
y_hat, aux_loss = model(X)

# --- 评估模式 ---
model.eval()
y_hat, aux_loss = model(X)

运行示例脚本

项目提供了两个现成的示例脚本供参考：

1. 基础示例（使用虚拟数据）：

   python example.py
   

2. CIFAR-10 数据集示例：

   python cifar10_example.py
   

   注：默认超参数下该示例准确率约为 39%，仅供流程参考，需进一步调优以达到收敛。