pytorch将tensor转为numpy用于图像绘制

目标: 将tensor(图像形式)转为numpy，然后绘制图像

1. 形如torch.Size([1, 3, 480, 856]) # B x C x H x W

分析: 这种形状的tensor一般是要输出的某种图像或者即将转化为要输出的某种特征图

绘制方法:

1. 首先压缩tensor的维度

tensor = tensor.squeeze(0) # torch.Size([3, 480, 856])

2. 然后转换tensor维度的顺序

tensor = tensor.permute(1, 2, 0) # torch.Size([480,856,3])

3. 然后将tensor从cuda转到cpu上(未使用cuda可以不写这句话)

tensor = tensor.cpu()

4. 最后将tensor形式转化为numpy形式

array = tensor.detach().numpy()

以上4步可以合并为一句话

array = tensor.squeeze(0).permute(1,2,0).cpu().detach().numpy()

5. 使用PIL库绘制图像

from PIL import Image
im = Image.fromarray(array) # 图像全黑则尝试 im = Image.fromarray(np.uint8(array*255))
im.show()

2. 形如torch.Size([1, 32, 56, 56]) # B x C x H x W

分析: 这种形状的tensor是中间层的特征图，可以一次绘制一个通道或者写循环绘制所有通道，本次绘制一个通道

1. 首先压缩tensor的维度

tensor = tensor.squeeze(0) # torch.Size([32, 56, 56])

2. 从32个通道中选取一个通道绘制，以第一个通道为例

tensor = tensor[0, :, :]  # torch.Size([56, 56])

3. 将tensor从cuda转到cpu上(未使用cuda可以不写这句话)

tensor = tensor.cpu()

4. 将tensor形式转化为numpy形式

array = tensor.detach().numpy()

一句话

array = tensor.squeeze(0)[0, :, :].cpu().detach().numpy()

5. 绘制

from PIL import Image
im = Image.fromarray(array) # 图像全黑则尝试 im = Image.fromarray(np.uint8(array*256))
im.show()