一、第三方库介绍
| 库/模块 | 功能 |
|---|---|
| torch | 提供张量操作、自动求导、优化算法、神经网络模块等基础设施。 |
| torchvision | 计算机视觉工具集，提供预训练模型、数据集、图像转换等功能。 |
| datasets (torchvision) | 用于加载常见数据集（如 ImageNet、CIFAR-10、MNIST）。 |
| transforms (torchvision) | 提供图像数据的预处理、数据增强操作（如大小调整、裁剪、转换为张量、归一化等）。 |
| nn (torch) | 用于定义和构建神经网络，包含各类网络层、损失函数等。 |
| optim (torch) | 提供优化算法（如 Adam、SGD、RMSprop）用于更新神经网络权重。 |
| DataLoader (torch.utils.data) | 用于批量加载数据，支持多线程加载数据，按批次读取数据。 |
| Image (PIL) | 用于图像处理（加载、裁剪、旋转、缩放、保存图像等）。 |
| ResNet18_Weights (torchvision.models) | 提供 ResNet18 模型的预训练权重，可用于迁移学习。 |
二、训练数据集介绍
三、原理简介
该代码使用PyTorch训练一个基于ResNet-18的猫狗分类模型。通过加载并处理数据、训练模型、调整最后输出层、使用Adam优化器进行反向传播，并在每个训练周期输出损失与准确率。训练完毕后，保存模型用于后续预测。
四、代码思路简介
- 加载数据 → 使用
datasets.ImageFolder加载猫狗数据集，并应用图像转换。 - 构建模型 → 使用预训练的 ResNet-18 模型，修改输出层以适应2类分类。
- 定义损失和优化器 → 使用交叉熵损失函数和 Adam 优化器。
- 训练模型 → 遍历数据集，前向传播、计算损失、反向传播、更新模型参数。
- 保存模型 → 训练完成后，保存模型权重。
- 预测图片 → 加载已训练模型，输入图片进行预测并输出分类结果。
五、代码
场景：使用pytorch识别猫狗
猫的图片路径：F:\pycharm\AIDEMO\data\cat
狗的图片路径：F:\pycharm\AIDEMO\data\dog
需要判断的图片：F:\pycharm\AIDEMO\01.jpeg
python
import torch
import torchvision
from torchvision import datasets, transforms
from torch import nn, optim
from torch.utils.data import DataLoader
from PIL import Image
from torchvision.models import ResNet18_Weights
# 定义transform类（视觉转换类，将图片格式转化为张量格式）
transform = transforms.Compose([
transforms.Resize((128, 128)), # 将图片缩放到统一大小
transforms.ToTensor(), # 转换为Tensor格式
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 标准化处理
])
def train_model(data_dir, num_epochs=10, batch_size=32, save_path='cat_dog_model.pth'):
"""
训练模型并保存。
:param data_dir: 数据路径，包含cat和dog文件夹
:param num_epochs: 训练周期，默认为10
:param batch_size: 批次大小，默认为32
:param save_path: 模型保存路径，默认为'cat_dog_model.pth'
"""
# 1. 加载训练数据
train_data = datasets.ImageFolder(
root=data_dir, # 数据路径
transform=transform
)
print(train_data.class_to_idx) # 输出文件夹编号，例如这里输出{'cat': 0, 'dog': 1}，表达0代表猫猫，1代表狗狗
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
# 2. 使用预训练的ResNet18模型
model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(model.fc.in_features, 2) # 修改输出层以适应2类分类（猫、狗）
# 3. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 4. 开始训练模型
for epoch in range(num_epochs):
model.train() # 设置模型为训练模式
running_loss = 0.0 # 初始化损失值
correct = 0 # 模型预测准确数
total = 0 # 模型预测总数
for images, labels in train_loader:
optimizer.zero_grad() # 清除之前的梯度
outputs = model(images) # 前向传播，得出预测结果
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
running_loss += loss.item()
# 计算准确率
_, predicted = torch.max(outputs, 1) # 获取预测结果
total += labels.size(0) # 累计总样本数
correct += (predicted == labels).sum().item() # 累计预测正确的样本数
# 输出训练周期的损失和准确率
accuracy = 100 * correct / total # 计算准确率
print(f'周期 [{epoch + 1}/{num_epochs}], 损失: {running_loss / len(train_loader):.4f}, 准确率: {accuracy:.2f}%')
if accuracy == 100: # 准确率达到100%就停止训练，避免过度拟合
break
# 保存训练模型
torch.save(model.state_dict(), save_path)
def predict_image(model_path, img_path):
"""
加载训练好的模型并进行图片预测。
:param model_path: 训练好的模型路径
:param img_path: 需要预测的图片路径
:return: 预测结果（猫或狗）
"""
# 加载模型
model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(model.fc.in_features, 2)
model.load_state_dict(torch.load(model_path))
model.eval() # 设置模型为评估模式
# 预测指定图片
img = Image.open(img_path)
img = transform(img).unsqueeze(0) # 将图片处理成张量输出并增加batch维度
# 模型预测
with torch.no_grad(): # 不需要梯度计算，只是进行模型预测
outputs = model(img)
_, predicted = torch.max(outputs, 1)
# 输出预测结果
return "这是猫的图片" if predicted.item() == 0 else "这是狗的图片"
if __name__ == "__main__":
# 01 训练出模型（若已训练出准确度较高模型，可注释下面两句话，直接用训练完毕的模型预测）
data_dir = 'F:/pycharm/AIDEMO/data' # 数据路径
train_model(data_dir, num_epochs=10, batch_size=32)
# 02 预测指定图片
img_path = 'F:/pycharm/AIDEMO/data/01.jpeg' # 图片路径
result = predict_image('cat_dog_model.pth', img_path)
print(result)六、输出结果展示
