pytorch保存张量为图片

这里用到的是torchvision中的save_image。

废话不多说，直接来代码：

import torch
from torchvision.utils import save_image
B, C, H, W = 64, 3, 32, 32
input_tensor = torch.randn(B, C, H, W) 

save_image(input_tensor, "hh.png", nrow=8)

首先，这里我们要明确 ( B , C , H , W ) (B, C, H, W) (B,C,H,W)的意思，我们有64张3通道的图片需要可视化。那么，我设定每排放8张。实现效果如图：

参考链接：https://pytorch.org/vision/stable/generated/torchvision.utils.save_image.html