PyTorch实战：基于Kaggle新冠影像数据集的四分类模型构建与优化

1. 项目背景与数据集介绍医学影像分类一直是人工智能在医疗领域的重要应用方向。最近几年随着深度学习技术的快速发展基于卷积神经网络(CNN)的医学影像分析取得了显著进展。特别是在肺部疾病诊断方面AI技术展现出了巨大潜力。Kaggle上的COVID-19放射学数据库是一个经过精心整理的公开数据集由卡塔尔大学、达卡大学等研究机构联合医生团队共同创建。这个数据集包含了四种类型的胸部X光图像COVID-19阳性病例正常病例肺不透明(非COVID-19肺感染)病毒性肺炎数据集的最新版本包含超过2万张图像其中COVID-19阳性病例3616张正常病例10192张肺不透明6012张病毒性肺炎1345张。每张图像都配有对应的肺部掩码这对于后续的病灶定位分析非常有帮助。提示使用这个数据集前需要注册Kaggle账号并同意相关使用条款数据集下载可能需要科学上网工具。2. 数据预处理实战2.1 数据集划分与目录结构医学影像项目的第一步永远是数据准备。我们需要将原始数据集划分为训练集、验证集和测试集比例通常为7:1:2。这种划分确保了模型训练时有足够的数据同时也能客观评估模型性能。import os import shutil # 定义数据集路径和输出路径 dataset_path COVID-19_Radiography_Dataset output_path processed_dataset # 创建目录结构 os.makedirs(os.path.join(output_path, train), exist_okTrue) os.makedirs(os.path.join(output_path, val), exist_okTrue) os.makedirs(os.path.join(output_path, test), exist_okTrue) categories [COVID, Lung_Opacity, Normal, Viral Pneumonia] for split in [train, val, test]: for category in categories: os.makedirs(os.path.join(output_path, split, category, images), exist_okTrue) os.makedirs(os.path.join(output_path, split, category, masks), exist_okTrue)2.2 数据增强与标准化医学影像数据通常存在样本不平衡问题COVID-19阳性样本相对较少。我们可以通过数据增强技术来缓解这个问题from torchvision import transforms # 训练集数据增强 train_transform transforms.Compose([ transforms.Resize(256), transforms.RandomRotation(10), transforms.RandomHorizontalFlip(), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 验证集和测试集只需基础预处理 val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])3. 模型构建与训练3.1 自定义CNN模型设计我们设计了一个包含4个卷积层的CNN模型每层后面都跟着批归一化和ReLU激活import torch.nn as nn class COVIDClassifier(nn.Module): def __init__(self, num_classes4): super(COVIDClassifier, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), nn.Conv2d(128, 256, kernel_size3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2) ) self.classifier nn.Sequential( nn.Linear(256*14*14, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(512, num_classes) ) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) x self.classifier(x) return x3.2 训练策略与超参数调优训练深度学习模型需要仔细选择超参数和优化策略import torch.optim as optim from torch.optim.lr_scheduler import ReduceLROnPlateau # 初始化模型 model COVIDClassifier(num_classes4).to(device) # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001, weight_decay1e-5) # 学习率调度器 scheduler ReduceLROnPlateau(optimizer, modemax, factor0.1, patience3, verboseTrue) # 早停机制 best_accuracy 0 patience 5 no_improve 0 for epoch in range(100): train_loss train_one_epoch(model, train_loader, criterion, optimizer) val_loss, val_acc evaluate(model, val_loader, criterion) # 更新学习率 scheduler.step(val_acc) # 早停检查 if val_acc best_accuracy: best_accuracy val_acc no_improve 0 torch.save(model.state_dict(), best_model.pth) else: no_improve 1 if no_improve patience: print(fEarly stopping at epoch {epoch}) break4. 模型评估与优化4.1 性能评估指标对于多分类问题我们不仅需要关注准确率还需要考察其他指标from sklearn.metrics import classification_report def get_classification_report(model, data_loader): model.eval() all_preds [] all_labels [] with torch.no_grad(): for inputs, labels in data_loader: inputs inputs.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) return classification_report(all_labels, all_preds, target_names[COVID, Lung_Opacity, Normal, Viral Pneumonia]) print(get_classification_report(model, test_loader))4.2 模型优化技巧为了提高模型性能我们可以尝试以下优化策略迁移学习使用预训练的ResNet或EfficientNet作为基础模型类别权重为样本较少的类别设置更高的权重混合精度训练使用AMP加速训练过程模型集成结合多个模型的预测结果# 使用预训练模型示例 from torchvision import models def get_pretrained_model(num_classes4): model models.efficientnet_b0(pretrainedTrue) # 冻结特征提取层 for param in model.features.parameters(): param.requires_grad False # 修改分类头 model.classifier[1] nn.Linear(model.classifier[1].in_features, num_classes) return model5. 部署与应用5.1 单张图像预测训练好的模型可以用于单张图像的分类预测from PIL import Image import matplotlib.pyplot as plt def predict_image(image_path, model, transform): image Image.open(image_path).convert(RGB) image_tensor transform(image).unsqueeze(0).to(device) with torch.no_grad(): output model(image_tensor) _, pred torch.max(output, 1) class_name [COVID, Lung_Opacity, Normal, Viral Pneumonia][pred.item()] plt.imshow(image) plt.title(fPrediction: {class_name}) plt.axis(off) plt.show() return class_name # 使用示例 predict_image(test_image.png, model, val_transform)5.2 模型轻量化与部署为了在实际应用中部署模型我们可以进行模型压缩# 模型量化示例 quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear}, dtypetorch.qint8 ) # 保存量化模型 torch.jit.save(torch.jit.script(quantized_model), quantized_covid_model.pt)在实际部署时可以考虑使用Flask或FastAPI构建Web服务或者将模型集成到移动应用中。对于医疗应用场景还需要特别注意模型的解释性可以使用Grad-CAM等技术可视化模型关注的区域。