目录
在这篇博客中，我们将展示如何使用PyTorch库来构建一个卷积神经网络（CNN）模型，用于食物图像的分类。我们将从头开始，包括数据准备、模型设计、训练过程、模型评估以及使用训练好的模型进行图像预测。
一、数据准备
首先，我们需要准备训练数据和测试数据。这里我们有两个文本文件train.txt和test.txt，分别包含训练集和测试集的图像路径和标签。每个文件中的每一行都包含一个图像路径和一个标签，用空格分隔。（本数据由合肥兰智数加学院李挺老师特别提供，提取码lzsj）
为了加载这些数据，我们定义了一个Dataset类food_dataset，它继承自torch.utils.data.Dataset。这个类负责从文本文件中读取图像路径和标签，并在请求时加载和预处理图像。
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
from torchvision import transforms
# 定义数据转换
data_transforms = {
'train':
transforms.Compose([
transforms.Resize([256,256]), # 调整图像大小为256x256像素
transforms.ToTensor(), # 将图像转换为Tensor格式
]),
'valid':
transforms.Compose([
transforms.Resize([256,256]), # 调整图像大小为256x256像素
transforms.ToTensor(), # 将图像转换为Tensor格式
]),
}
class food_dataset(Dataset):
def __init__(self, file_path, transform = None):
self.file_path = file_path
self.imgs = []
self.labels = []
self.transform = transform
with open(self.file_path) as f:
samples = [x.strip().split(' ') for x in f.readlines()] # 读取文件中的样本数据
for img_path, label in samples:
self.imgs.append(img_path) # 将图像路径添加到列表中
self.labels.append(label) # 将标签添加到列表中
def __len__(self):
return len(self.imgs) # 返回数据集的长度
def __getitem__(self, idx):
image = Image.open(self.imgs[idx]) # 打开图像文件
if self.transform:
image = self.transform(image) # 应用数据转换
label = self.labels[idx] # 获取图像对应的标签
label = torch.from_numpy(np.array(label, dtype=np.int64)) # 将标签转换为Tensor格式
return image, label
# 创建训练集和测试集的实例
training_data = food_dataset(file_path = 'train.txt', transform = data_transforms['train'])
test_data = food_dataset(file_path = 'test.txt', transform = data_transforms['valid'])
# 创建数据加载器
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True) # 训练集数据加载器
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True) # 测试集数据加载器
二、模型设计
接下来，我们设计一个适用于食物图像分类的CNN模型。这个模型包含几个卷积层、ReLU激活层和最大池化层，最后是一个全连接层用于输出分类结果。
import torch.nn as nn
# 检测GPU是否可用，优先选择cuda，然后mps，否则选择cpu
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
# 定义卷积神经网络模型
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
# 第一个卷积层
self.conv1 = nn.Sequential(
nn.Conv2d(
in_channels = 3, # 输入通道数为3（彩色图像）
out_channels = 16, # 输出通道数为16
kernel_size = 5, # 卷积核大小为5x5
stride=1, # 步长为1
padding=2, # 填充为2
),
nn.ReLU(), # ReLU激活函数
nn.MaxPool2d(kernel_size = 2), # 最大池化层，池化核大小为2x2
)
# 第二个卷积层
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2), # 输入通道数16，输出通道数32，卷积核大小5x5，填充2
nn.ReLU(), # ReLU激活函数
nn.Conv2d(32, 32, 5, 1, 2), # 输入通道数32，输出通道数32，卷积核大小5x5，填充2
nn.ReLU(), # ReLU激活函数
nn.MaxPool2d(2), # 最大池化层，池化核大小为2x2
)
# 第三个卷积层
self.conv3 = nn.Sequential(
nn.Conv2d(32, 128, 5, 1, 2), # 输入通道数32，输出通道数128，卷积核大小5x5，填充2
nn.ReLU(), # ReLU激活函数
)
# 全连接层
self.out = nn.Linear(128 * 64 * 64, 20) # 输入大小为128*64*64，输出大小为20（类别数）
def forward(self, x):
x = self.conv1(x) # 第一层卷积
x = self.conv2(x) # 第二层卷积
x = self.conv3(x) # 第三层卷积
x = x.view(x.size(0), -1) # 展平多维的卷积图像
output = self.out(x) # 输出层
return output
# 创建模型实例，并移动到对应的设备（GPU或CPU）
model = CNN().to(device)
print(model) # 打印模型结构
三、训练过程
在训练过程中，我们使用定义的train函数来迭代训练数据，更新模型权重。这个函数计算每个批次的损失，并使用反向传播来优化模型。
# 定义训练函数
def train(dataloader, model, loss_fn, optimizer):
model.train() # 将模型设为训练模式
batch_size_num = 1 # 记录批次号，用于打印输出
for X, y in dataloader:
X, y = X.to(device), y.to(device) # 将数据移动到设备（GPU或CPU）
pred = model.forward(X) # 前向传播得到预测值
loss = loss_fn(pred, y) # 计算损失
optimizer.zero_grad() # 梯度清零
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新参数
loss_value = loss.item() # 获取损失值
print(f"loss: {loss_value:>7f} [number:{batch_size_num}]") # 打印损失值
batch_size_num += 1 # 更新批次号
# 定义测试函数
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset) # 数据集大小
num_batches = len(dataloader) # 数据批次数
model.eval() # 将模型设为评估模式，不更新梯度
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device) # 将数据移动到设备
pred = model.forward(X) # 前向传播得到预测值
test_loss += loss_fn(pred, y).item() # 计算测试集损失
correct += (pred.argmax(1) == y).type(torch.float).sum().item() # 计算正确预测数量
test_loss /= num_batches # 平均每个批次的损失
correct /= size # 计算准确率
print(f"Test result: \n Accuracy: {(100*correct)}%，Avg loss: {test_loss}") # 打印测试结果
loss_fn = nn.CrossEntropyLoss() # 损失函数为交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=0.000175) # Adam优化器，学习率为0.000175
# 训练模型一次
train(train_dataloader, model, loss_fn, optimizer)
# 在测试集上评估模型表现
test(test_dataloader, model, loss_fn)
epochs = 30 # 定义总的训练轮数
for t in range(epochs):
print(f"Epoch {t+1}\n--------------------------")
train(train_dataloader, model, loss_fn, optimizer) # 训练模型
print("Done!")
test(test_dataloader, model, loss_fn) # 最终在测试集上评估模型性能
五、模型保存与加载
训练好的模型需要被保存下来以便将来使用。我们可以使用torch.save函数来保存模型的状态字典，并在需要时加载它。
model_path = 'fooddataset2_cnn_model.pth' # 定义模型保存路径和文件名
torch.save(model, model_path) # 使用PyTorch的torch.save函数保存模型到指定路径
# 注释: 将训练好的CNN模型保存到文件'fooddataset2_cnn_model.pth'
六、使用模型进行预测
最后，定义一个函数来使用训练好的模型对新的图像进行预测。这个函数加载模型，对图像进行预处理，并通过模型进行前向传播来获取预测结果。
1、构建CNN模型
在PyTorch中，模型通常是通过继承nn.Module类来构建的。下面是一个简单的CNN模型处理大小为(3, 256, 256)的RGB图像，并输出20个类别的预测结果。
import torch
from PIL import Image
from torchvision import transforms
import torch.nn as nn
class CNN(nn.Module):#类的名称。
def __init__(self): #输入大小(3, 256, 256)
super(CNN, self).__init__()
self.conv1 = nn.Sequential( # 将多个层组合成一起。
nn.Conv2d( #2d一般用于图像，3d用于视频数据(多一个时间维度)，1d -一般用于结构化的序列数据
in_channels = 3, # 图像通道个数，1表示灰度图(确定了卷积核组中的个数)
out_channels = 16, # 要得到多少个特征图，卷积核的个数
kernel_size = 5, # 卷积核大小, 5*5
stride=1, # 步长
padding=2, # 参数5，1，2可以使卷积核处理后的结果大小与处理前的数据大小相同
),
nn.ReLU(), # relu层，不会改变特征图的大小
nn.MaxPool2d(kernel_size = 2), # 进行池化操作(2x2区域)
)
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2),
nn.ReLU(), # relu层
nn.Conv2d(32, 32, 5, 1, 2),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.conv3 = nn.Sequential( # 输入 (32，64，64)
nn.Conv2d(32, 128, 5, 1, 2),
nn.ReLU(),
)
self.out = nn.Linear(128 * 64 * 64, 20) # 全连接层得到的结果
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1)
output = self.out(x)
return output2、模型加载与预测
上文中训练好的模型保存为.pth的文件，现在加载模型并使用它对单张图像进行预测。
# 定义一个函数来进行预测
def predict_image(image_path, model_path='fooddataset2_cnn_model.pth'):
# 加载模型
model = torch.load(model_path, map_location=torch.device('cpu')) # 加载模型并确保在CPU上运行
model.eval() # 将模型设为评估模式，不更新梯度
# 图像预处理
preprocess = transforms.Compose([
transforms.Resize([256, 256]), # 调整图像大小为256x256像素
transforms.ToTensor(), # 将图像转换为Tensor格式
])
# 读取图像
image = Image.open(image_path) # 使用PIL库打开图像
# 预处理图像
image_tensor = preprocess(image) # 对图像进行预处理
image_tensor = image_tensor.unsqueeze_(0) # 添加一个维度作为batch维度
# 使用模型进行预测
with torch.no_grad(): # 在预测阶段不需要计算梯度
output = model(image_tensor) # 将预处理后的图像输入模型进行预测
# 获取预测类别
_, predicted_idx = torch.max(output, 1) # 获取预测结果中概率最高的类别索引
food_classes = {
0: "哈密瓜", 1: "圣女果", 2: "巴旦木", 3: "坚果", 4: "汉堡",
5: "火龙果", 6: "炸鸡", 7: "瓜子", 8: "澳洲牛肉", 9: "白萝卜",
10: "胡萝卜", 11: "草莓", 12: "菠萝", 13: "薯条", 14: "鸡蛋",
15: "蛋挞", 16: "菠菜", 17: "骨肉相连", 18: "鸡翅", 19: "八宝粥"
}
# 获取食物类别名称
predicted_label = predicted_idx.item()
food_name = food_classes[predicted_label]
return food_name
# 使用函数预测图像
image_path = r'C:\Users\tonyzhu\Desktop\哈密瓜.jpg' # 替换为你的图像路径
predicted_label = predict_image(image_path) # 调用预测函数获取预测标签
print(f"The predicted label is: {predicted_label}") # 打印预测标签
# 注释: 定义了一个函数来预测输入图像的类别，包括模型加载、图像预处理和预测过程。
3、结果展示
下课！！！
1654
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
