on
使用SLURM管理机器学习实验
搞机器学习,尤其是我们这些做深度学习的人,不可避免的要做大量的实验,很多时候做着做着就不记得之前的实验结果是怎样的了。有时候翻出来一些实验记录,但是又回想不起来,是在具体怎样的实验条件下做出来的结果,非常纠结,以至于常常要重新做对比实验。而深度学习的实验又往往动辄几天,非常的费力。
因此,我一直在寻找某种“科研管理软件”,要求也很简单:
- 记录所有的实验结果,包括训练好的模型以及控制台输出等
- 记录所有的的实验条件,方便复现实验
- 任务队列机制,用来管理待完成的实验,在设备空闲时自动开始实验
这其实都是非常简单的功能,我还曾经一度萌生过自己写一个此类软件的念头。但是,我也始终在纠结,这样的功能应该是所有科研工作者都需要啊,为什么没有一个成熟的开源或者闭源的解决方案呢?
直到有一天,我发现还有一类软件叫做集群调度软件。使用这些软件可以高效的管理一个计算机集群甚至超级计算机。因为涉及HPC领域,像我这样以前做工程的人并没有太多机会去接触,只是略有耳闻。
然而现在才发现,这不就是我一直想要的“科研管理软件”嘛。这类软件的一个基本功能就是任务调度,这就解决了我们深度学习经常会遇到的排队使用显卡的问题。任务调度可以提供一个队列,逐个完成我们预先定义好的一组任务。这对于我来说简直是一个福音。有了任务队列,我就可以把所有需要的对比实验一次性加入队列,然后愉快的出去玩啦,集群管理软件会帮我保存下所有的控制台输出。至于实验条件这块,我的方案就是配合git,在每个实验中,输出所使用代码的git hash,并在保存所使用的脚本,就可以保存下复现实验所有必需的信息啦。
下面说一下我的解决方案。
软件选型方面,我选择了SLURM(Simple Linux® Utility for Resource Management)。SLURM 是一种可用于大型计算节点集群的高度可伸缩和容错的集群管理器和作业调度系统。SLURM 维护着一个待处理工作的队列并管理此工作的整体资源利用。作为开源集群管理软件中最受欢迎的一款,SLURM提供了包括GPU调度,CPU绑定,优先级调度等非常全面的功能。虽然用SLURM这样在超级计算上使用的软件有点杀鸡用牛刀的感觉,但是鉴于他安装也不麻烦,还是非常好的一个选择。而且,SLURM本身作为集群管理软件,还提供了跨计算机调度的功能,虽然我们也就几台服务器,但是配合并行文件系统,还是可以实现对服务器非常高效率的使用的。
管理方式上,我的方法是,每个实验建立一个文件夹,里面主要包含一个slrum脚本,使用sbatch提交,所有的参数在脚本中使用#SBATCH指定,不在sbatch命令行带入任何参数,方便实验重现。脚本中输出所用代码的git hash,并将所有的控制台输出,训练好的模型等,保存在该文件夹下,简化管理。
有了这么一组配置,只要队列里有一堆实验要跑,我就可以放心的出去玩啦!