1. 医疗影像分类模型定制
任务目标
使用公开胸部X光数据集（COVID-19 Radiography Database），训练一个轻量CNN模型，实现肺炎（PNEUMONIA）与正常（NORMAL）的二分类，并在测试集上评估准确率。本案例聚焦于完整可复现的最小训练闭环，不涉及部署或超参数调优。
环境与依赖
- Python 3.9+，PyTorch 2.0+，torchvision，Pillow，numpy，matplotlib（可选）
- 数据集：COVID-19 Radiography Database（下载后解压，目录结构需为
train/NORMAL、train/PNEUMONIA、test/NORMAL、test/PNEUMONIA）
数据准备说明
原始数据集包含COVID、Lung Opacity、Normal、Viral Pneumonia四类。本案例仅使用其中的 NORMAL 与 PNEUMONIA（即Viral Pneumonia）两个子目录。请将原始数据按下述结构整理后，直接运行下方脚本即可完成训练与评估。
关键完整代码
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
import os
import time
# ---------- 配置 ----------
DATA_ROOT = "./data" # 包含 train/ 和 test/ 的根目录
BATCH_SIZE = 32
EPOCHS = 5
LR = 1e-3
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
NUM_CLASSES = 2
# ---------- 数据增强与归一化 ----------
train_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
test_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# ---------- 加载数据 ----------
train_dataset = datasets.ImageFolder(os.path.join(DATA_ROOT, "train"), transform=train_transform)
test_dataset = datasets.ImageFolder(os.path.join(DATA_ROOT, "test"), transform=test_transform)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
print(f"训练集样本数: {len(train_dataset)}, 测试集样本数: {len(test_dataset)}")
print(f"类别映射: {train_dataset.class_to_idx}")
# ---------- 模型：ResNet18迁移学习 ----------
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)
model = model.to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)
# ---------- 训练 ----------
def train_one_epoch(model, loader, criterion, optimizer):
model.train()
total_loss, correct, total = 0, 0, 0
for images, labels in loader:
images, labels = images.to(DEVICE), labels.to(DEVICE)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * images.size(0)
_, preds = torch.max(outputs, 1)
correct += (preds == labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
# ---------- 评估 ----------
def evaluate(model, loader, criterion):
model.eval()
total_loss, correct, total = 0, 0, 0
with torch.no_grad():
for images, labels in loader:
images, labels = images.to(DEVICE), labels.to(DEVICE)
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item() * images.size(0)
_, preds = torch.max(outputs, 1)
correct += (preds == labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
# ---------- 执行训练 ----------
print(f"训练设备: {DEVICE}")
for epoch in range(EPOCHS):
train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)
val_loss, val_acc = evaluate(model, test_loader, criterion)
print(f"Epoch {epoch+1}/{EPOCHS} | "
f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
f"Test Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
# ---------- 最终测试集评估 ----------
test_loss, test_acc = evaluate(model, test_loader, criterion)
print(f"\n最终测试集准确率: {test_acc:.4f} ({test_acc*100:.2f}%)")
运行方式
- 将数据集整理为如下结构：
data/
├── train/
│ ├── NORMAL/ # 正常胸片
│ └── PNEUMONIA/ # 肺炎胸片
└── test/
├── NORMAL/
└── PNEUMONIA/
- 安装依赖：
pip install torch torchvision pillow numpy - 将上述代码保存为
train_pneumonia.py，与data目录同级，执行python train_pneumonia.py
验证方法与通过条件
- 运行成功标志：控制台输出训练集/测试集样本数、类别映射，且每个Epoch均打印训练与测试的Loss和Accuracy，无报错中断。
- 判定通过标准：最终测试集准确率 ≥ 85%。若低于该值，可增加训练轮次至10轮或降低学习率至1e-4后重试。
- 预期输出示例（基于代码逻辑，非实测值）：训练约5轮后，测试准确率通常在0.88～0.95区间，具体数值取决于数据划分与随机初始化。
常见问题
- 显存不足：将
BATCH_SIZE降至16或8。 - 数据集路径错误：确认
DATA_ROOT指向包含train和test的上级目录，而非train本身。 - 类别顺序：ImageFolder按文件夹名称字母序映射类别，
NORMAL为0，PNEUMONIA为1，代码中已自动处理。