寻找可批量展示GPU中灰度张量图像的torch.imshow()类命令

当然有办法实现你想要的这种便捷批量图像查看需求！不用每次手动写一堆cpu().detach()和单张显示的代码，下面给你两种实用的方案：

方案一：自定义封装一个专属显示函数 #

这是最灵活的方式，你可以根据自己的需求定制显示的行列数、风格等。比如写一个类似torch_imshow的函数：

import matplotlib.pyplot as plt
import torch
import math

def torch_imshow(tensor, num_imgs=8, cmap='gray'):
    # 处理GPU张量，转到CPU并分离计算图
    img_tensor = tensor.cpu().detach()
    # 针对灰度图去掉通道维度（如果是RGB则不需要这一步）
    img_tensor = img_tensor.squeeze(1)
    # 确保不超过批量总数
    num_imgs = min(num_imgs, img_tensor.shape[0])
    # 计算子图的行列数，这里按每行4张排列，你可以自己调整
    rows = math.ceil(num_imgs / 4)
    cols = min(4, num_imgs)
    
    fig, axes = plt.subplots(rows, cols, figsize=(cols*3, rows*3))
    # 处理单行列的情况，避免axes是一维数组
    axes = axes.flatten() if rows > 1 or cols > 1 else [axes]
    
    for i in range(num_imgs):
        axes[i].imshow(img_tensor[i], cmap=cmap)
        axes[i].axis('off')  # 关闭坐标轴，更美观
    
    # 隐藏多余的子图（如果num_imgs不是行列的整数倍）
    for j in range(num_imgs, len(axes)):
        axes[j].axis('off')
    
    plt.tight_layout()
    plt.show()

调用的时候就像你想要的那样简单：

# 假设你的GPU张量是image，尺寸[32,1,256,256]
torch_imshow(image, 8, 'gray')

方案二：用TorchVision的官方工具make_grid #

TorchVision自带的make_grid可以把批量图像拼成一张网格图，非常适合快速预览：

import matplotlib.pyplot as plt
import torch
from torchvision.utils import make_grid

def torch_grid_show(tensor, num_imgs=8, cmap='gray'):
    # 取前num_imgs张图像，处理GPU张量
    img_slice = tensor[:num_imgs].cpu().detach()
    # 拼接成网格，normalize=True可以自动调整像素到0-1范围（如果你的张量像素不在这个范围的话）
    grid = make_grid(img_slice, nrow=4, normalize=True)
    # PyTorch张量是[C, H, W]格式，plt需要[H, W, C]，所以转置维度
    grid = grid.permute(1, 2, 0)
    plt.imshow(grid, cmap=cmap)
    plt.axis('off')
    plt.show()

调用方式同样简洁：

torch_grid_show(image, 8, 'gray')

两种方案都能满足你的需求，自定义函数更灵活可控，make_grid则更轻量化，你可以根据自己的习惯选择~

内容的提问来源于stack exchange，提问作者Adar Cohen