トーチテンソルはjpg/png画像を保存します

序文

場合によっては、トレーニング中に検証されたいくつかの特徴マップまたは画像出力結果を保存する必要があります。
後で遊びに行ってください、前書きは書かないでください

保存_img

この方法は単純で失礼です

from torchvision.utils import save_image

save_image(tensor, path)

枕.保存

まず 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))