make_grid和save_image的用法
每次将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")