使用PyTorch TensorBoard SummaryWriter从NumPy数组add_images报错KeyError: 'PNG'

问题复现 #

运行以下代码测试TensorBoard的add_images功能时，出现KeyError: 'PNG'错误：

from torch.utils.tensorboard import SummaryWriter
import numpy as np

logs_base_dir = "logs"
img_batch = np.zeros((16, 3, 100, 100))
for i in range(16): # 生成16张图像
    img_batch[i, 0] = np.arange(0, 10000).reshape(100, 100) / 10000 / 16 * i
    img_batch[i, 1] = (1 - np.arange(0, 10000).reshape(100, 100) / 10000) / 16 * i

writer = SummaryWriter(logs_base_dir)
writer.add_images('my_image_batch', img_batch, 0)
writer.close()

原因分析 #

这个错误通常是由于缺少图像编码依赖库Pillow（PIL），或者Pillow版本过低导致TensorBoard无法将图像编码为PNG格式；另外，输入数据的类型/格式不符合要求也可能引发问题。

解决方案 #

• 步骤1：安装/升级Pillow库
  打开终端执行以下命令，确保Pillow版本满足TensorBoard的要求：

  pip install --upgrade pillow
  

• 步骤2：确保输入数据格式与类型正确
  TensorBoard的add_images要求输入数据为(N, C, H, W)格式（N是图像数量，C是通道数，H/W是高/宽），且数据范围需符合规范：

  • 浮点类型数值应在[0, 1]区间内
  • 整数类型数值应在[0, 255]区间内
    可以将numpy数组转换为torch张量，提升兼容性，修改后的代码如下：

  from torch.utils.tensorboard import SummaryWriter
  import numpy as np
  import torch
  
  logs_base_dir = "logs"
  img_batch = np.zeros((16, 3, 100, 100))
  for i in range(16): # 生成16张图像
      img_batch[i, 0] = np.arange(0, 10000).reshape(100, 100) / 10000 / 16 * i
      img_batch[i, 1] = (1 - np.arange(0, 10000).reshape(100, 100) / 10000) / 16 * i
  
  # 转换为torch浮点张量
  img_batch = torch.from_numpy(img_batch).float()
  
  writer = SummaryWriter(logs_base_dir)
  writer.add_images('my_image_batch', img_batch, 0)
  writer.close()
  

• 步骤3：验证修复效果
  重新运行代码，然后启动TensorBoard查看图像：

  tensorboard --logdir=logs
  

  打开浏览器访问输出的地址，检查my_image_batch是否正常显示。

内容的提问来源于stack exchange，提问作者조해피