数据集准备

从飞桨官网上下载数据集，然后放在本地或者上传到colab。我使用的环境是colab，首先提取出训练集，然后删除无关的文件夹。

#提取数据集
import zipfile
import os
import shutil
zip_files = ['/content/archive_train.zip']
extract_dir = '/content/data'

# 创建提取目录
os.makedirs(extract_dir, exist_ok=True)

# 遍历每个zip文件
for zip_file in zip_files:
    # 打开zip文件
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        # 解压缩文件到目标目录
        zip_ref.extractall(extract_dir)

# 输出提取完成的信息
print("文件提取完成！")
# 删除__MACOSX
items = os.listdir(extract_dir)
for item in items:
  item_path = os.path.join(extract_dir,item)
  #检查文件夹名称是否为"_MACOSX"
  if os.path.isdir(item_path) and item == "__MACOSX":
    shutil.rmtree(item_path)

这个时候得到的文件夹，里面有各个子文件夹，子文件夹里面有n张图片用于训练。

查看训练集信息

#查看数据集信息
items=os.listdir(extract_dir)
for item in items:
  if item==".ipynb_checkpoints":
    shutil.rmtree(os.path.join(extract_dir,item))
    break
  item_path = os.path.join(extract_dir,item)
  jpg_files = [file for file in os.listdir(item_path) if file.lower().endswith(".jpg")]
  print(f"目录{item}下面有{len(jpg_files)}个文件")

目录Zircon下面有33个文件
目录Rhodochrosite下面有29个文件
目录Variscite下面有30个文件
目录Almandine下面有31个文件
目录Iolite下面有32个文件
目录Labradorite下面有40个文件
目录Carnelian下面有33个文件
目录Sapphire Blue下面有34个文件
目录Diamond下面有31个文件
目录Tanzanite下面有36个文件
目录Garnet Red下面有36个文件
目录Hessonite下面有30个文件
目录Quartz Beer下面有35个文件
目录Cats Eye下面有31个文件
目录Fluorite下面有32个文件
目录Kunzite下面有32个文件
目录Jade下面有28个文件
目录Pearl下面有33个文件
目录Emerald下面有36个文件
目录Beryl Golden下面有36个文件
目录Malachite下面有28个文件
目录Benitoite下面有31个文件
目录Danburite下面有32个文件
目录Alexandrite下面有34个文件
目录Onyx Black下面有28个文件

构建dataset

import torchvision.transforms as transforms
import torchvision.datasets as datasets
data_dir="/content/data"
data_transforms=transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
])
dataset=datasets.ImageFolder(data_dir,transform=data_transforms)
print(f"一共有{len(dataset.classes)}个类")
print(f"一共有{len(dataset)}个图片")

构建dataloader，模型，开始训练

import torch
import torch.nn as nn
import torch.optim as optim
#创建数据加载器
batch_size=32
data_loader=torch.utils.data.DataLoader(dataset,batch_size=batch_size,shuffle=True)
#构建模型
model = nn.Sequential(
    nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(32 * 56 * 56, 128),
    nn.ReLU(),
    nn.Linear(128, len(dataset.classes))  # num_classes是分类的类别数
)
#定义损失函数
criterion = nn.CrossEntropyLoss()

#定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
#训练模型
num_epochs = 10

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    running_loss = 0.0
    for images, labels in data_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(data_loader)
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}")

保存模型，并进行预测

#保存模型
import torch
import torchvision.datasets as datasets
torch.save(model.state_dict(),"model.pt")
#进行测试
test_data="/content/data/test"
os.makedirs(test_data, exist_ok=True)
with zipfile.ZipFile("/content/archive_test.zip",'r') as zip_ref:
  zip_ref.extractall(test_data)
#删除测试集中无关内容
for item in os.listdir(test_data):
  if not item.endswith(".jpg"):
    shutil.rmtree(os.path.join(test_data,item))
#创建测试集，将测试集图片放入到类别中，方便ImageFolder导入
for item in os.listdir(test_data):
  label=item.split("_")[0].capitalize()
  os.makedirs(os.path.join(test_data,label),exist_ok=True)
  shutil.move(os.path.join(test_data,item),os.path.join(test_data,label,item))
# 定义图像转换操作
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载测试集
test_dataset = datasets.ImageFolder(test_data, transform=transform)
model.eval()
for image, label in test_dataset:
    # 做前向传播预测
    output = model(image.unsqueeze(0))  # 添加批次维度
    _, predicted = torch.max(output, dim=1)

    # 获取预测结果和标签名称
    predicted_class = test_dataset.classes[predicted.item()]
    true_class = test_dataset.classes[label]

    # 打印预测结果和标签
    print(f"Predicted: {predicted_class}, True: {true_class}")

问题

我写的代码是真的丑陋

本章内容来自百度松果