Save PyTorch Tensor to Image File: A Step Guide – PyTorch Tutorial

By | May 20, 2022

When we are using pytorch to process images, we may need save some image tensors to image files. In this tutorial, we will introduce you how to do.

For example, if you are using pytorch to classify mnist dataset, you may read an image to a tensor: image_feature.

In order to save this tensor to image file, we can do by these steps.

Step 1: import some libraries

In this example, we will use matplotlib to display and save a tensor.

import numpy as np
import matplotlib.pyplot as plt

In order to know how to display an image using matplotlib, you can view:

Understand matplotlib.pyplot.imshow(): Display Data as an Image – Matplotlib Tutorial

Step 2: convert pytorch tensor to numpy

Here is an example:

image_data= image_feature.data.cpu().numpy()

Here image_feature is 28*28

Convert PyTorch Tensor to NumPy: A Step Guide – PyTorch Tutorial

Step 3: use matplotlib display and save tensor to an image

After having convert a pytorch tensor to numpy ndarray, we can display and save it easily.

plt.imshow(image_data, cmap = "gray")
plt.savefig("test.png", bbox_inches = "tight", pad_inches = 0.0)

Then, we can find pytorch tensor image_feature is saved into an image file test.png

If you are using tensorflow, you can do as follows:

Matplotlib Show Images Processed by CNN Networks – Deep Learning Tutorial

Leave a Reply