自定义Dataset实现ImageFolder转DataLoader偶发张量尺寸不匹配报错

为什么会出现4通道图像？ #

报错里的[4,128,128]是RGBA格式的图像，第四个通道是透明度（Alpha通道）。你的数据集里混进了少量这类格式的图片，大部分图片是常规的RGB三通道，所以才会出现偶发报错——当DataLoader随机采样的batch里刚好包含RGBA图时，就会因为通道数不匹配触发stack报错，没抽到的时候就正常运行。

解决办法 #

在自定义Dataset的__getitem__方法里，强制将所有图像转换为3通道RGB格式，不管原图是RGBA、灰度图还是其他格式。直接修改图像读取的代码：

原来的图像读取代码可能是这样：

from PIL import Image

img = Image.open(img_path)
img = img.resize((128, 128))

改成下面这样，加个.convert('RGB')：

from PIL import Image

img = Image.open(img_path).convert('RGB')  # 强制转成3通道RGB
img = img.resize((128, 128))

这样所有图像都会统一成3通道，DataLoader堆叠batch时就不会再出现尺寸不匹配的问题了。

另外也可以提前批量检查并转换数据集里的RGBA图，但代码里实时处理更稳妥，能避免后续新加入的异常格式图片引发问题。

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