将分类的类别也应用在训练程序中,更好的解耦。
This commit is contained in:
+2
-1
@@ -3,6 +3,7 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision import transforms
|
||||
from PIL import Image
|
||||
from settings import settings
|
||||
|
||||
|
||||
class FoodCNN(nn.Module):
|
||||
@@ -50,7 +51,7 @@ class FoodCNN(nn.Module):
|
||||
# 全连接层
|
||||
self.fc1 = nn.Linear(128 * 4 * 4, 512)
|
||||
self.dropout4 = nn.Dropout(0.5)
|
||||
self.fc2 = nn.Linear(512, 3) # 3分类
|
||||
self.fc2 = nn.Linear(512, settings.NUM_CLASSES)
|
||||
|
||||
def preprocess_image(self, image):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user