Skip to content

Instantly share code, notes, and snippets.

@Multihuntr

Multihuntr/pytorch_tensor_to_image.py

Created September 2, 2019 03:35
Show Gist options
  • Save Multihuntr/6c9ed3938fcf2d58954bf538e521b6ea to your computer and use it in GitHub Desktop.
Save Multihuntr/6c9ed3938fcf2d58954bf538e521b6ea to your computer and use it in GitHub Desktop.
Download ZIP
One liner to save a GPU pytorch tensor to disk as an image using PIL
Raw
pytorch_tensor_to_image.py
from PIL import Image
# Assumes tensor is shaped [c, h, w]
Image.fromarray(tensor.detach().cpu().numpy().transpose([1, 2, 0])).save('test.png')
# Assumes tensor is shaped [n, c, h, w]
Image.fromarray(tensor[0].detach().cpu().numpy().transpose([1, 2, 0])).save('test.png')
# Yeah, pretty simple, but annoying to remember...
# You can use these from within the python debugger, by the way
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment