### 使用自定义数据集在PyTorch中实现图像分类模型的训练与推理 #### 准备工作环境 为了确保能够顺利加载并处理自定义的数据集，首先需要安装必要的库。通常情况下，`torch`, `torchvision`, 和 `PIL` 是必不可少的部分。 ```bash pip install torch torchvision pillow ``` #### 创建自定义数据集类 创建一个新的Python文件来定义一个继承自`torch.utils.data.Dataset`的新类，用于读取本地磁盘上的图片路径列表以及对应的标签信息。通过重写`__init__()`,`__getitem__()`和`__len__()`这三个函数来自定义行为[^1]。 ```python from PIL import Image import os from torch.utils.data import Dataset class CustomImageDataset(Dataset): def __init__(self, annotations_file, img_dir, transform=None, target_transform=None): self.img_labels = pd.read_csv(annotations_file) self.img_dir = img_dir self.transform = transform self.target_transform = target_transform def __len__(self): return len(self.img_labels) def __getitem__(self, idx): img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0]) image = Image.open(img_path).convert('RGB') label = self.img_labels.iloc[idx, 1] if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) sample = {"image": image, "label": label} return sample ``` #### 加载预处理后的数据集 利用之前定义好的`CustomImageDataset`类实例化对象，并将其传递给`DataLoader`以便于后续迭代获取batch大小的小批量样本进行训练或者测试操作[^4]。 ```python from torchvision.transforms import ToTensor from torch.utils.data import DataLoader train_dataset = CustomImageDataset( annotations_file='path/to/your/train_annotations.csv', img_dir='path/to/your/images/', transform=ToTensor() ) val_dataset = CustomImageDataset( annotations_file='path/to/your/validation_annotations.csv', img_dir='path/to/your/images/', transform=ToTensor() ) test_dataset = CustomImageDataset( annotations_file='path/to/your/test_annotations.csv', img_dir='path/to/your/images/', transform=ToTensor()) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) validation_loader = DataLoader(val_dataset, batch_size=64, shuffle=False) test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False) ``` #### 构建神经网络架构 可以选择现有的经典卷积神经网络结构如ResNet系列作为基础模型来进行微调(fine-tuning)，也可以设计全新的网络层组合方式满足特定需求[^2]。 ```python model = models.resnet18(pretrained=True) num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, num_classes) # 替换最后一层全连接层适应新的类别数 ``` #### 训练过程配置 设置损失函数（Loss Function）、优化器(Optimizer)以及其他超参数调节项以指导反向传播算法更新权重矩阵直至收敛至局部最优解附近位置停止学习进程[^3]。 ```python criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) scheduler = StepLR(optimizer, step_size=7, gamma=0.1) device = 'cuda' if torch.cuda.is_available() else 'cpu' model.to(device=device); for epoch in range(num_epochs): train_one_epoch(epoch, model, criterion, optimizer, scheduler, device, train_loader) validate_one_epoch(epoch, model, criterion, device, validation_loader) ``` #### 执行推断任务 一旦完成了上述所有准备工作之后就可以开始执行实际的任务了——即对于未知输入给出预测结果。此时只需简单地遍历待测集合中的每一个元素并通过已训练完毕的模型获得其最有可能所属的那个类别编号即可完成整个流程[^5]。 ```python def predict_image_class(image_tensor, model, class_names): with torch.no_grad(): output = model(image_tensor.unsqueeze_(dim=0)) _, predicted_idx = torch.max(output, dim=-1) prediction = class_names[predicted_idx.item()] return prediction # Example usage of the function defined above. predicted_label = predict_image_class(some_test_sample['image'], trained_model, list_of_all_possible_labels) print(predicted_label) ```