将GPU上的特征矩阵保存为图片

QT-Smile

已于 2023-03-03 21:37:31 修改

阅读量240

点赞数

分类专栏：经验分享文章标签：矩阵 python 深度学习

于 2023-03-03 21:35:21 首次发布

本文链接：https://blog.csdn.net/qq_34848334/article/details/129327715

版权

经验分享专栏收录该内容

15 篇文章 0 订阅

订阅专栏

将GPU上的特征矩阵保存为图片

import torch
from skimage import img_as_ubyte
import cv2


# filepath：保存图像的路径
# img：需要保存的图像（img在代码中是特征矩阵）
def save_img(filepath, img):
    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))


restored_preds_decoder_level4 = torch.ones([1, 3, 256, 256])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
restored_preds_decoder_level4.to(device)
# 深度学习模型一般都运行在GPU上。首先，将restored_preds_decoder_level4放在CPU上；其次，更换通道的顺序（一般深度学习模型都会更改通道顺序，当没有更改通达顺序时，这一步可以不做），并将4维压缩为3维；最后把特征矩阵变换为数组（numpy）
restored_preds_decoder_level4 = torch.clamp(restored_preds_decoder_level4, 0, 1).cpu().detach().permute(0, 2, 3, 1).squeeze(0).numpy()
result_dir_preds_decoder_level4 = 'C:\YXL\Restormer-mainV201\YXL_dir\photo/photo_preds_decoder_level4.png'
# 保存图像
save_img(result_dir_preds_decoder_level4, img_as_ubyte(restored_preds_decoder_level4))