diff --git a/net/food_net.py b/net/food_net.py index d5ebb46..5e8d221 100644 --- a/net/food_net.py +++ b/net/food_net.py @@ -9,16 +9,26 @@ class FoodCNN(nn.Module): """ 食物分类CNN模型 基于CIFAR10结构,适配3分类任务 + use_internal_preprocess 是否在模型内部预处理 + 训练的时候, """ - def __init__(self): + def __init__(self, use_internal_preprocess=False): super(FoodCNN, self).__init__() - # 图片预处理变换 + self.use_internal_preprocess = use_internal_preprocess + + # 图片预处理变换(仅在推理时使用) self.preprocess = transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) + # 内部预处理变换(用于已经是tensor但未归一化的数据) + self.internal_preprocess = transforms.Compose([ + # 不包含Resize,因为训练时已经在DataLoader中处理了 + transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) + ]) + # 第一个卷积块 self.conv1 = nn.Conv2d(3, 32, 3, padding=1) self.conv2 = nn.Conv2d(32, 32, 3, padding=1) @@ -65,6 +75,10 @@ class FoodCNN(nn.Module): return processed.unsqueeze(0) def forward(self, x): + # 如果启用内部预处理且输入是tensor + if self.use_internal_preprocess and isinstance(x, torch.Tensor): + x = self.internal_preprocess(x) + # 第一个卷积块 x = F.relu(self.conv1(x)) x = F.relu(self.conv2(x)) @@ -94,14 +108,17 @@ class FoodCNN(nn.Module): return x -def create_food_cnn(): +def create_food_cnn(use_internal_preprocess=False): """ 创建食物分类CNN模型 + Args: + use_internal_preprocess: 是否在forward中进行预处理 + Returns: FoodCNN: 网络模型实例 """ - return FoodCNN() + return FoodCNN(use_internal_preprocess=use_internal_preprocess) if __name__ == "__main__": diff --git a/train/train_food_classifier.py b/train/train_food_classifier.py index 2f99be3..cd9504a 100644 --- a/train/train_food_classifier.py +++ b/train/train_food_classifier.py @@ -29,20 +29,20 @@ print(f"使用设备: {device}") -# 数据预处理 +# 数据预处理 - 必须包含Resize以保证batch中tensor尺寸一致,只有Normalize由模型内部完成 transform_train = transforms.Compose([ - transforms.Resize((32, 32)), + transforms.Resize((32, 32)), # 必须保留,确保batch中tensor尺寸一致 transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(10), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1), transforms.ToTensor(), - transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) + # 注意:只有Normalize由模型内部处理 ]) transform_test = transforms.Compose([ - transforms.Resize((32, 32)), + transforms.Resize((32, 32)), # 必须保留,确保batch中tensor尺寸一致 transforms.ToTensor(), - transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) + # 注意:只有Normalize由模型内部处理 ]) # 训练函数 @@ -153,8 +153,8 @@ if __name__ == '__main__': print(f"验证集大小: {len(val_dataset)}") print(f"测试集大小: {len(test_dataset)}") - # 创建模型 - model = create_food_cnn().to(device) + # 创建模型 - 启用内部预处理 + model = create_food_cnn(use_internal_preprocess=True).to(device) print(f"模型参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}") # 定义损失函数和优化器