ch01_resnet152_finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "525e220d-fb89-42ba-90f0-f0a4c43d878c",
   "metadata": {},
   "source": [
    "# 基于ResNet152微调实现胸部X-Ray图像肺炎分类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faf76373-3b99-4a0b-8d56-0271ed4d3264",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import kagglehub\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torchvision\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from PIL import Image\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961de779",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f67d0e7a",
   "metadata": {},
   "source": [
    "## 一、下载数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd821d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# enter kaggle username and token\n",
    "kagglehub.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06af3e6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download dataset to .cache\n",
    "download_path = kagglehub.dataset_download(\"jtiptj/chest-xray-pneumoniacovid19tuberculosis\")\n",
    "print(download_path)\n",
    "# then copy files to data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeec139d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChestXrayDataset(Dataset):\n",
    "    def __init__(self, data_dir, transform=None):\n",
    "        \"\"\"\n",
    "        Dataset init.\n",
    "        :param data_dir: Dataset root directory ('train''test''val')\n",
    "        :param transform: Data preprocessing\n",
    "        \"\"\"\n",
    "        self.data_dir = data_dir\n",
    "        self.transform = transform\n",
    "        self.classes = ['COVID19', 'NORMAL', 'PNEUMONIA', 'TURBERCULOSIS']\n",
    "        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}\n",
    "        self.samples = self._load_samples()\n",
    "\n",
    "    def _load_samples(self):\n",
    "        samples = []\n",
    "        for cls in self.classes:\n",
    "            cls_dir = os.path.join(self.data_dir, cls)\n",
    "            if not os.path.exists(cls_dir):\n",
    "                continue\n",
    "            for img_name in os.listdir(cls_dir):\n",
    "                img_path = os.path.join(cls_dir, img_name)\n",
    "                if img_path.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
    "                    samples.append((img_path, self.class_to_idx[cls]))\n",
    "        return samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path, label = self.samples[idx]\n",
    "        image = Image.open(img_path).convert('RGB')\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b182052",
   "metadata": {},
   "source": [
    "## 二、数据预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a145a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize((224, 224)),\n",
    "    torchvision.transforms.RandomHorizontalFlip(),\n",
    "    torchvision.transforms.RandomRotation(10),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6217d6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = ChestXrayDataset(\n",
    "    data_dir='data/chest-xray-pneumoniacovid19tuberculosis/train', transform=transform)\n",
    "val_dataset = ChestXrayDataset(\n",
    "    data_dir='data/chest-xray-pneumoniacovid19tuberculosis/val', transform=transform)\n",
    "test_dataset = ChestXrayDataset(\n",
    "    data_dir='data/chest-xray-pneumoniacovid19tuberculosis/test', transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b73347",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size,\n",
    "                          shuffle=True, drop_last=True, num_workers=4)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size,\n",
    "                        shuffle=False, drop_last=True, num_workers=4)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size,\n",
    "                         shuffle=False, drop_last=False, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c06ee38",
   "metadata": {},
   "outputs": [],
   "source": [
    "for images, labels in train_loader:\n",
    "    print(f\"Batch images shape: {images.shape}\")  # 打印图像张量形状\n",
    "    print(f\"Batch labels: {labels}\")  # 打印标签\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42eb7fa1",
   "metadata": {},
   "source": [
    "## 三、定义微调模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6da4965a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fine tuning the resnet34\n",
    "def get_net(device):\n",
    "    finetune_net = nn.Sequential()\n",
    "    finetune_net.features = torchvision.models.resnet152(\n",
    "        weights=torchvision.models.ResNet152_Weights.DEFAULT,\n",
    "    )\n",
    "    finetune_net.output_new = nn.Sequential(\n",
    "        nn.Linear(1000, 256),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(256, 4),\n",
    "    )\n",
    "    finetune_net = finetune_net.to(device)\n",
    "    \n",
    "    for param in finetune_net.features.parameters():\n",
    "        param.requires_grad = False\n",
    "    return finetune_net\n",
    "\n",
    "model = get_net(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0998b7",
   "metadata": {},
   "source": [
    "## 四、模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f4640e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(model, train_loader, optimizer, criterion, device):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    for inputs, labels in train_loader:\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item() * inputs.size(0)\n",
    "        _, predicted = outputs.max(1)\n",
    "        correct += predicted.eq(labels).sum().item()\n",
    "        total += labels.size(0)\n",
    "    \n",
    "    epoch_loss = running_loss / total\n",
    "    epoch_acc = correct / total\n",
    "    return epoch_loss, epoch_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "255b1a2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(model, val_loader, criterion, device):\n",
    "    model.eval()\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in val_loader:\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            \n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            running_loss += loss.item() * inputs.size(0)\n",
    "            _, predicted = outputs.max(1)\n",
    "            correct += predicted.eq(labels).sum().item()\n",
    "            total += labels.size(0)\n",
    "    \n",
    "    epoch_loss = running_loss / total\n",
    "    epoch_acc = correct / total\n",
    "    return epoch_loss, epoch_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c21a2f9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_training_history(history, epoch, epochs):\n",
    "    \"\"\"\n",
    "    动态绘制训练和验证的损失与准确率曲线\n",
    "    :param history: 包含训练和验证指标的字典\n",
    "    :param epoch: 当前训练的轮次\n",
    "    :param epochs: 总训练轮次\n",
    "    \"\"\"\n",
    "    clear_output(wait=True)\n",
    "    plt.figure(figsize=(12, 4))\n",
    "    \n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(history['train_loss'], label='Train Loss')\n",
    "    plt.plot(history['val_loss'], label='Val Loss')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title(f'Epoch {epoch+1}/{epochs}')\n",
    "    plt.legend()\n",
    "    \n",
    "    # 绘制准确率曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(history['train_acc'], label='Train Acc')\n",
    "    plt.plot(history['val_acc'], label='Val Acc')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title(f'Epoch {epoch+1}/{epochs}')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    # plt.savefig(f'assets/epoch_{epoch+1}.png')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08e004fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, train_loader, val_loader, optimizer, criterion, scheduler, device, epochs=10):\n",
    "    best_acc = 0.0\n",
    "    history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion, device)\n",
    "        val_loss, val_acc = validate(model, val_loader, criterion, device)\n",
    "        scheduler.step()\n",
    "        \n",
    "        history['train_loss'].append(train_loss)\n",
    "        history['train_acc'].append(train_acc)\n",
    "        history['val_loss'].append(val_loss)\n",
    "        history['val_acc'].append(val_acc)\n",
    "        \n",
    "        if val_acc > best_acc:\n",
    "            best_acc = val_acc\n",
    "            torch.save(model.state_dict(), 'assets/model/best_model.pth')\n",
    "        \n",
    "        plot_training_history(history, epoch, epochs)\n",
    "        \n",
    "        print(f'Epoch {epoch+1}/{epochs}')\n",
    "        print(f'Train Loss: {train_loss:.4f} Acc: {train_acc:.4f}')\n",
    "        print(f'Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}\\n')\n",
    "        time.sleep(0.1)\n",
    "    \n",
    "    print(f'Best Val Acc: {best_acc:.4f}')\n",
    "    return history\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(\n",
    "    (param for param in model.parameters() if param.requires_grad),  # 仅优化全连接层\n",
    "    lr=1e-3,\n",
    ")\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)\n",
    "\n",
    "history = train_model(\n",
    "    model,\n",
    "    train_loader,\n",
    "    val_loader,\n",
    "    optimizer,\n",
    "    criterion,\n",
    "    scheduler,\n",
    "    device,\n",
    "    epochs=20\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa50786d",
   "metadata": {},
   "source": [
    "## 五、预测与评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc817e72",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, test_loader, criterion, device):\n",
    "    model.eval()\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in test_loader:\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "            \n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            running_loss += loss.item() * inputs.size(0)\n",
    "            _, predicted = outputs.max(1)\n",
    "            correct += predicted.eq(labels).sum().item()\n",
    "            total += labels.size(0)\n",
    "            \n",
    "            all_preds.extend(predicted.cpu().numpy())\n",
    "            all_labels.extend(labels.cpu().numpy())\n",
    "    \n",
    "    test_loss = running_loss / total\n",
    "    test_acc = correct / total\n",
    "    \n",
    "    return test_loss, test_acc, all_preds, all_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5750ee7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading best model\n",
    "model.load_state_dict(torch.load('assets/model/resnet152_finetune_epoch17.pth', weights_only=True))\n",
    "model.eval()\n",
    "\n",
    "# evaluate\n",
    "test_loss, test_acc, all_preds, all_labels = test_model(model, test_loader, criterion, device)\n",
    "print(f'Test Loss: {test_loss:.4f}')\n",
    "print(f'Test Accuracy: {test_acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "069ff195",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 反归一化函数\n",
    "def unnormalize(img, mean, std):\n",
    "    img = img.clone().cpu().numpy().transpose((1, 2, 0))\n",
    "    img = std * img + mean\n",
    "    img = np.clip(img, 0, 1)\n",
    "    return img\n",
    "\n",
    "\n",
    "# 归一化参数\n",
    "mean = np.array([0.485, 0.456, 0.406])\n",
    "std = np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "# 随机采样num_images个测试样本\n",
    "num_images = 6\n",
    "indices = np.random.choice(len(test_dataset), num_images, replace=False)\n",
    "samples = [test_dataset[i] for i in indices]\n",
    "images, labels = zip(*samples)\n",
    "images = torch.stack(images)  # shape: [num_images, C, H, W]\n",
    "labels = torch.tensor(labels)\n",
    "\n",
    "# 模型预测\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    outputs = model(images.to(device))\n",
    "    _, preds = torch.max(outputs, 1)\n",
    "target_names = ['COVID19', 'NORMAL', 'PNEUMONIA', 'TUBERCULOSIS']\n",
    "\n",
    "# 绘制预测结果\n",
    "fig, axes = plt.subplots(2, 3, figsize=(12, 8))\n",
    "axes = axes.flatten()\n",
    "for i in range(num_images):\n",
    "    ax = axes[i]\n",
    "    # 对每张图像进行反归一化处理\n",
    "    img = unnormalize(images[i], mean, std)\n",
    "    ax.imshow(img)\n",
    "    ax.set_title(f\"Pred: {target_names[preds[i].item()]}\\nTrue: {target_names[labels[i].item()]}\")\n",
    "    ax.tick_params(axis='both', which='both', labelsize=10)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bc38794",
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.metrics import classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e3815d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 计算混淆矩阵\n",
    "cm = confusion_matrix(all_labels, all_preds)\n",
    "\n",
    "# 可视化混淆矩阵\n",
    "plt.figure(figsize=(8, 6))\n",
    "sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', \n",
    "            xticklabels=['COVID19', 'NORMAL', 'PNEUMONIA', 'TUBERCULOSIS'], \n",
    "            yticklabels=['COVID19', 'NORMAL', 'PNEUMONIA', 'TUBERCULOSIS'])\n",
    "plt.xlabel('Predicted')\n",
    "plt.ylabel('True')\n",
    "plt.title('Confusion Matrix')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b1994e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印分类报告\n",
    "print(classification_report(\n",
    "    all_labels,\n",
    "    all_preds,\n",
    "    target_names=['COVID19', 'NORMAL', 'PNEUMONIA', 'TUBERCULOSIS']\n",
    "))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}