在模型定义中,增加了移动端的前向传播过程。修改了toAndroid的代码。
This commit is contained in:
+2
-2
@@ -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']
|
__all__ = ['FoodCNN', 'create_food_cnn','create_mobile_food_cnn']
|
||||||
@@ -74,6 +74,29 @@ class FoodCNN(nn.Module):
|
|||||||
# 添加batch维度
|
# 添加batch维度
|
||||||
return processed.unsqueeze(0)
|
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):
|
def forward(self, x):
|
||||||
# 如果启用内部预处理且输入是tensor
|
# 如果启用内部预处理且输入是tensor
|
||||||
if self.use_internal_preprocess and isinstance(x, torch.Tensor):
|
if self.use_internal_preprocess and isinstance(x, torch.Tensor):
|
||||||
@@ -107,6 +130,22 @@ class FoodCNN(nn.Module):
|
|||||||
|
|
||||||
return 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):
|
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)
|
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__":
|
if __name__ == "__main__":
|
||||||
# 测试网络
|
# 测试网络
|
||||||
model = create_food_cnn()
|
model = create_food_cnn()
|
||||||
|
|||||||
+69
-12
@@ -1,18 +1,75 @@
|
|||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
from net import create_food_cnn, create_mobile_food_cnn
|
||||||
import torch.nn.functional as F
|
import os
|
||||||
from net import create_food_cnn
|
|
||||||
|
|
||||||
|
|
||||||
# 1. 初始化模型
|
def main():
|
||||||
model = create_food_cnn()
|
print("开始转换模型为移动端格式...")
|
||||||
# 2. 加载训练好的权重
|
|
||||||
model.load_state_dict(torch.load("../model/06/best_food_model.pth", map_location='cpu'))
|
|
||||||
model.eval() # 设置为推理模式
|
|
||||||
|
|
||||||
# 3. 创建示例输入 (假设输入是 3x224x224 的图片)
|
# 1. 加载训练好的基础模型权重
|
||||||
example_input = torch.randn(1, 3, 224, 224)
|
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
|
# 4. 转换为 TorchScript
|
||||||
traced_script_module = torch.jit.trace(model, example_input)
|
try:
|
||||||
traced_script_module.save("../model/06/best_food_model_mobile.pt")
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user