make_grid和save_image的用法

Posted 2021-02-18 hypnus-ly

每次将tensor保存为图片时，总是要先将tensor转换到cpu，然后再转换到numpy，最后再保存图片，过程太繁琐。今天介绍个pytorch的原生api，可以直接将cuda tensor 保存为图片。

import torch
import torchvision
import numpy as np
from PIL import Image

# 读取rgb格式的图片
img = Image.open("3.jpg")
# (h,w,c)--(c,h,w)，pytorch的输入格式必须为(c,h,w)
img = np.array(img).transpose((2,0,1))
# 执行了transpose后，numpy数组的内存不连续了，转换到tensor时会报错，需要先执行如下操作
img = np.ascontiguousarray([img,img,img,img])
img = torch.Tensor(img)

# 以下两句代码可以注释，save_image()函数里已经包含了make_grid()操作
# img_grid = torchvision.utils.make_grid(img)
# print(img_grid.shape)

# img如果没有归一化，必须要除以255。
torchvision.utils.save_image(img/255.0,"test.jpg")