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")

 

posted @ 2020-05-20 11:29  blackx  阅读(5724)  评论(0编辑  收藏  举报