### 创建自定义数据集进行图像分类 为了使用 PyTorch 框架创建自定义数据集并实现图像分类，可以遵循以下方法： #### 导入必要库 在编写代码之前，需确保已安装 `torch` 和 `torchvision` 库。接着导入所需的模块来处理图像和构建模型。 ```python import os from torch import nn, optim import torchvision.transforms as transforms from torchvision import models, datasets from torch.utils.data import DataLoader ``` #### 定义自定义数据集类 通过继承 `torch.utils.data.Dataset` 类来自定义数据集。重写 `__init__`, `__len__` 及 `__getitem__` 方法以适应特定的数据结构[^2]。 ```python class CustomDataset(datasets.VisionDataset): def __init__(self, root, transform=None, target_transform=None): super(CustomDataset, self).__init__( root, transform=transform, target_transform=target_transform) # 加载图片路径列表与对应标签 self.samples = [] with open(os.path.join(root, 'annotations.txt')) as f: lines = f.readlines() for line in lines: img_path, label = line.strip().split(' ') self.samples.append((img_path, int(label))) def __getitem__(self, index): path, target = self.samples[index] sample = Image.open(path).convert('RGB') if self.transform is not None: sample = self.transform(sample) if self.target_transform is not None: target = self.target_transform(target) return sample, target def __len__(self): return len(self.samples) ``` 此段代码展示了如何读取包含图像路径及其相应类别索引的文本文件，并将其转换为可用于训练的形式[^3]。 #### 设置数据增强及预处理操作 利用 `transforms.Compose()` 函数组合多个变换函数应用于输入数据上，比如调整大小、裁剪、翻转等随机变化可增加模型泛化能力。 ```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]) ]) } ``` #### 构建DataLoader实例用于批量加载数据 借助于 `torch.utils.data.DataLoader` 来高效地迭代整个数据集，在每次迭代过程中返回一批次样本以及对应的标签。 ```python image_datasets = {x: CustomDataset( root=f'data/{x}', transform=data_transforms[x]) for x in ['train', 'val']} dataloaders = {x: DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4) for x in ['train', 'val']} dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} ``` 上述过程完成了从原始图像到适合神经网络输入形式之间的转换工作流设计。 #### 修改ResNet最后的全连接层适配新任务 对于迁移学习场景下，通常只需修改最后一层即全连接层即可完成新的分类任务设置。这里假设目标数据集中共有 N 种不同类型的对象，则应将 ResNet 预训练权重下载下来之后更改其顶层架构使之匹配当前需求[^1]。 ```python model_ft = models.resnet18(pretrained=True) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, N) ``` 至此已经介绍了完整的流程，包括但不限于定义自定义数据集类、配置数据增强策略、建立批处理机制以及针对具体应用场景定制化的微调方案。