SpeedTorch

收录理由

在 PyTorch 里训练大规模 embedding 表，常常要为 GPU 显存发愁。SpeedTorch 换了个思路：把 Cupy 张量固定在 CPU 内存上，在 CPU 与 GPU 之间搬运数据时，很多情况下比 PyTorch 自带的 pinned 张量快得多，这样就能把暂时不用的参数留在系统内存里，稀疏训练时只把活跃的那部分推上 GPU。它提供的工厂类可以构建模型和优化器，在 CUDA、pinned CPU 或 Cupy 后端之间做选择，还把稀疏训练扩展到了 Adam、AdamW 这类通常只处理稠密梯度的优化器上。如果你做 NLP 或其他 embedding 密集的工作，这个库值得一试。不过要注意，项目自 2020 年起就没再更新，加速效果也取决于你的 CPU 核心数，最好先在自己的硬件上跟原生 PyTorch 跑个基准对比。