适配MNIST设计的深度聚类网络至大尺寸RGB数据集的问题

背景 #

原深度聚类算法针对MNIST单通道28x28图像设计，现需适配416x416的3通道RGB图像，当前数据集初始化脚本如下：

class CachedMNIST(Dataset):
    def __init__(self, train, cuda, testing_mode=False):
        img_transform = transforms.Compose([transforms.Lambda(self._transformation)])
        # img_transform = transforms.Compose([transforms.Resize((28*28)), transforms.ToTensor(), transforms.Grayscale()])
        self.ds = torchvision.datasets.ImageFolder(root=train, transform=img_transform)
        self.cuda = cuda
        self.testing_mode = testing_mode
        self._cache = dict()

    @staticmethod
    def _transformation(img):
        return (torch.ByteTensor(torch.ByteStorage.from_buffer(img.tobytes())).float()
                * 0.02
        )

当前遇到的问题 #

错误1：张量形状不匹配 #

未修改图像时，_transformation输出张量形状为torch.Size([256,519168])，与自动编码器输入要求的784维度不兼容，报错：

RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x519168 and 784x500)

错误2：CUDA内存不足 #

调整图像尺寸后得到4D张量torch.Size([256,1,784,748])，即便将Batch Size调至极小，仍触发内存溢出：

RuntimeError: CUDA out of memory.

当前自动编码器结构 #

StackedDenoisingAutoEncoder(
  (encoder): Sequential(
    (0): Sequential(
      (linear): Linear(in_features=784, out_features=500, bias=True)
      (activation): ReLU()
    )
    (1): Sequential(
      (linear): Linear(in_features=500, out_features=500, bias=True)
      (activation): ReLU()
    )
    (2): Sequential(
      (linear): Linear(in_features=500, out_features=2000, bias=True)
      (activation): ReLU()
    )
    (3): Sequential(
      (linear): Linear(in_features=2000, out_features=10, bias=True)
    )
  )
  (decoder): Sequential(
    (0): Sequential(
      (linear): Linear(in_features=10, out_features=2000, bias=True)
      (activation): ReLU()
    )
    (1): Sequential(
      (linear): Linear(in_features=2000, out_features=500, bias=True)
      (activation): ReLU()
    )
    (2): Sequential(
      (linear): Linear(in_features=500, out_features=500, bias=True)
      (activation): ReLU()
    )
    (3): Sequential(
      (linear): Linear(in_features=500, out_features=784, bias=True)
    )
  )
)

解决方案方向 #

1. 图像预处理适配（低成本快速解决） #

直接将RGB图像转换为与MNIST规格一致的输入，完美匹配原自动编码器：
修改img_transform为：

img_transform = transforms.Compose([
    transforms.Resize((28, 28)),  # 缩放到28x28
    transforms.Grayscale(num_output_channels=1),  # 转为单通道灰度图
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.flatten())  # 展平为784维向量
])

同时移除自定义的_transformation函数，用PyTorch内置transform处理更规范，避免手动字节流操作导致的形状错误。

2. 模型结构改造（保留高分辨率信息） #

若需保留416x416图像的细节，需将全连接自动编码器改为卷积自动编码器，大幅降低参数数量与内存占用：

import torch.nn as nn

class ConvAutoEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        # 编码器：卷积+池化提取特征，最后映射到10维聚类特征
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*52*52, 10)  # 416经过3次2倍池化后尺寸为52x52
        )
        # 解码器：反卷积恢复图像尺寸
        self.decoder = nn.Sequential(
            nn.Linear(10, 64*52*52),
            nn.Unflatten(1, (64,52,52)),
            nn.ConvTranspose2d(64,32,2,stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(32,16,2,stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(16,3,2,stride=2),
            nn.Sigmoid()
        )
    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

3. 内存优化辅助技巧 #

• 启用梯度检查点：使用torch.utils.checkpoint.checkpoint减少训练时的中间内存占用；
• 混合精度训练：用torch.cuda.amp.GradScaler+autocast降低内存消耗，同时保证精度；
• 进一步降低Batch Size至8或4，同步调整学习率适配小批量训练。

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