目标：
对文件夹中的食物图片使用卷积神经网络训练实现分类。
代码思路：
1、创建文件夹写入图片地址。2、定义处理原始图片获得张量的函数。3、定义卷积神经网络。4、定义测试函数与训练函数。5、选择损失函数和优化器。6、训练与测试。
代码设计：
1、创建文件夹写入图片地址：
python
# 任务: 给图片自动生成train.txt,以及test.txt。对应的标签,就是用os默认的
import os #
def train_test_file(root,dir):
file_txt = open(dir+'.txt','w')
path = os.path.join(root,dir)#拼接路径
for roots,directories, files in os.walk(path):#os.list_dir()
if len(directories) != 0 :
dirs = directories
else:
now_dir = roots.split('\\')
for file in files:
path_1 = os.path.join(roots,file)
print(path_1)
file_txt.write(path_1+' '+str(dirs.index(now_dir[-1]))+'\n')
#str(dirs.index(now_dir[-1]))作为图片的分类标签
file_txt.close()
root = r'.\食物分类\food_dataset'
train_dir = 'train'
test_dir = 'test'
train_test_file(root,train_dir)
train_test_file(root,test_dir)
# 备注: 文件也可以是TXT、csv、json等格式os.walk():
是Python 内置os模块的核心函数，作用是从指定的path路径开始，递归地遍历该路径下的所有目录和文件。
遍历顺序：
split:
split(分隔符) 是 Python 字符串的内置函数，作用是：把一个字符串按指定的 "分隔符" 拆分成列表 。在代码中的效果：
执行结果：
2、把图片和标签转变为张量
代码：
python
'''创建数据集的类'''
import torch
from torch.utils.data import Dataset,DataLoader #用于处理数据集的
import numpy as np
from PIL import Image #
from torchvision import transforms #对数据进行处理工具 转换
data_transforms = { #机器学习的时候, 数据进行归一化？？ 图片做归一化 ->0~1
'train':
transforms.Compose([
transforms.Resize([300,300]), #是图像变换大小 opencv int8 0~
transforms.RandomRotation(45),#随机旋转, -45到45度之间随机选
transforms.CenterCrop(256),#从中心开始裁剪[256,256]
transforms.RandomHorizontalFlip(p=0.5),#随机水平翻转 选择一个概率概率
transforms.RandomVerticalFlip(p=0.5),#随机垂直翻转
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),#参数1为亮度, 参数2为对比度, 参数3为饱和度, 参
transforms.RandomGrayscale(p=0.1),#概率转换成灰度率, 3通道就是R=G=B
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])#标准化, 均值, 标准差,
]),
'valid':#验证集 不需要对图像进行数据增强
transforms.Compose([
transforms.Resize([256,256]), #
transforms.ToTensor(),
transforms.Normalize( [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
}#数组增强
class food_dataset(Dataset): # food_dataset是自己创建的类名称，可以改为你需要的名称 2用法
def __init__(self, file_path,transform=None): #类的初始化,解析数据文件txt
self.file_path = file_path
self.imgs = []
self.labels = []
self.transform = transform
with open(self.file_path) as f:#是把train.txt文件中图片的路径保存在 self.imgs，train.txt文件中标签保存在 self.labels
samples = [x.strip().split(' ') for x in f.readlines()]
for img_path, label in samples:
self.imgs.append(img_path) #图像的路径
self.labels.append(label) #标签,还不是tensor
#初始化：把图片目录加载到self.imgs.
def __len__(self): #类实例化对象后，可以使用len函数测量对象的个数 ls=[12,3,4,4] len(training_data)
return len(self.imgs)
#training_data[1]
def __getitem__(self, idx): #关键，可通过索引的形式获取每一个图片数据及标签
image = Image.open(self.imgs[idx]) #读取到图片数据，还不是tensor,BGR
if self.transform: #将PIL图像数据转换为tensor
image = self.transform(image) #图像处理为256×256，转换为tenor
label = self.labels[idx] #label还不是tensor
label = torch.from_numpy(np.array(label,dtype = np.int64)) #label也转换为tensor，
return image, label
#training_data包含了本次需要训练的全部数据集？
training_data = food_dataset(file_path = './train.txt',transform = data_transforms['train']) #
test_data = food_dataset(file_path = './test.txt',transform = data_transforms['valid'])
#training_data需要具备索引的功能，还要确保数据是tensor
train_dataloader = DataLoader(training_data, batch_size=64,shuffle=True)#64张图片为一个包，
test_dataloader = DataLoader(test_data,batch_size=64,shuffle=True)data_transforms:
用字典的形式定义了训练集和验证集的形式：
子类class food_dataset：
创建子类继承Datasetj处理图片数据，初始化过程中通过samples = [x.strip().split(' ') for x in f.readlines()]和下面的循环结构获取存储图片位置的列表和标签的列表。.strip()去除字符串两头的空格以及换行符。.split(' ')以空格分隔字符串。执行效果：
Dataloader:
Dataloader打包批次时会自动生成索引，调用__getitem__将图片和标签转化为张量的形式。到这里我们就完成了对原始数据的处理，获得了可以训练的张量形式数据，接下来就是定义神经网络进行训练，和上一篇博文的方法是一样的。
3、定义神经网络
代码：
python
import torch
import torch.nn as nn
''' 定义神经网络 '''
class CNN(nn.Module):#类的名称。 2用法
def __init__(self): # 输入大小
super(CNN, self).__init__()#初始化父类
self.conv1 = nn.Sequential( #将多个层组合成一起。创建了一个容器，将多个网络合在一起
nn.Conv2d( #2d一般用于图像，3d用于视频数据（多一个时间维度），1d一般用于结构化的序列数据28*28
in_channels=3, # 图像通道个数，1表示灰度图（确定了卷积核 组中的个数），
out_channels=16, # 要得到多少个特征图，卷积核的个数
kernel_size=5, # 卷积核大小，5×5
stride=1, # 步长
padding=2, # 一般希望卷积核处理后的结果大小与处理前的数据大小相同，效果会比较好。那padding该如何设计呢？建议stride为1
), # 输出的特征图为16*128*128
nn.ReLU(), # relu层，不会改变特征图的大小
nn.MaxPool2d(kernel_size=2), # 进行池化操作（2x2 区域），输出结果为16*14*14
)
self.conv2 = nn.Sequential( # 输入16*14*14
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2), # 输出32*128*128
nn.ReLU(), # relu层
nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2), # 输出32*128*128
nn.ReLU(),
nn.MaxPool2d(2), #32*64*64
)
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=128, kernel_size=5, stride=1, padding=2), #128*64*64
nn.ReLU(), # 输出
)
self.out = nn.Linear(128*64 *64, out_features=20) # 全连接层得到的结果
def forward(self, x):#这里必须要写 forward是来自于父类nn里面的函数 要继承父类的功能
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)#输出（64,64, 7, 7）
x = x.view(x.size(0), -1) # flatten操作，结果为：(batch_size, 64 * 7 * 7)
output = self.out(x)
return output
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)#类的初始化完成，就会创建一个对象，model
print(model)4、定义训练函数和测试函数
代码：
python
# 5. 定义训练函数
def train(dataloader, model, loss_fn, optimizer):
model.train() # 切换到训练模式
batch_size_num = 1 # 统计batch数量
for X, y in dataloader: # dataloader是函数形参，调用时传入train_dataloader/test_dataloader；迭代返回批次级数据：X为批次图片张量(shape[batch_size,1,28,28])，y为批次标签张量(shape[batch_size])，对应批次内所有样本的图片和标签
# 数据移动到设备
X, y = X.to(device), y.to(device)
# 前向传播计算预测值
pred = model(X) # 可省略.forward，model(X)会自动调用forward
loss = loss_fn(pred, y) # 计算损失
# 反向传播更新参数
optimizer.zero_grad() # 梯度清零
loss.backward() # 反向传播计算梯度（原代码中Loss大写，修正为loss）
optimizer.step() # 更新模型参数
# 每100个batch打印一次损失
loss_value = loss.item()
if batch_size_num % 100 == 0:
print(f"Loss: {loss_value:>7f} [batch: {batch_size_num}]")
batch_size_num += 1
# 6. 定义测试函数
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset) # 测试集总样本数
num_batches = len(dataloader) # 测试集batch数量
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(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):>0.1f}%, Avg loss: {test_loss:>8f} \n")5、选择损失函数和优化器
代码：
python
# 7. 初始化损失函数和优化器
loss_fn = nn.CrossEntropyLoss() # 交叉熵损失（适用于分类任务）
# optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # SGD优化器
# 原代码的optimizer替换为Adam
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam默认lr=0.001即可6、训练和测试
代码：
python
epochs = 10 # 训练10轮
for t in range(epochs):
print(f"\n训练轮数 {t+1}/{epochs}")
train(train_dataloader, model, loss_fn, optimizer)
test(test_dataloader, model, loss_fn) # 每轮训练后测试
print("训练完成！")
print("开始测试：")
test(test_dataloader, model, loss_fn)7、运行结果
保存模型
我们不会每次使用模型都重新训练一次，这里给出两种保存模型的方法：仅保存模型的参数和保存完整的模型（包括卷积神经网络的定义）。
以保存完整模型为例，在测试函数中加入保存模型的语句：
python
def test(dataloader, model, loss_fn):
global best_acc # 声明使用全局变量
size = len(dataloader.dataset)
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(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
# 计算准确率（0~1之间）
correct /= size
# 保存最优模型
if correct > best_acc:
best_acc = correct
# 保存TorchScript完整模型
script_model = torch.jit.script(model)
torch.jit.save(script_model, "best12.pth")
# 仅保存模型参数方法:torch.save(model.state_dict(), path) (w,b)
# # print(model.state_dict().keys()) # 输出模型参数名称 cnn
# # torch.save(model.state_dict(), "best2025-30.pth")
print(f"Test: Accuracy {(100*correct):>0.1f}%, Best Accuracy {(100*best_acc):>0.1f}%")加载模型：
保存完整模型后，我们使用保存的模型，不需要构建神经网络，只需要完成数据的处理，使用保存的模型best12.pth的完整代码：
python
# -*- coding: gbk -*-
# 后续代码...
import torch
from torch.utils.data import Dataset,DataLoader #用于处理数据集的
import numpy as np
from PIL import Image #
from torchvision import transforms #对数据进行处理工具 转换
# 选择设备（与保存时一致或按需调整）
device = "cuda" if torch.cuda.is_available() else "cpu"
# 加载TorchScript模型
model = torch.jit.load("best12.pth", map_location=device)
# 设置模型为评估模式（用于预测，禁用训练相关层的行为）
model.eval()
class food_dataset(Dataset): # food_dataset是自己创建的类名称，可以改为你需要的名称 2用法
def __init__(self, file_path,transform=None): #类的初始化,解析数据文件txt
self.file_path = file_path
self.imgs = []
self.labels = []
self.transform = transform
with open(self.file_path) as f:#是把train.txt文件中图片的路径保存在 self.imgs，train.txt文件中标签保存在 self.labels
samples = [x.strip().split(' ') for x in f.readlines()]
for img_path, label in samples:
self.imgs.append(img_path) #图像的路径
self.labels.append(label) #标签,还不是tensor
#初始化：把图片目录加载到self.imgs.
def __len__(self): #类实例化对象后，可以使用len函数测量对象的个数 ls=[12,3,4,4] len(training_data)
return len(self.imgs)
#training_data[1]
def __getitem__(self, idx): #关键，可通过索引的形式获取每一个图片数据及标签
image = Image.open(self.imgs[idx]) #读取到图片数据，还不是tensor,BGR
if self.transform: #将PIL图像数据转换为tensor
image = self.transform(image) #图像处理为256×256，转换为tenor
label = self.labels[idx] #label还不是tensor
label = torch.from_numpy(np.array(label,dtype = np.int64)) #label也转换为tensor，
return image, label
data_transforms = { #机器学习的时候, 数据进行归一化？？ 图片做归一化 ->0~1
'train':
transforms.Compose([
transforms.Resize([300,300]), #是图像变换大小 opencv int8 0~
transforms.RandomRotation(45),#随机旋转, -45到45度之间随机选
transforms.CenterCrop(256),#从中心开始裁剪[256,256]
transforms.RandomHorizontalFlip(p=0.5),#随机水平翻转 选择一个概率概率
transforms.RandomVerticalFlip(p=0.5),#随机垂直翻转
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),#参数1为亮度, 参数2为对比度, 参数3为饱和度, 参
transforms.RandomGrayscale(p=0.1),#概率转换成灰度率, 3通道就是R=G=B
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])#标准化, 均值, 标准差,
]),
'valid':#验证集 不需要对图像进行数据增强
transforms.Compose([
transforms.Resize([256,256]), #
transforms.ToTensor(),
transforms.Normalize( [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
}
test_data = food_dataset(file_path = './test.txt',transform = data_transforms['valid'])
test_dataloader = DataLoader(test_data, batch_size=1, shuffle=True)
result = []#保存的预测的结果
Lables = []#真实结果
def test_true(dataloader, model):
with torch.no_grad(): #一个上下文管理器，关闭梯度计算。当你确认不会调用Tensor.backward()的时候。这可以减少计算所用内存消耗。
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)#预测之后的结果。
result.append(pred.argmax(1).item())
Lables.append(y.item())
test_true(test_dataloader,model)
print('预测值：\t',result)
print('真实值：\t',Lables)运行结果：
准确率还是很低哈，我们今天主要保证代码能跑通。先跑通代码，再调优。
如果我们使用仅保存模型参数的方法，使用时需要写出与原始模型参数一样的神经网络，这里不再赘述，如想了解可以问豆包等。
欧克，今天我们也是完成了基于卷积神经网络对实物图片的分类，从图片数据整理转化为张量，到构建卷积神经网络，到最后保存模型。完整完美实现了我们的案例项目。