Dataset 三个方法 · DataLoader 组批

一、预处理:统一尺寸,转成张量

食物分类.py —— 数据预处理流水线

data_transforms = {

        'train':

            transforms.Compose([

            transforms.Resize((256,256)),

            transforms.ToTensor(),

        ]),

        'valid':

    transforms.Compose([

        transforms.Resize((256,256)),

        transforms.ToTensor(),

    ])

}

transforms.Compose 把多个变换按顺序串成流水线。Resize((256,256)) 统一缩放到 256×256(卷积网络要求输入尺寸固定);ToTensor() 把 PIL 图片转成张量——像素 0~255 归一化到 0~1,通道顺序从 H×W×C 调整为 PyTorch 要求的 C×H×W。训练集和验证集各一套,方便后面用不同策略(见第 5 篇)。

二、自定义 Dataset:模型取数据的入口

食物分类.py —— 自定义数据集类

class food_dataset(Dataset):

    def __init__(self,file_path,transform=None):

        self.file_path = file_path

        self.imgs = []

        self.labels = []

        self.transform = transform

        with open(self.file_path, encoding='utf-8') as f:

            samples = [x.strip().split(' ') for x in f.readlines()]

            for img_path,label in samples:

                self.imgs.append(img_path)

                self.labels.append(label)

    def __len__(self):

        return len(self.imgs)

    def __getitem__(self, idx):

        image = Image.open(self.imgs[idx]).convert('RGB')

        if self.transform:

            image = self.transform(image)

        label = self.labels[idx]

        label = int(label)

        return image, label

继承 PyTorch 的 Dataset,只需实现三个方法:

• __init__(构造):打开第 1 篇生成的 txt,readlines() 逐行读,strip() 去换行,split(' ') 按空格切成「路径,标签」,分别存进 self.imgs、self.labels。读文件同样写了 encoding='utf-8',与生成端一致;

• __len__:返回样本总数,DataLoader 靠它计算一个 epoch 要跑多少批;

• __getitem__:按下标取第 idx 个样本。Image.open 打开图片后紧跟 .convert('RGB') 统一转成三通道——数据集混有 PNG 和灰度图,不转的话通道数对不上卷积层会报错;接着套 transform;标签字符串用 int(label) 转整数,最终返回 (image, label)。

三、DataLoader:自动组批

食物分类.py —— 打包成批

training_data = food_dataset(file_path='food_dataset/train.txt',transform=data_transforms['train'])

test_data = food_dataset(file_path='food_dataset/test.txt',transform=data_transforms['valid'])

train_loader = DataLoader(training_data,batch_size=64,shuffle=True)

test_loader = DataLoader(test_data,batch_size=64,shuffle=True)

DataLoader 包装 Dataset,负责三件事:batch_size=64 一次取 64 张打包成批;shuffle=True 每轮随机打乱(防止模型记住数据顺序);多进程并行加载(可选)。返回对象可直接 for 迭代,一次吐出一个 batch。

卷积/BN/池化 · 全局平均池化

一、整体结构

食物分类.py —— CNN 网络定义

class CNN(nn.Module):

    def __init__(self):

        super(CNN, self).__init__()

        self.conv1 = nn.Sequential(

            nn.Conv2d(3, 16, 5, 1, 2),

            nn.BatchNorm2d(16),

            nn.ReLU(),

            nn.MaxPool2d(2),          # 256 -> 128

        )

        self.conv2 = nn.Sequential(

            nn.Conv2d(16, 32, 5, 1, 2),

            nn.BatchNorm2d(32),

            nn.ReLU(),

            nn.Conv2d(32, 32, 5, 1, 2),

            nn.BatchNorm2d(32),

            nn.ReLU(),

            nn.MaxPool2d(2),          # 128 -> 64

        )

        self.conv3 = nn.Sequential(

            nn.Conv2d(32, 128, 5, 1, 2),

            nn.BatchNorm2d(128),

            nn.ReLU(),

            nn.MaxPool2d(2),          # 64 -> 32

        )

        self.pool = nn.AdaptiveAvgPool2d(1)      # 32x32 -> 1x1

        self.out = nn.Linear(128, 20)

    def forward(self, x):

        x = self.conv1(x)

        x = self.conv2(x)

        x = self.conv3(x)

        x = self.pool(x)

        x = x.view(x.size(0), -1)

        return self.out(x)

网络分四个部分:

• conv1:16 个 5×5 卷积核 + BN + ReLU + 最大池化。Conv2d(3,16,5,1,2) 的 padding=2 让输出尺寸不变,池化把 256×256 减半成 128×128;

• conv2:两连卷积(16→32、32→32),通道翻倍,再池化,128→64;

• conv3:32→128 通道,再池化,64→32;

• pool + out:全局平均池化把 128 张 32×32 特征图压成 128×1×1,view 拉平成 128 维,送进 nn.Linear(128,20) 输出 20 个分数,一一对应 20 类。

二、重点设计:全局平均池化与参数量对比

全局平均池化 + 分类头 —— 参数量对比的关键

self.pool = nn.AdaptiveAvgPool2d(1)      # 128x32x32 -> 128x1x1

self.out = nn.Linear(128, 20)           # 参数只有 128x20+20 = 2580

如果没有全局平均池化,特征图 128×32×32=131072 维直接接全连接,仅 Linear(131072,20) 就有约 262 万个参数(项目最早版本 128×64×64 直接展开,高达约 1048 万);改成先池化再接 Linear(128,20)(约 2560 个参数),参数量骤降,训练更快、更不易过拟合。BatchNorm2d 把每批特征分布拉回零均值单位方差附近,加速收敛、稳定训练,还带一点正则化效果。

三、设备选择

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = CNN().to(device)

torch.cuda.is_available() 检测有没有可用 GPU,有就用 cuda,没有自动退回 CPU,任何机器都能直接跑。