训练的时候缩放,推理的时候,在模型中缩放!

This commit is contained in:
zhanghuan
2025-09-11 14:13:44 +08:00
parent ee829dc8c9
commit d9280fd0a0
2 changed files with 28 additions and 11 deletions
+21 -4
View File
@@ -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__":