增加成三分类。

This commit is contained in:
zhanghuan
2025-09-10 14:11:11 +08:00
parent 068c6977c9
commit 2a3431b286
2 changed files with 77 additions and 6 deletions
+6 -6
View File
@@ -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_total = list(0. for i in range(2))
class_correct = list(0. for i in range(3))
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):