每次將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")