在模型定义中,增加了移动端的前向传播过程。修改了toAndroid的代码。
This commit is contained in:
@@ -74,6 +74,29 @@ class FoodCNN(nn.Module):
|
||||
# 添加batch维度
|
||||
return processed.unsqueeze(0)
|
||||
|
||||
def mobile_preprocess(self, x):
|
||||
"""
|
||||
移动端预处理(TorchScript兼容)
|
||||
|
||||
Args:
|
||||
x: 输入tensor,形状为 [batch, 3, height, width],值范围 0-255
|
||||
|
||||
Returns:
|
||||
torch.Tensor: 预处理后的tensor,形状为 [batch, 3, 32, 32]
|
||||
"""
|
||||
# 归一化到 [0, 1]
|
||||
x = x.float() / 255.0
|
||||
|
||||
# 缩放到 32x32
|
||||
x = F.interpolate(x, size=(32, 32), mode='bilinear', align_corners=False)
|
||||
|
||||
# ImageNet标准化
|
||||
mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).view(1, 3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225], device=x.device).view(1, 3, 1, 1)
|
||||
x = (x - mean) / std
|
||||
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
# 如果启用内部预处理且输入是tensor
|
||||
if self.use_internal_preprocess and isinstance(x, torch.Tensor):
|
||||
@@ -106,6 +129,22 @@ class FoodCNN(nn.Module):
|
||||
x = self.fc2(x)
|
||||
|
||||
return x
|
||||
|
||||
def forward_mobile(self, x):
|
||||
"""
|
||||
移动端前向传播(包含预处理)
|
||||
|
||||
Args:
|
||||
x: 输入tensor,形状为 [batch, 3, height, width],值范围 0-255
|
||||
|
||||
Returns:
|
||||
torch.Tensor: 模型输出
|
||||
"""
|
||||
# 移动端预处理
|
||||
x = self.mobile_preprocess(x)
|
||||
|
||||
# 标准前向传播
|
||||
return self.forward(x)
|
||||
|
||||
|
||||
def create_food_cnn(use_internal_preprocess=False):
|
||||
@@ -121,6 +160,37 @@ def create_food_cnn(use_internal_preprocess=False):
|
||||
return FoodCNN(use_internal_preprocess=use_internal_preprocess)
|
||||
|
||||
|
||||
def create_mobile_food_cnn():
|
||||
"""
|
||||
创建移动端食物分类CNN模型
|
||||
|
||||
Returns:
|
||||
FoodCNN: 配置为移动端使用的网络模型实例
|
||||
"""
|
||||
|
||||
class MobileFoodCNN(FoodCNN):
|
||||
"""
|
||||
移动端食物分类模型
|
||||
重写forward方法以包含预处理
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__(use_internal_preprocess=False)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
移动端前向传播(自动包含预处理)
|
||||
|
||||
Args:
|
||||
x: 输入tensor,形状为 [batch, 3, height, width],值范围 0-255
|
||||
|
||||
Returns:
|
||||
torch.Tensor: 模型输出
|
||||
"""
|
||||
return self.forward_mobile(x)
|
||||
|
||||
return MobileFoodCNN()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试网络
|
||||
model = create_food_cnn()
|
||||
|
||||
Reference in New Issue
Block a user