diff --git a/net/food_net.py b/net/food_net.py index cdafcdd..30bac5c 100644 --- a/net/food_net.py +++ b/net/food_net.py @@ -97,11 +97,16 @@ class FoodCNN(nn.Module): return x - def forward(self, x): - # 如果启用内部预处理且输入是tensor - if self.use_internal_preprocess and isinstance(x, torch.Tensor): - x = self.internal_preprocess(x) + def _forward_network(self, x): + """ + 网络的核心前向传播(不包含预处理) + Args: + x: 预处理后的tensor,形状为 [batch, 3, 32, 32] + + Returns: + torch.Tensor: 模型输出 + """ # 第一个卷积块 x = F.relu(self.conv1(x)) x = F.relu(self.conv2(x)) @@ -130,6 +135,13 @@ class FoodCNN(nn.Module): return x + def forward(self, x): + # 如果启用内部预处理且输入是tensor + if self.use_internal_preprocess and isinstance(x, torch.Tensor): + x = self.internal_preprocess(x) + + return self._forward_network(x) + def forward_mobile(self, x): """ 移动端前向传播(包含预处理) @@ -143,8 +155,8 @@ class FoodCNN(nn.Module): # 移动端预处理 x = self.mobile_preprocess(x) - # 标准前向传播 - return self.forward(x) + # 执行网络前向传播 + return self._forward_network(x) def create_food_cnn(use_internal_preprocess=False):