增加成三分类。

toAndroid/toAndroid.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class FoodCNN(nn.Module):
+    def __init__(self):
+        super(FoodCNN, self).__init__()
+        # 第一个卷积块
+        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
+        self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
+        self.pool1 = nn.MaxPool2d(2, 2)
+        self.dropout1 = nn.Dropout2d(0.25)
+
+        # 第二个卷积块
+        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
+        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
+        self.pool2 = nn.MaxPool2d(2, 2)
+        self.dropout2 = nn.Dropout2d(0.25)
+
+        # 第三个卷积块
+        self.conv5 = nn.Conv2d(64, 128, 3, padding=1)
+        self.conv6 = nn.Conv2d(128, 128, 3, padding=1)
+        self.pool3 = nn.MaxPool2d(2, 2)
+        self.dropout3 = nn.Dropout2d(0.25)
+
+        # 全连接层
+        self.fc1 = nn.Linear(128 * 4 * 4, 512)
+        self.dropout4 = nn.Dropout(0.5)
+        self.fc2 = nn.Linear(512, 2)  # 2分类
+
+    def forward(self, x):
+        # 第一个卷积块
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = self.pool1(x)
+        x = self.dropout1(x)
+
+        # 第二个卷积块
+        x = F.relu(self.conv3(x))
+        x = F.relu(self.conv4(x))
+        x = self.pool2(x)
+        x = self.dropout2(x)
+
+        # 第三个卷积块
+        x = F.relu(self.conv5(x))
+        x = F.relu(self.conv6(x))
+        x = self.pool3(x)
+        x = self.dropout3(x)
+
+        # 展平
+        x = x.view(-1, 128 * 4 * 4)
+
+        # 全连接层
+        x = F.relu(self.fc1(x))
+        x = self.dropout4(x)
+        x = self.fc2(x)
+
+        return x
+
+# 1. 初始化模型
+model = FoodCNN()
+# 2. 加载训练好的权重
+model.load_state_dict(torch.load("../model/01/best_food_model.pth", map_location='cpu'))
+model.eval()  # 设置为推理模式
+
+# 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/01/best_food_model_mobile.pt")

train/food_classifier.py

@@ -18,7 +18,7 @@ plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"使用设备: {device}")
 
-# 定义CNN模型（基于CIFAR10结构，输出层改为2分类）
+# 定义CNN模型（基于CIFAR10结构，输出层改为3分类）
 class FoodCNN(nn.Module):
     def __init__(self):
         super(FoodCNN, self).__init__()
@@ -43,7 +43,7 @@ class FoodCNN(nn.Module):
         # 全连接层
         self.fc1 = nn.Linear(128 * 4 * 4, 512)
         self.dropout4 = nn.Dropout(0.5)
-        self.fc2 = nn.Linear(512, 2)  # 改为2分类
+        self.fc2 = nn.Linear(512, 3)  # 改为2分类
         
     def forward(self, x):
         # 第一个卷积块
@@ -150,8 +150,8 @@ def test(model, test_loader, device, class_names):
     model.eval()
     correct = 0
     total = 0
-    class_correct = list(0. for i in range(2))
+    class_correct = list(0. for i in range(3))
-    class_total = list(0. for i in range(2))
+    class_total = list(0. for i in range(3))
     
     with torch.no_grad():
         test_bar = tqdm(test_loader, desc='测试中')
@@ -174,7 +174,7 @@ def test(model, test_loader, device, class_names):
             })
     
     print(f'\n测试集总体准确率: {100.*correct/total:.2f}%')
-    for i in range(2):
+    for i in range(3):
         if class_total[i] > 0:
             print(f'{class_names[i]} 准确率: {100.*class_correct[i]/class_total[i]:.2f}%')
     
@@ -216,7 +216,7 @@ if __name__ == '__main__':
     val_accuracies = []
 
     best_val_acc = 0.0
-    best_model_path = '../model/01/best_food_model.pth'
+    best_model_path = '../model/02/best_food_model.pth'
 
     print("开始训练...")
     for epoch in range(num_epochs):