### 使用PyTorch和ResNet50模型训练自定义数据集 #### 准备工作环境 为了确保能够顺利运行代码，需先安装必要的Python库。如果尚未安装这些依赖项，则可以通过pip命令来完成。 ```bash pip install torch torchvision tqdm matplotlib ``` #### 导入所需模块 导入用于构建神经网络、加载数据以及辅助功能所需的各类库文件。 ```python import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms, models import numpy as np import matplotlib.pyplot as plt import os from tqdm.auto import tqdm # 显示进度条 ``` #### 数据预处理与增强 对于图像分类任务来说，适当的数据增广可以提高泛化能力并防止过拟合现象的发生。这里采用了一些常见的变换操作如随机裁剪、水平翻转等，并将所有输入图片调整到统一尺寸以便于后续计算。 ```python 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]) ]), } ``` #### 加载数据集 根据个人需求指定路径读取本地存储的图片资料，并创建相应的迭代器对象方便批量获取样本及其标签信息。 ```python data_dir = './path_to_your_dataset' # 替换成实际存放数据的位置 image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']} dataloaders = {x: DataLoader(image_datasets[x], batch_size=32, shuffle=True, num_workers=4) for x in ['train', 'val']} dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 判断是否有GPU可用 ``` #### 修改ResNet50的最后一层 由于原始的ResNet50是在ImageNet上预训练得到的一个大型通用特征提取器，默认情况下其最后一层是一个具有1000个节点的全连接层用来做类别预测；而当面对新的特定领域问题时往往只需改变这部分结构使之适应目标数量即可达到很好的迁移学习效果[^3]。 ```python model_ft = models.resnet50(pretrained=True) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, len(class_names)) # 将最后输出改为对应新任务类别的数目 model_ft = model_ft.to(device) criterion = nn.CrossEntropyLoss() optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1) def train_model(model, criterion, optimizer, scheduler, num_epochs=25): best_model_wts = model.state_dict() best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch}/{num_epochs - 1}') print('-' * 10) for phase in ['train', 'val']: if phase == 'train': model.train() else: model.eval() running_loss = 0.0 running_corrects = 0 for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) if phase == 'train': 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 / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = model.state_dict() print() model.load_state_dict(best_model_wts) return model if __name__ == '__main__': model_trained = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25) ```