From 2f6ce5ffc39a7f0920f4bebd197ea1eabf147cd8 Mon Sep 17 00:00:00 2001 From: zhanghuan <1262329256@qq.com> Date: Thu, 11 Sep 2025 14:35:46 +0800 Subject: [PATCH] =?UTF-8?q?=E5=9C=A8=E6=A8=A1=E5=9E=8B=E5=AE=9A=E4=B9=89?= =?UTF-8?q?=E4=B8=AD=EF=BC=8C=E5=A2=9E=E5=8A=A0=E4=BA=86=E7=A7=BB=E5=8A=A8?= =?UTF-8?q?=E7=AB=AF=E7=9A=84=E5=89=8D=E5=90=91=E4=BC=A0=E6=92=AD=E8=BF=87?= =?UTF-8?q?=E7=A8=8B=E3=80=82=E4=BF=AE=E6=94=B9=E4=BA=86toAndroid=E7=9A=84?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- net/__init__.py | 4 +- net/food_net.py | 70 +++++++++++++++++++++++++++++++++++ toAndroid/toAndroid.py | 83 +++++++++++++++++++++++++++++++++++------- 3 files changed, 142 insertions(+), 15 deletions(-) diff --git a/net/__init__.py b/net/__init__.py index fb4fe63..900351c 100644 --- a/net/__init__.py +++ b/net/__init__.py @@ -2,6 +2,6 @@ 食物分类网络模型包 """ -from .food_net import FoodCNN, create_food_cnn +from .food_net import FoodCNN, create_food_cnn,create_mobile_food_cnn -__all__ = ['FoodCNN', 'create_food_cnn'] \ No newline at end of file +__all__ = ['FoodCNN', 'create_food_cnn','create_mobile_food_cnn'] \ No newline at end of file diff --git a/net/food_net.py b/net/food_net.py index 5e8d221..cdafcdd 100644 --- a/net/food_net.py +++ b/net/food_net.py @@ -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() diff --git a/toAndroid/toAndroid.py b/toAndroid/toAndroid.py index f30b533..0b7485f 100644 --- a/toAndroid/toAndroid.py +++ b/toAndroid/toAndroid.py @@ -1,18 +1,75 @@ import torch -import torch.nn as nn -import torch.nn.functional as F -from net import create_food_cnn +from net import create_food_cnn, create_mobile_food_cnn +import os -# 1. 初始化模型 -model = create_food_cnn() -# 2. 加载训练好的权重 -model.load_state_dict(torch.load("../model/06/best_food_model.pth", map_location='cpu')) -model.eval() # 设置为推理模式 +def main(): + print("开始转换模型为移动端格式...") + + # 1. 加载训练好的基础模型权重 + base_model = create_food_cnn() + model_path = "../model/07/best_food_model.pth" + + if not os.path.exists(model_path): + print(f"错误：模型文件不存在 {model_path}") + return + + base_model.load_state_dict(torch.load(model_path, map_location='cpu')) + print("✓ 基础模型权重加载成功") + + # 2. 创建移动端模型并复制权重 + mobile_model = create_mobile_food_cnn() + mobile_model.load_state_dict(base_model.state_dict()) + mobile_model.eval() + print("✓ 移动端模型创建成功") + + # 3. 测试模型 + # 创建示例输入 (模拟Android端输入：1x3x224x224，值范围0-255) + example_input = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.float32) + + with torch.no_grad(): + output = mobile_model(example_input) + print(f"✓ 模型测试成功，输出形状: {output.shape}") + print(f" 输出值范围: [{output.min():.3f}, {output.max():.3f}]") + + # 4. 转换为 TorchScript + try: + traced_model = torch.jit.trace(mobile_model, example_input) + + # 保存模型 + output_path = "../model/07/best_food_model_mobile.pt" + traced_model.save(output_path) + print(f"✓ TorchScript模型保存成功: {output_path}") + + # 验证保存的模型 + loaded_model = torch.jit.load(output_path) + with torch.no_grad(): + loaded_output = loaded_model(example_input) + print(f"✓ 保存的模型验证成功，输出形状: {loaded_output.shape}") + + # 检查输出是否一致 + if torch.allclose(output, loaded_output, atol=1e-6): + print("✓ 模型输出一致性验证通过") + else: + print("⚠ 警告：模型输出存在差异") + + except Exception as e: + print(f"✗ TorchScript转换失败: {e}") + return + + print("\n" + "="*50) + print("移动端模型转换完成！") + print("="*50) + print("使用说明：") + print("1. Android端输入图片应为 [batch, 3, height, width] 格式") + print("2. 像素值范围：0-255 (uint8 或 float32)") + print("3. 模型会自动处理：") + print(" - 归一化到 [0,1]") + print(" - 缩放到 32x32") + print(" - ImageNet标准化") + print("4. 输出：3个类别的logits") + print("="*50) -# 3. 创建示例输入 (假设输入是 3x224x224 的图片) -example_input = torch.randn(1, 3, 224, 224) -# 4. 转换为 TorchScript -traced_script_module = torch.jit.trace(model, example_input) -traced_script_module.save("../model/06/best_food_model_mobile.pt") +if __name__ == "__main__": + main() \ No newline at end of file