torch tensor保存jpg/png图片

最新推荐文章于 2025-03-09 18:17:42 发布

Andy Dennis

最新推荐文章于 2025-03-09 18:17:42 发布

阅读量2.6k

点赞数 2

文章标签：深度学习图像生成 python pytorch

本文链接：https://blog.csdn.net/weixin_43850253/article/details/130887705

版权

文章介绍了两种方法来保存训练或验证过程中的图像数据。一种是使用torchvision.utils的save_image函数，另一种涉及将张量转换为numpy数组，调整范围至0-255，再利用Pillow库进行保存。代码示例展示了如何处理批量图像并保存到指定目录。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

就是有时候需要保存一些feature map或者训练过程中验证的图像输出结果。
待会儿出去玩，先不写前言

save_img

这种方式简单粗暴

from torchvision.utils import save_image

save_image(tensor, path)

pillow.save

先转成numpy, 映射回0到255，然后再转成pillow, 最后保存图片

def save_batch_img(batch_img, output_dir, prefix=None):
    """
    img: tensor, [b, c, h, w] , [0, 1]
    """
    for i in range(batch_img.shape[0]):
        img = batch_img[i]
        img = img.permute(1, 2, 0) # shape HWC
        img = img.detach().cpu().numpy()
        img = (img * 255).astype(np.uint8)
        img = Image.fromarray(img)

        # # 保存图像
        os.makedirs(output_dir, exist_ok=True)
        if prefix is not None:
            img.save('{}/{}_{}.jpg'.format(output_dir, prefix, i))
        else:
            img.save('{}/{}.jpg'.format(output_dir, i))