Pytorch分布式训练print()使用技巧

在分布式训练场景中,有时我们可能会需要使用print函数(虽然大部分情况下大多会用logging进行信息输出)在终端打印相关信息。但由于同时运行多个进程,如果不进行限制,每个进程都会打印信息,不但影响观感,而且可能会造成阻塞。

通常的解决方法是利用if条件语句进行限制,只在主进程中进行打印,如下:

python 复制代码
# 当前为主进程
if args.rank == 0:
    print('Train message')

但最近在学习目标检测模型DINO源码时,我发现作者采用重写内置print函数 的方式实现了相同的功能,即只在主进程中启用print函数,在其他进程中禁用print函数。

函数源码如下:

python 复制代码
def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__

    # 得到内置的print函数
    builtin_print = __builtin__.print

    
    # 重写print函数
    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        # 在主进程或者强制条件下才调用内置print输出
        if is_master or force:
            builtin_print(*args, **kwargs)

    # 用重写后的print函数替换内置的print函数
    __builtin__.print = print

该方法具体的调用位置是在初始化多进程组之后,示例如下:

python 复制代码
import torch

args.rank = int(os.environ["RANK"])
args.world_size = int(os.environ['WORLD_SIZE'])
args.dist_backend = 'nccl'
args.dist_url = 'env://'
torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                         world_size=args.world_size, rank=args.rank)
# 只在主进程启用print
setup_for_distributed(args.rank == 0)

实测好用,且思路清奇,果然学习永无止境。在此做一个学习记录,也分享给需要的人。

相关推荐
Phoenixtree_DongZhao12 分钟前
ICLM 2025 Time Series 时间序列论文汇总(论文链接)
人工智能·时间序列
eve杭20 分钟前
网络安全细则[特殊字符]
大数据·人工智能·5g·网络安全
图学习的小张1 小时前
Windows安装mamba全流程(全网最稳定最成功)
人工智能·windows·深度学习·语言模型
编程指南针1 小时前
2026新选题-基于Python的老年病医疗数据分析系统的设计与实现(数据采集+可视化分析)
开发语言·python·病历分析·医疗病历分析
lisw051 小时前
数据科学与AI的未来就业前景如何?
人工智能·机器学习·软件工程
索西引擎1 小时前
AI 智能体的运行模式
人工智能·ai智能体
reasonsummer1 小时前
【办公类-117-01】20250924通义万相视频2.5——三个小人(幼儿作品动态化)
人工智能·音视频·通义万相
常州晟凯电子科技2 小时前
君正T32开发笔记之固件烧写
人工智能·笔记·嵌入式硬件·物联网
元宇宙时间2 小时前
SYN VISION韩国发布会:获评非小号Alpha,战略合作PrompTale
人工智能·web3·区块链
王哥儿聊AI2 小时前
告别人工出题!PromptCoT 2.0 让大模型自己造训练难题,7B 模型仅用合成数据碾压人工数据集效果!
人工智能·深度学习·算法·机器学习·软件工程