pytorch训练过程中存图片的一种方式

个人随笔记录，方便使用

from einops import rearrange
from PIL import Image
x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
img = Image.fromarray(x_sample.astype(np.uint8))
img.save("test.png")

用其他的opencv或者plt应该也是可以的