不更改一行AI模型的代码加速你的模型训练过程——AI模型训练加速库Nebulgym简介

标签:#AI训练##模型加速# 时间:2022/06/17 09:06:51 作者:小木

前几天初创AI企业Nebuly开源了一个AI加速库nebulgym,它最大的特点是不更改你现有AI模型的代码,但是可以将训练速度提升2倍。


一、nebulgym简介

nebulgym是一个Python库。它包含了最先进的优化技术,以帮助我们加速AI模型的训练过程。nebulgym大大减少了人工智能模型的训练时间,而不需要对训练设置进行任何修改。nebulgym优化了整个训练计算栈,从有效的数据加载,到更快的前向和后向传递,到更早的收敛等。

它的使用非常简单,如下所示:

  1. @accelerate_dataset()
  2. class MyDataset{…}
  3. @accelerate_model()
  4. class MyModel{…}
  5. #Train your model as you usually do

也就是说,只需要在你先有载入数据集或者模型类上加上nebulgym装饰器即可。其他不用任何更改即可加速。下表是其测试结果:

Hardware Not-optimized Accelerated Speedup
M1 Pro 632.05 347.52 1.8x
Intel Xeon 788.05 381.01 2.1x
AMD EPYC 1547.35 1034.37 1.5x
NVIDIA T4 258.88 127.32 2.0x
__ _ _ _

这个测试是他们自己写的一个模型,epoch为10,batch size是8,可以看到,在不同的硬件平台上都有很好的加速表现。

二、nebulgym的特点

除了这些加速外,nebulgym特点包括:

  1. 速度很快🚀:对于不同的硬件和模型,加速1.5-2倍。
  2. 使用非常简单☘️:只需要在你的代码里加上nebugym的类装饰器即可,目前支持两类加速的装饰器,一个是针对数据集载入加速的@accelerate_dataset,一个是针对模型加速的@accelerate_model
  3. 多技术支持💥:可以与任何模型、trainer或者其它训练技术相结合。
  4. 多框架支持🦾:目前仅支持PyTorch,未来预计会针对TensorFlow、Hugging Face、Jax等提供支持。
  5. 多模型支持💻:支持目前所有流行的深度学习架构,包括transformers、LSTM、CNN和FCN等。
  6. 多硬件支持🤖:这个库的目标是支持所有AI硬件,从通用的GPU、CPU到专有硬件加速器,如FPGA、ASIC等。目前仅支持CPU和GPU。

三、总结

总的来说,这个库还是很有前景的,因为它真的很简单,使用conda或者pip都可以直接安装。它目前使用的技术就是在数据加载过程中进行加速以及通过稀疏策略(sparse stategy)和高效编译技术对前向和后向传播进行加速。我看了一下,目前代码完全由python实现,代码量不多。有兴趣的童鞋可以去GitHub上瞅瞅,学习试用一下。

GitHub地址:https://github.com/nebuly-ai/nebulgym

欢迎大家关注DataLearner官方微信,接受最新的AI技术推送
相关博客
Back to Top