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.imgsidx).convert('RGB')
if self.transform:
image = self.transform(image)
label = self.labelsidx
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,任何机器都能直接跑。