使用PyTorch DataLoader加载单通道灰度图出现多余3维度问题求助

问题根源 #

该维度异常是torchvision.datasets.ImageFolder默认使用的加载器逻辑导致的：

• ImageFolder默认的PIL图像加载器会自动将单通道灰度TIFF图像转换为3通道RGB模式，你单独测试单张图像时直接调用Image.open读取的是原生单通道数据，和ImageFolder内部加载的输入数据本身就存在维度差异
• 你自定义transform里的np.array(image)拿到的ImageFolder输入已经是(2294, 1914, 3)的3通道数组，执行unsqueeze(0)后得到(1, 2294, 1914, 3)，按batch拼接后自然就多了最后一维的3通道维度

解决方案 #

你可以任选以下一种方式修复：

方案1：自定义ImageFolder的加载器，强制读取单通道 #

def gray_loader(path):
    # 强制以单通道模式打开图像
    return Image.open(path).convert('L')

train_dataset = torchvision.datasets.ImageFolder(
    root = training_dataset_path, 
    transform = training_transforms,
    loader = gray_loader
)

方案2：在transform中主动提取单通道数据 #

调整你的transform逻辑，不管输入是几通道，只保留第一个通道：

training_transforms = transforms.Compose([
    transforms.Lambda(lambda image: torch.from_numpy(np.array(image)[:,:,0] \
    .astype(np.float32)).unsqueeze(0)),
    transforms.Lambda(lambda image: image / 65535)
])

内容的提问来源于stack exchange，提问作者David Yu