mixture of experts

Pytorch专家层的稀疏门控混合物

该存储库包含Pytorch重新实现Pytorch的纸张中稀疏门控的MOE层。

 from moe import MoE
import torch

# instantiate the MoE layer
model = MoE ( input_size = 1000 , output_size = 20 , num_experts = 10 , hidden_size = 66 , k = 4 , noisy_gating = True )

X = torch . rand ( 32 , 1000 )

#train
model . train ()
# forward
y_hat , aux_loss = model ( X )

# evaluation

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实现。

引用

 @misc{rau2019moe,
    title={Sparsely-gated Mixture-of-Experts PyTorch implementation},
    author={Rau, David},
    journal={https://github.com/davidmrau/mixture-of-experts},
    year={2019}
}