基于自定义数据的带标签图像RNN分类技术求助

我完全懂你找资料的痛苦——用RNN做图像分类的示例确实大多围着MNIST转，换自定义数据集就没那么多现成参考了。下面给你几个实用的方向和代码思路，帮你把教程里的CNN顺利换成RNN：

核心思路：把图像转成RNN能处理的序列 #

RNN天生是处理序列数据的，所以第一步要把2D图像转换成序列。常见的做法有两种：

• 按行/列拆分：比如把64x64的RGB图像拆成64个长度为3*64=192的序列（每行作为一个时间步），或者拆成64个按列的序列
• 按块拆分：把图像分成若干个小方块（比如16x16），每个方块展平后作为一个时间步，适合大尺寸图像

完整代码示例（PyTorch） #

1. 自定义数据集加载 #

首先实现一个能处理自定义带标签图像的Dataset类，同时完成图像到序列的转换：

import os
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import torchvision.transforms as transforms

class CustomImageDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        # 遍历文件夹获取所有图像路径
        self.img_paths = [
            os.path.join(img_dir, f) 
            for f in os.listdir(img_dir) 
            if f.lower().endswith(('.png', '.jpg', '.jpeg'))
        ]
        # 假设标签从文件夹结构获取（比如data/cat/xxx.jpg，标签是cat）
        # 如果你是文件名带标签，就改成从文件名提取
        self.labels = [os.path.basename(os.path.dirname(p)) for p in self.img_paths]
        # 标签转索引（方便训练）
        self.label_map = {label: idx for idx, label in enumerate(sorted(set(self.labels)))}

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        image = Image.open(img_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
        
        # 把图像转成序列：(通道数, 高度, 宽度) -> (高度, 通道数*宽度)
        # 这里按行作为时间步，seq_len=高度，每个时间步的特征是一行的所有像素
        seq_image = image.permute(1, 0, 2).flatten(start_dim=1)
        label = self.label_map[self.labels[idx]]
        
        return seq_image, label

2. RNN模型定义 #

用LSTM实现一个简单的图像分类模型，和CNN的分类头逻辑类似，只是前面换成LSTM层：

import torch
import torch.nn as nn

class ImageLSTMClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes, dropout=0.2):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # LSTM层：batch_first=True表示输入形状是(batch_size, seq_len, input_size)
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0
        )
        # 分类输出层
        self.fc = nn.Linear(hidden_size, num_classes)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 初始化隐藏状态和细胞状态
        device = x.device
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        
        # LSTM前向传播，取最后一个时间步的输出
        out, _ = self.lstm(x, (h0, c0))  # out形状: (batch_size, seq_len, hidden_size)
        out = self.dropout(out[:, -1, :])  # 取最后一个时间步的特征
        out = self.fc(out)
        
        return out

3. 训练循环（和CNN逻辑一致） #

这部分和CNN的训练流程几乎一样，只是输入变成了序列格式：

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 参数设置
img_size = 64
input_size = 3 * img_size  # 每个时间步的特征数：RGB通道*图像宽度
hidden_size = 256
num_layers = 2
num_classes = len(dataset.label_map)  # 自动获取类别数
batch_size = 32
lr = 0.001
epochs = 15

# 数据变换
transform = transforms.Compose([
    transforms.Resize((img_size, img_size)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据
dataset = CustomImageDataset(img_dir='your_data_root_path', transform=transform)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

# 初始化模型、损失函数、优化器
model = ImageLSTMClassifier(input_size, hidden_size, num_layers, num_classes).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

# 训练
model.train()
for epoch in range(epochs):
    total_loss = 0.0
    correct = 0
    total = 0
    
    for seq_imgs, labels in train_loader:
        seq_imgs, labels = seq_imgs.to(device), labels.to(device)
        
        # 前向传播
        outputs = model(seq_imgs)
        loss = criterion(outputs, labels)
        
        # 反向传播+优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 统计损失和准确率
        total_loss += loss.item() * seq_imgs.size(0)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    epoch_loss = total_loss / len(dataset)
    epoch_acc = 100 * correct / total
    print(f'Epoch [{epoch+1}/{epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')

扩展与优化建议 #

• 尝试不同的序列划分方式：如果按行效果不好，可以试试按列拆分，或者把图像分成更大的块（比如16x16的块作为时间步），找到适合你数据集的方式
• 双向LSTM：把模型里的LSTM改成nn.LSTM(..., bidirectional=True)，同时把分类层的输入改成hidden_size*2，能利用前后序列的信息
• 迁移学习：如果你的数据集不大，可以先在ImageNet预训练的CNN上提取特征，再把特征序列喂给RNN，这样效果会更好
• 对比CNN的效果：RNN在图像分类上通常不如CNN高效，尤其是大尺寸图像，所以如果最终效果不好，也可以考虑用CNN+RNN的混合模型（比如CNN提取空间特征，RNN处理序列特征）

内容的提问来源于stack exchange，提问作者Albab A. Khan