PyTorch中TensorBoard展示DeepLabV3图像反归一化颜色异常问题

你遇到的问题大概率是通道维度不匹配、设备不一致或数据格式不符合TensorBoard要求导致的，以下是修正方案和关键解析：

修正后的反归一化代码 #

import torch

def denormalize(img, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
    # 确保均值、标准差与输入图像在同一设备（CPU/GPU）
    mean = torch.tensor(mean, device=img.device).view(1, 3, 1, 1)
    std = torch.tensor(std, device=img.device).view(1, 3, 1, 1)
    # 执行反归一化计算
    img = img * std + mean
    # 缩放至0-255范围并截断超出值
    img = torch.clamp(img * 255, 0, 255)
    # 转换通道顺序：PyTorch默认CHW -> TensorBoard要求的HWC
    img = img.permute(0, 2, 3, 1)
    # 转换为TensorBoard兼容的uint8类型
    return img.to(torch.uint8)

关键问题解析 #

1. 设备一致性：如果输入图像在GPU上，但均值/标准差是CPU张量，会导致隐性计算错误，必须通过device=img.device对齐设备。
2. 通道顺序转换：PyTorch处理图像用(batch, channel, height, width)（CHW）格式，但TensorBoard显示需要(batch, height, width, channel)（HWC）格式，permute操作是核心修正点。
3. 数据类型要求：TensorBoard对float32类型的图像显示支持不佳，必须将0-255的浮点值转换为uint8类型才能正常渲染。

使用示例 #

处理后直接传入TensorBoard：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter()
# 假设normalized_img是经过归一化的模型输入
denorm_img = denormalize(normalized_img)
writer.add_images("Denormalized Images", denorm_img, global_step=0)
writer.close()

内容的提问来源于stack exchange，提问作者Diego Alejandro Gómez Pardo