如何屏蔽PyTorch Profiler的Completed Stage系列日志输出？

可以通过以下两种方式屏蔽这些日志输出：

• 调整PyTorch全局日志级别
  这些STAGE日志属于PyTorch C++层的INFO级输出，通过设置更高的日志级别即可过滤：

  import torch
  import torchvision.models as models
  from torch.profiler import profile, ProfilerActivity
  
  # 将日志级别设为ERROR，仅保留严重错误信息
  torch.set_logging_level(torch.logging.ERROR)
  
  model = models.resnet18()
  inputs = torch.randn(5, 3, 224, 224)
  
  with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
      model(inputs)
  

  可选日志级别包括DEBUG、INFO、WARNING、ERROR、CRITICAL，设置为ERROR及以上就能屏蔽STAGE类的INFO日志。

• 临时重定向输出流
  如果不想全局修改日志级别，可以仅在Profiler运行期间将标准输出/错误重定向到空设备：

  import torch
  import torchvision.models as models
  from torch.profiler import profile, ProfilerActivity
  import os
  import sys
  
  model = models.resnet18()
  inputs = torch.randn(5, 3, 224, 224)
  
  # 保存原始输出流
  original_stdout = sys.stdout
  original_stderr = sys.stderr
  
  try:
      # 重定向到空设备
      sys.stdout = open(os.devnull, 'w')
      sys.stderr = open(os.devnull, 'w')
      with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
          model(inputs)
  finally:
      # 恢复原始输出流
      sys.stdout.close()
      sys.stderr.close()
      sys.stdout = original_stdout
      sys.stderr = original_stderr
  

  这种方式仅在Profiler执行阶段屏蔽输出，不会影响其他代码的日志显示。

内容的提问来源于stack exchange，提问作者Anonymous