修复toAndroid.py,避免递归调用。
This commit is contained in:
+18
-6
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user