import torch
import torch.nn as nn
import torch.optim as optim
from datasets import load_dataset
from torch.utils.data import DataLoader
from torchvision import models, transforms
from PIL import Image
# ===============================
# 1. 加载 Hugging Face 数据集
# ===============================
# 加载 Hugging Face 蘑菇数据集
dataset = load_dataset("mushrooms")
# 检查数据集的结构
print(dataset)
# 查看训练集的一个样本
print(dataset['train'][0])
# ===============================
# 2. 数据预处理
# ===============================
# 定义图像转换
transform = transforms.Compose([
transforms.Resize((224, 224)), # 调整图像大小
transforms.ToTensor(), # 转换为 PyTorch 张量
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 标准化
])
# 自定义函数，将数据集中的图片进行处理
def transform_examples(examples):
examples['image'] = [transform(Image.fromarray(img)) for img in examples['image']]
return examples
# 应用转换到训练集和测试集
train_dataset = dataset['train'].map(transform_examples, batched=True)
test_dataset = dataset['test'].map(transform_examples, batched=True)
# ===============================
# 3. 创建 DataLoader
# ===============================
# 创建 PyTorch DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# 打印一个 batch 的数据，确保数据处理正确
for batch in train_loader:
images, labels = batch['image'], batch['label']
print(images.shape) # (batch_size, 3, 224, 224)
print(labels.shape) # (batch_size)
break
# ===============================
# 4. 加载预训练的 ResNet18 模型并进行迁移学习
# ===============================
# 加载预训练的 ResNet18 模型
model = models.resnet18(pretrained=True)
# 替换最后一层全连接层，输出类别数为 2（可食用和有毒）
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)
# 将模型移到 GPU（如果可用）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# ===============================
# 5. 训练模型
# ===============================
num_epochs = 10 # 训练轮数
for epoch in range(num_epochs):
model.train() # 训练模式
running_loss = 0.0
correct = 0
total = 0
for batch in train_loader:
images, labels = batch['image'].to(device), batch['label'].to(device)
optimizer.zero_grad()
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播和优化
loss.backward()
optimizer.step()
# 统计训练损失和准确率
running_loss += loss.item() * images.size(0)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
epoch_loss = running_loss / len(train_dataset)
epoch_acc = correct / total
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.4f}')
# ===============================
# 6. 验证模型
# ===============================
model.eval() # 评估模式
correct = 0
total = 0
with torch.no_grad():
for batch in test_loader:
images, labels = batch['image'].to(device), batch['label'].to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Test Accuracy: {correct / total:.4f}')
# ===============================
# 7. 保存模型
# ===============================
torch.save(model.state_dict(), 'mushroom_classifier.pth')
在本项目中，我们将介绍如何使用 PyTorch 框架和 Hugging Face 数据集来构建一个能够识别蘑菇是否可食用的分类器。通过迁移学习，我们将利用预训练的 ResNet18 模型，进行模型训练和评估。以下是整个过程的详细步骤，包括数据加载、预处理、模型训练、验证和最终保存模型。
1. 加载数据集
数据集简介
首先，我们加载了 Hugging Face 提供的蘑菇数据集。这个数据集包含了大量蘑菇的图像，并标注了每种蘑菇的类别，主要分为可食用和有毒两类。这一数据集非常适合用于训练图像分类模型，因为它不仅提供了丰富的图像样本，还包含了必要的标签信息。
加载数据集
我们利用 datasets 库中的 load_dataset 函数来加载数据集，并打印出数据集的结构和第一个样本，以确认数据是否成功加载。通过这种方式，我们可以快速了解数据集中包含的信息类型和格式。
2. 数据预处理
图像转换
在加载数据集后，我们需要对图像进行预处理，以确保它们适合输入到深度学习模型中。我们使用 torchvision.transforms 模块来定义图像转换操作，包括：
- 调整图像大小：将所有图像调整为统一的 224x224 像素大小，以符合 ResNet18 模型的输入要求。
- 转换为张量：将图像数据转换为 PyTorch 张量，便于后续的处理和计算。
- 标准化：使用 ImageNet 数据集的均值和标准差对图像进行标准化，确保模型在训练过程中能够更好地收敛。
应用转换
通过自定义函数，我们将定义的图像转换应用到训练集和测试集上。这个过程会遍历数据集中的所有图像，并逐一应用转换，从而生成适合模型输入格式的新数据集。
3. 创建 DataLoader
为了高效地处理数据，我们使用 PyTorch 的 DataLoader 来创建训练和测试数据的批次。`DataLoader` 能够自动处理数据的分批、打乱和多线程加载等操作，使得训练过程更加高效。我们设置批量大小为 32，并对训练集进行随机打乱，以增强模型的泛化能力。
4. 加载预训练模型并进行迁移学习
选择模型
在本项目中，我们选择了预训练的 ResNet18 模型，这是一种经典的卷积神经网络（CNN），在许多图像分类任务中表现优异。通过使用预训练模型，我们可以利用在大型数据集上获得的知识，加速我们的训练过程。
修改模型结构
由于我们的任务是二分类（可食用与有毒），我们需要修改 ResNet18 的最后一层全连接层，将其输出数量改为 2。这样，模型就能够正确地输出每种类别的概率。
设备配置
接下来，我们将模型移到 GPU（如果可用），这能显著加速训练过程。我们检查计算设备，并将模型和输入数据都转移到相应的设备上。
5. 训练模型
设置损失函数和优化器
在训练之前，我们需要定义损失函数和优化器。我们选择了交叉熵损失（Cross-Entropy Loss），这是多分类任务中常用的损失函数。同时，我们使用 Adam 优化器进行参数更新，这种优化器在训练深度学习模型时通常表现良好。
训练过程
在每个训练周期中，我们将模型设置为训练模式，并初始化运行损失和正确分类的计数。通过遍历训练数据加载器中的每个批次，进行以下步骤：
1. 前向传播：将输入图像通过模型，得到输出。
2. 计算损失：使用交叉熵损失函数计算模型输出与真实标签之间的差距。
3. 反向传播和优化：通过反向传播算法计算梯度，并使用优化器更新模型参数。
我们在每个周期结束时计算平均损失和分类准确率，并打印结果，帮助我们监控模型的训练进度。
6. 验证模型
训练完成后，需要验证模型的性能。我们将模型设置为评估模式，并关闭梯度计算（使用 torch.no_grad()），以减少内存消耗。在这一阶段，我们遍历测试数据加载器，进行类似于训练的步骤，计算模型在测试集上的分类准确率，以评估其泛化能力。