1.背景及准备

书接上文【以图搜图代码实现】–犬类以图搜图示例 总结了一下可以优化的点，其中提到使用自己的数据集训练网络，而不是单纯使用预训练的模型，这不就来了！！

使用11类犬类微调resnet18网络模型：
1. 数据准备
【数据集】11种犬类，共1089张
链接：百度网盘链接
提取码：qlrt

2. 数据集划分
按照train和val8:2的比例进行划分，划分代码如下：

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：ImageRec 
@File    ：split_data.py
@IDE     ：PyCharm 
@Author  ：菜菜2024
@Date    ：2024/9/30
'''
import os
import shutil
import randomdef split_images_into_train_test(source_directory, train_directory, val_directory, train_ratio=0.8):"""将源文件夹下的图片按照指定比例分为训练集和测试集，并分别复制到train和val文件夹下。"""# 确保train和test目录存在，如果不存在则创建os.makedirs(train_directory, exist_ok=True)os.makedirs(val_directory, exist_ok=True)# 获取源文件夹中所有图片文件的列表image_files = [f for f in os.listdir(source_directory) if os.path.isfile(os.path.join(source_directory, f))]image_files = [f for f in image_files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]# 打乱图片文件列表的顺序random.shuffle(image_files)total_images = len(image_files)train_images_count = int(train_ratio * total_images)# 将图片复制到对应的文件夹下for i, image_file in enumerate(image_files):source_path = os.path.join(source_directory, image_file)if i < train_images_count:dest_path = os.path.join(train_directory, image_file)else:dest_path = os.path.join(val_directory, image_file)shutil.copy2(source_path, dest_path)  print(f"Copied {image_file} to {os.path.dirname(dest_path)}")if __name__ == '__main__':source_directory = "E:\\xxx\\datas\\imgs"train_directory = "E:\\xxx\\datas\\pet_dog\\train"val_directory = "E:\\xxx\\datas\\pet_dog\\val"file_list = os.listdir(source_directory)for file in file_list:source=os.path.join(source_directory, file)val = os.path.join(val_directory, file)train = os.path.join(train_directory, file)split_images_into_train_test(source, train, val)

最终效果：

train和val下的目录结果都是如下图所示，只是数量不一样。

2.代码实现

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：ImageRec 
@File    ：train.py
@IDE     ：PyCharm 
@Author  ：菜菜2024
@Date    ：2024/9/30  
'''
import torch
import os
import torch.optim as optim
from torch.optim import lr_scheduler
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
import argparsedef train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs=5):"""参数:model: torch.nn.Module - 要训练的模型实例。dataloaders: dict - 包含训练集和验证集的数据加载器，例如{'train': train_loader, 'val': val_loader}。criterion: nn.Module - 用于计算损失的函数。optimizer: torch.optim.Optimizer - 用于更新模型参数的优化器。scheduler: torch.optim.lr_scheduler._LRScheduler - 学习率调度器。num_epochs: int - 训练的总轮数。"""device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model.to(device)for epoch in range(num_epochs):print(f"Epoch {epoch + 1}/{num_epochs}")for phase in ['train', 'val']:if phase == 'train':model.train()  # 设置模型为训练模式else:model.eval()   # 设置模型为评估模式running_loss = 0.0running_corrects = 0i = 0for inputs, labels in dataloaders[phase]:i+=1inputs, labels = inputs.to(device), labels.to(device)# 前向传播with torch.set_grad_enabled(phase == 'train'):outputs = model(inputs)loss = criterion(outputs, labels)if i%10==0:print(f"{phase} Loss: {loss:.4f}")_, preds = torch.max(outputs, 1)# 反向传播与优化（仅在训练阶段）if phase == 'train':optimizer.zero_grad()loss.backward()optimizer.step()running_loss += loss.item() * inputs.size(0)running_corrects += torch.sum(preds == labels.data)if phase == 'train':scheduler.step()epoch_loss = running_loss / len(dataloaders[phase].dataset)epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)print(f"{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")print("Training complete.")# 训练完成后，保存模型状态字典（包含权重）torch.save(model.state_dict(), './weights/resnet18_dog.pth')def main():# 创建参数解析器parser = argparse.ArgumentParser(description='使用自己的数据集训练resnet18')# 添加参数parser.add_argument('--data_dir', type=str, default="E:\HWR_files\datas\pet_dog",help='Path to the dataset directory')parser.add_argument('--batch_size', type=int, default=16, help='Input batch size for training (default: 16)')parser.add_argument('--num_workers', type=int, default=2, help='Number of workers for data loading (default: 2)')parser.add_argument('--learning_rate', type=float, default=0.001, help='Learning rate (default: 0.001)')parser.add_argument('--num_epochs', type=int, default=5, help='Number of epochs to train (default: 25)')args = parser.parse_args()# 数据预处理和加载data_transforms = {'train': transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),'val': transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),}image_datasets = {x: datasets.ImageFolder(os.path.join(args.data_dir, x), data_transforms[x]) for x in['train', 'val']}dataloaders_dict = {x: DataLoader(image_datasets[x], batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) for xin ['train', 'val']}# 使用ResNet18模型model = models.resnet18(pretrained=True)# 加载之前保存的权重# model.load_state_dict(torch.load('./weights/resnet18_dog.pth'))num_features = model.fc.in_featuresmodel.fc = torch.nn.Linear(num_features, len(image_datasets['train'].classes))  # 修改最后一层全连接层device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = model.to(device)# 定义损失函数和优化器criterion = torch.nn.CrossEntropyLoss()optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9)scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)# 训练模型train_model(model, dataloaders_dict, criterion, optimizer, scheduler, args.num_epochs)if __name__ == '__main__':main()

3.代码测试

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：ImageRec 
@File    ：test.py
@IDE     ：PyCharm 
@Author  ：菜菜2024
@Date    ：2024/9/30 
'''
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, modelsdef test_model(weights_path, val_root, batch_size=4):"""使用验证集测试模型性能。参数:- weights_path: str, 训练好的模型权重文件路径- val_root: str, 验证数据集的根目录- batch_size: int, 数据加载时的批次大小"""# 设定数据预处理transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])# 加载验证集val_dataset = datasets.ImageFolder(root=val_root, transform=transform)val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)# 初始化模型并加载权重device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 修改最后一层全连接层确保num_classes与实际类别数匹配model = models.resnet18()num_features = model.fc.in_featuresmodel.fc = torch.nn.Linear(num_features, len(val_dataset.classes))model.load_state_dict(torch.load(weights_path, map_location=device))model.to(device)model.eval()  # 设置模型为评估模式# 测试循环correct = 0total = 0with torch.no_grad():for images, labels in val_loader:images, labels = images.to(device), labels.to(device)outputs = model(images)_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()# 计算准确率并打印结果accuracy = 100 * correct / totalprint(f'Accuracy on validation set: {accuracy}%')if __name__ == '__main__':weights_path = './weights/resnet18_dog.pth'val_root = "E:\\xxx\\datas\\pet_dog\\val"test_model(weights_path, val_root)

结果图：

4.效果对比

书接上篇的图像检索：【以图搜图代码实现】–犬类以图搜图示例
来看看有没有准一点的

使用预训练的resnet18:

离谱了，匹配的前三个都是吉娃娃

看看使用微调之后的resnet18:
对应在上一篇种，模型加载和最后一层的输出个数变成类别数，这里是11。

哇哇哇！效果显著呀！！！

可以可以，下次尝试使用faiss喽