告别训练中断：用run_with_submitit.py实现DINO模型的Slurm任务调度-优快云博客

告别训练中断：用run_with_submitit.py实现DINO模型的Slurm任务调度

【免费下载链接】dino PyTorch code for Vision Transformers training with the Self-Supervised learning method DINO 项目地址: https://gitcode.com/gh_mirrors/di/dino

你是否曾因服务器维护、网络波动或资源限制导致Vision Transformer模型训练中断？作为AI工程师，80%的时间都耗费在任务调度和环境配置上，真正留给模型调优的时间寥寥无几。本文将带你深入解析DINO项目中的run_with_submitit.py脚本，掌握Slurm（简单Linux效用资源管理器）任务调度的核心技巧，让你的自监督训练像时钟一样精准可靠。

读完本文你将获得：

3分钟上手的DINO训练任务提交流程
解决GPU资源争用的5个高级参数配置
任务中断自动恢复的实现原理
可视化的Slurm任务生命周期管理指南

DINO训练的任务调度痛点

在分布式训练Vision Transformer模型时，我们常面临三大挑战：

资源抢占：多用户共享集群时，GPU资源随时可能被更高优先级任务抢占
训练中断：单节点故障或网络超时导致数天的训练成果丢失
环境一致性：不同节点间的依赖库版本差异引发"works on my machine"问题

run_with_submitit.py作为DINO项目的任务调度入口，通过与submitit库和Slurm调度系统深度集成，完美解决了这些痛点。其核心设计思路体现在main()函数中，采用"参数解析→执行器配置→任务提交"的三段式架构，确保训练任务在复杂集群环境中稳健运行。

从源码解析到实战：3步提交DINO训练任务

1. 参数配置：定制你的训练需求

parse_args()函数定义了10+个关键参数，让我们聚焦最常用的5个：

参数名	类型	默认值	实用场景
--ngpus	int	8	单节点GPU数量，根据集群实际配置调整
--nodes	int	2	分布式训练节点数，DINO推荐至少2节点
--timeout	int	2800	任务超时时间(分钟)，长训练建议设为1440(24小时)
--partition	str	learnfair	集群分区名称，需替换为你的集群实际分区
--use_volta32	bool	False	启用32GB大显存GPU，训练ViT-L/16以上模型必选

典型配置示例：

python run_with_submitit.py \
  --ngpus 8 \
  --nodes 4 \
  --timeout 1440 \
  --partition my_gpu_partition \
  --use_volta32 \
  --comment "DINO ViT-B/16 pretraining"

2. 执行器初始化：Slurm任务的"智能管家"

main()函数的核心是创建submitit执行器实例：

executor = submitit.AutoExecutor(folder=args.output_dir, slurm_max_num_timeout=30)

这行代码初始化了一个能自动处理Slurm任务超时的执行器，最多允许30次超时重试。关键配置在executor.update_parameters()调用中，其中：

mem_gb=40 * num_gpus_per_node：为每GPU分配40GB内存，避免显存溢出
tasks_per_node=num_gpus_per_node：确保每个GPU对应一个独立任务
slurm_signal_delay_s=120：接收终止信号后延迟120秒再退出，为 checkpoint 保存争取时间

3. 任务提交与生命周期管理

Trainer类的设计是任务稳健运行的关键。当执行executor.submit(trainer)时，发生了这些事情：

call()方法触发实际训练，调用main_dino.train_dino()
checkpoint()方法实现任务中断恢复机制，通过创建新的DelayedSubmission对象实现自动重排队
_setup_gpu_args()动态配置分布式环境，确保每个节点正确获取GPU资源

高级技巧：让你的任务"永不宕机"

分布式训练的初始化文件机制

get_init_file()函数创建了一个唯一的初始化文件：

init_file = get_shared_folder() / f"{uuid.uuid4().hex}_init"

这个UUID命名的文件作为分布式训练的临时协调点，确保所有节点正确加入同一个训练进程组。文件存储在/checkpoint/<username>/experiments目录，这是Facebook AI Research集群的标准路径，在其他集群使用时需修改get_shared_folder()函数。

任务状态监控与管理

提交任务后，通过以下命令监控训练状态：

# 查看任务队列
squeue -u $USER

# 查看任务日志
cat <output_dir>/<job_id>/log.out

# 取消任务
scancel <job_id>

run_with_submitit.py会在提交时显示日志路径：Logs and checkpoints will be saved at: <path>，建议立即记录这个路径以便后续追踪。

解决常见调度问题的5个锦囊

资源不足：添加--comment "priority job for paper deadline"获取优先级提升
显存溢出：启用--use_volta32并减少main_dino.py中的batch_size
任务频繁超时：检查utils.py中的数据加载效率，增加预加载线程数
节点间通信慢：在vision_transformer.py中启用梯度检查点
结果文件丢失：修改get_shared_folder()使用NFS共享目录

任务调度流程可视化

mermaid

上图展示了run_with_submitit.py实现的任务生命周期管理流程。关键在于Trainer类的checkpoint()方法，当Slurm发送终止信号（如抢占通知）时，会自动创建新的任务提交请求，实现训练的无缝衔接。

从新手到专家：进阶学习路径

要深入掌握DINO项目的任务调度机制，建议按以下路径学习：

基础层：理解run_with_submitit.py的参数传递流程，特别是与main_dino.py的交互
中间层：研究submitit库文档，掌握Trainer类的__call__和checkpoint协议
高级层：结合eval_linear.py和eval_image_retrieval.py，实现评估任务的自动化调度
专家层：修改visualize_attention.py，添加Slurm任务依赖，实现"训练完成后自动运行注意力可视化"的工作流

总结与最佳实践

通过run_with_submitit.py实现稳健的DINO训练任务调度，核心在于：

合理配置资源参数，避免过度请求或资源不足
利用checkpoint机制实现任务中断自动恢复
规范管理输出目录，便于追踪多个实验结果

最佳实践清单：

始终指定--output_dir参数，避免默认路径被清理
长训练任务设置--timeout 1440（24小时）并启用自动重试
提交前用--dry-run检查参数是否正确（需修改源码添加该功能）
关键实验添加--comment说明，便于集群管理员理解任务重要性

掌握这些技巧后，你将从繁琐的任务管理中解放出来，专注于提升Vision Transformer模型的性能。下一篇我们将深入探讨vision_transformer.py中的注意力机制可视化，敬请期待！

如果觉得本文对你有帮助，请点赞收藏，你的支持是我们持续产出优质内容的动力！

【免费下载链接】dino PyTorch code for Vision Transformers training with the Self-Supervised learning method DINO 项目地址: https://gitcode.com/gh_mirrors/di/dino

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考