2 XDTLMI-Net training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a2e0e92",
   "metadata": {},
   "source": [
    "# Transfer learning Fine-tuning on the `train_dataset`\n",
    "\n",
    "This notebook demonstrates how a pre-trained CNN can be fine-tuned using randomly selected train splits. Here, we'll demonstrate fine-tuning a pre-trained CNN on the medical image classification task. In this example, fine-tuning serves to update the CNN to new measurement parameters. This code illustrates the procedure described in the `3.Materials and methods`. \n",
    "\n",
    "Copyright (C) 2023, Zhao Bingqiang, All Rights Reserved\n",
    "\n",
    "Email: zbqherb@163.com\n",
    "\n",
    "2023-07-02"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b8e5fc",
   "metadata": {},
   "source": [
    "## Import Libs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36670c66",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "import wandb\n",
    "import random\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from datetime import datetime\n",
    "from matplotlib import colors as mcolors\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torchvision import datasets\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from torchvision import models\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "\n",
    "from sklearn.metrics import precision_score\n",
    "from sklearn.metrics import recall_score\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.metrics import f1_score\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "print('device', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77bb8dac",
   "metadata": {},
   "source": [
    "## Image data preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d84a5ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COVID-19 CT\n",
    "train_transform = transforms.Compose([transforms.RandomResizedCrop(512),\n",
    "                                      transforms.RandomHorizontalFlip(),\n",
    "                                      transforms.ToTensor(),\n",
    "                                      transforms.Normalize((0.5,), (0.5,))])\n",
    "\n",
    "test_transform = transforms.Compose([transforms.Resize(1000),\n",
    "                                     transforms.CenterCrop(512),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize((0.5,), (0.5,))])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d827c0a9",
   "metadata": {},
   "source": [
    "## Load image dataset `data_split`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f25e787",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dir = 'data_split'\n",
    "train_path = os.path.join(dataset_dir, 'train')\n",
    "test_path = os.path.join(dataset_dir, 'val')\n",
    "print('train_dataset Path', train_path)\n",
    "print('test_dataset Path', test_path)\n",
    "\n",
    "# Load train_dataset\n",
    "train_dataset = datasets.ImageFolder(train_path, train_transform)\n",
    "# Load test_dataset\n",
    "test_dataset = datasets.ImageFolder(test_path, test_transform)\n",
    "\n",
    "print('train_dataset number', len(train_dataset))\n",
    "print('train_dataset class number', len(train_dataset.classes))\n",
    "print('train_dataset class name', train_dataset.classes)\n",
    "\n",
    "print('test_dataset number', len(test_dataset))\n",
    "print('test_dataset class number', len(test_dataset.classes))\n",
    "print('test_dataset class name', test_dataset.classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5508d5fb",
   "metadata": {},
   "source": [
    "## Class and Index mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51dd6a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping：Class to Index\n",
    "class_to_idx = train_dataset.class_to_idx\n",
    "# Mapping：Index to Class \n",
    "idx_to_labels = {y:x for x,y in train_dataset.class_to_idx.items()}\n",
    "# Save mapping files\n",
    "np.save('table/idx_to_labels.npy', idx_to_labels)\n",
    "np.save('table/labels_to_idx.npy', class_to_idx)\n",
    "\n",
    "print(class_to_idx)\n",
    "print(idx_to_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2147b31",
   "metadata": {},
   "source": [
    "## Define DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b91ce4e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32\n",
    "\n",
    "# train_dataset DataLoader\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size = BATCH_SIZE,\n",
    "                          shuffle = True,\n",
    "                          num_workers = 4\n",
    "                         )\n",
    "\n",
    "# test_dataset DataLoader\n",
    "test_loader = DataLoader(test_dataset,\n",
    "                         batch_size = BATCH_SIZE,\n",
    "                         shuffle = False,\n",
    "                         num_workers = 4\n",
    "                        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f27a39ad",
   "metadata": {},
   "source": [
    "## Visualize the image and annotation of a batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c0cb00a",
   "metadata": {},
   "outputs": [],
   "source": [
    "images, labels = next(iter(train_loader))\n",
    "print(images.shape)\n",
    "print(labels)\n",
    "\n",
    "# Tensor to np.array\n",
    "images = images.numpy()\n",
    "plt.figure(figsize = (15,10))\n",
    "plt.hist(images[10].flatten(), bins = 100)\n",
    "plt.tick_params(labelsize = 25)\n",
    "plt.savefig('figure/hist.tif', dpi = 300, bbox_inches = 'tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d0bfcef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocessed images in the batch \n",
    "idx = 1\n",
    "label = labels[idx].item()\n",
    "# Preprocessed image\n",
    "plt.subplot(121)\n",
    "plt.imshow(images[idx].transpose((1,2,0)))\n",
    "plt.axis('off')\n",
    "plt.title('Preprocessed:'+ idx_to_labels[label], fontsize = 10)\n",
    "\n",
    "# Original image\n",
    "plt.subplot(122)\n",
    "mean = np.array((0.5,))\n",
    "std = np.array((0.5,))\n",
    "# mean = np.array([0.485, 0.456, 0.406])\n",
    "# std = np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "plt.imshow(np.clip(images[idx].transpose((1,2,0)) * std + mean, 0, 1))\n",
    "plt.axis('off')\n",
    "plt.title('Original:'+ idx_to_labels[label], fontsize = 10)\n",
    "plt.savefig('figure/image visualization2.tif', dpi = 300, bbox_inches = 'tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a16b3612",
   "metadata": {},
   "source": [
    "## Transfer learning fine-tuning options\n",
    "\n",
    "https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html\n",
    "\n",
    "Now we set up a ResNet CNN and load weights that previously trained for the `ImageNet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b3bf5cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_class = len(train_dataset.classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48a7504f",
   "metadata": {},
   "source": [
    "### 1. Fine-tuning fully connected layer only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8325576",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Load pre_trained image classification model\n",
    "# model = models.resnet18(pretrained = True) \n",
    "# # New layer default (requires_grad = True)\n",
    "# model.fc = nn.Linear(model.fc.in_features, n_class)\n",
    "# optimizer = optim.Adam(model.fc.parameters())\n",
    "\n",
    "# print(model.fc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d0c8cdf",
   "metadata": {},
   "source": [
    "### 2. Fine-tunning all layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c1d5846",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.resnet18(pretrained = True)\n",
    "model.fc = nn.Linear(model.fc.in_features, n_class)\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "print(model.fc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69bf56ae",
   "metadata": {},
   "source": [
    "### 3. Initialize all model weights randomly and train all layers from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f688529",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Only the model structure is loaded, not the pre-training weight parameters\n",
    "# model = models.resnet18(pretrained=False) \n",
    "# model.fc = nn.Linear(model.fc.in_features, n_class)\n",
    "# optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c11a04",
   "metadata": {},
   "source": [
    "## Model training Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fedf17f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(device)\n",
    "criterion = nn.CrossEntropyLoss() \n",
    "\n",
    "EPOCHS = 100\n",
    "\n",
    "lr_scheduler = lr_scheduler.StepLR(optimizer, step_size = 50, gamma = 0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91bf867e",
   "metadata": {},
   "source": [
    "## Function: Train on the `train_dataset`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18d25776",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_one_batch(images, labels):\n",
    "    \n",
    "    '''\n",
    "    Train a batch, returns the Training Log of the current batch\n",
    "    '''\n",
    "    \n",
    "    images = images.to(device)\n",
    "    labels = labels.to(device)\n",
    "    \n",
    "    # Forward propagation\n",
    "    outputs = model(images) \n",
    "    # Calculate the average cross-entropy loss function value \n",
    "    # of each sample in the current batch\n",
    "    loss = criterion(outputs, labels) \n",
    "    \n",
    "    # Back propagation, optimize and update the weight\n",
    "    optimizer.zero_grad() # gradient to zero\n",
    "    loss.backward() # Back propagation\n",
    "    optimizer.step() # Update parameters\n",
    "    \n",
    "    # Gets the label ID and predicted ID for the current batch\n",
    "    # Gets the predicted ID for all images in the current batch\n",
    "    _, preds = torch.max(outputs, 1) \n",
    "    preds = preds.cpu().numpy()\n",
    "    loss = loss.detach().cpu().numpy()\n",
    "    outputs = outputs.detach().cpu().numpy()\n",
    "    labels = labels.detach().cpu().numpy()\n",
    "    \n",
    "    log_train = {}\n",
    "    log_train['epoch'] = epoch\n",
    "    log_train['batch'] = batch_idx\n",
    "    # Classification evaluation index on train_dataset\n",
    "    log_train['train_loss'] = loss\n",
    "    log_train['train_accuracy'] = accuracy_score(labels, preds)\n",
    "    log_train['train_precision'] = precision_score(labels, preds, average = 'macro')\n",
    "    log_train['train_recall'] = recall_score(labels, preds, average = 'macro')\n",
    "    log_train['train_f1-score'] = f1_score(labels, preds, average = 'macro')\n",
    "    \n",
    "    return log_train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b6628f0",
   "metadata": {},
   "source": [
    "## Function: Evaluate on the `test_dataset`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc8d27ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_testset():\n",
    "    '''\n",
    "    Evaluate the test_dataset, returns the Test log of current epoch\n",
    "    '''\n",
    "\n",
    "    loss_list = []\n",
    "    labels_list = []\n",
    "    preds_list = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader: # Generate a batch of data and annotations\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images) # forward prediction\n",
    "\n",
    "            # Obtain the label and predicted ID for the test_dataset\n",
    "            # Gets the predicted ID for all images in the current batch\n",
    "            _, preds = torch.max(outputs, 1) \n",
    "            preds = preds.cpu().numpy()\n",
    "            # Calculate the average cross-entropy loss function value \n",
    "            # of each sample in the current batch via logit\n",
    "            loss = criterion(outputs, labels) \n",
    "            loss = loss.detach().cpu().numpy()\n",
    "            outputs = outputs.detach().cpu().numpy()\n",
    "            labels = labels.detach().cpu().numpy()\n",
    "\n",
    "            loss_list.append(loss)\n",
    "            labels_list.extend(labels)\n",
    "            preds_list.extend(preds)\n",
    "        \n",
    "    log_test = {}\n",
    "    log_test['epoch'] = epoch\n",
    "    \n",
    "    # Classification evaluation index on test_dataset\n",
    "    log_test['test_loss'] = np.mean(loss_list)\n",
    "    log_test['test_accuracy'] = accuracy_score(labels_list, preds_list)\n",
    "    log_test['test_precision'] = precision_score(labels_list, preds_list, average = 'macro')\n",
    "    log_test['test_recall'] = recall_score(labels_list, preds_list, average = 'macro')\n",
    "    log_test['test_f1-score'] = f1_score(labels_list, preds_list, average = 'macro')\n",
    "    \n",
    "    return log_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88b53698",
   "metadata": {},
   "source": [
    "## Log recoed before training starts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "819f5f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 0\n",
    "batch_idx = 0\n",
    "best_test_accuracy = 0\n",
    "\n",
    "# log_train - train_dataset\n",
    "df_train_log = pd.DataFrame()\n",
    "log_train = {}\n",
    "log_train['epoch'] = 0\n",
    "log_train['batch'] = 0\n",
    "images, labels = next(iter(train_loader))\n",
    "log_train.update(train_one_batch(images, labels))\n",
    "df_train_log = df_train_log._append(log_train, ignore_index = True)\n",
    "\n",
    "df_train_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e9fa9f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# log_train - test_dataset\n",
    "df_test_log = pd.DataFrame()\n",
    "log_test = {}\n",
    "log_test['epoch'] = 0\n",
    "log_test.update(evaluate_testset())\n",
    "df_test_log = df_test_log._append(log_test, ignore_index = True)\n",
    "\n",
    "df_test_log"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "156fccfe",
   "metadata": {},
   "source": [
    "## wandb visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b62659",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# wandb.init(project = 'COVID', name = time.strftime('%m%d%H%M%S'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abc8e14c",
   "metadata": {},
   "source": [
    "## Train start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8539d137",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for epoch in range(1, EPOCHS+1):\n",
    "    \n",
    "    print(f'{datetime.now()}, Epoch {epoch}/{EPOCHS}')\n",
    "    \n",
    "    ## Train phase\n",
    "    model.train()\n",
    "    for images, labels in tqdm(train_loader): \n",
    "        batch_idx += 1\n",
    "        log_train = train_one_batch(images, labels)\n",
    "        df_train_log = df_train_log._append(log_train, ignore_index = True)\n",
    "#         wandb.log(log_train)\n",
    "        \n",
    "    lr_scheduler.step()\n",
    "\n",
    "    ## Test phase\n",
    "    model.eval()\n",
    "    log_test = evaluate_testset()\n",
    "    df_test_log = df_test_log._append(log_test, ignore_index = True)\n",
    "#     wandb.log(log_test)\n",
    "    \n",
    "    # Save the latest best model file\n",
    "    if log_test['test_accuracy'] > best_test_accuracy: \n",
    "        # Delete old best model files (if any)\n",
    "        old_best_checkpoint_path = 'checkpoint/best-{:.3f}.pth'.format(best_test_accuracy)\n",
    "        if os.path.exists(old_best_checkpoint_path):\n",
    "            os.remove(old_best_checkpoint_path)\n",
    "        # Save the new best model file\n",
    "        best_test_accuracy = log_test['test_accuracy']\n",
    "        new_best_checkpoint_path = 'checkpoint/best-{:.3f}.pth'.format(log_test['test_accuracy'])\n",
    "        torch.save(model, new_best_checkpoint_path)\n",
    "        print('save best model', 'checkpoint/best-{:.3f}.pth'.format(best_test_accuracy))\n",
    "            \n",
    "\n",
    "df_train_log.to_csv('table/Train Log-train_dataset.csv', index = False)\n",
    "df_test_log.to_csv('table/Train Log-test_dataset.csv', index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea1b5581",
   "metadata": {},
   "source": [
    "## Evaluate on `test_dataset`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45d4d8c7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = torch.load('checkpoint/best-{:.3f}.pth'.format(best_test_accuracy))\n",
    "\n",
    "model.eval()\n",
    "print(evaluate_testset())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "919b9e8a",
   "metadata": {},
   "source": [
    "## Visualize Train Log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a9b9be",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv('table/Train Log-train_dataset.csv')\n",
    "df_test = pd.read_csv('table/Train Log-test_dataset.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68a09405",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90e47c11",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "034b9eb1",
   "metadata": {},
   "source": [
    "### `train_dataset` loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46299fd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (15, 10))\n",
    "\n",
    "x = df_train['batch']\n",
    "y = df_train['train_loss']\n",
    "\n",
    "plt.plot(x, y, linewidth = 2)\n",
    "\n",
    "plt.tick_params(labelsize = 25)\n",
    "plt.xlabel('Batch', fontsize = 25)\n",
    "plt.ylabel('Train Loss', fontsize = 25)\n",
    "# plt.title('train_dataset Loss', fontsize=25)\n",
    "plt.savefig('figure/train_dataset Loss.tif', dpi = 300, bbox_inches = 'tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81c13078",
   "metadata": {},
   "source": [
    "### `train_dataset` accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f379d850",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (15, 10))\n",
    "\n",
    "x = df_train['batch']\n",
    "y = df_train['train_accuracy']\n",
    "\n",
    "plt.plot(x, y, linewidth = 2)\n",
    "\n",
    "plt.tick_params(labelsize = 25)\n",
    "plt.xlabel('Batch', fontsize = 25)\n",
    "plt.ylabel('Train Accuracy', fontsize = 25)\n",
    "# plt.title('train_dataset Accuracy', fontsize=25)\n",
    "plt.savefig('figure/train_dataset Accuracy.tif', dpi = 300, bbox_inches = 'tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c151197",
   "metadata": {},
   "source": [
    "### `test_dataset` loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c279a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (15, 10))\n",
    "\n",
    "x = df_test['epoch']\n",
    "y = df_test['test_loss']\n",
    "\n",
    "plt.plot(x, y, linewidth = 2.5)\n",
    "\n",
    "plt.tick_params(labelsize = 25)\n",
    "plt.xlabel('Epoch', fontsize = 25)\n",
    "plt.ylabel('Test Loss', fontsize = 25)\n",
    "# plt.title('test_dataset Loss', fontsize=25)\n",
    "plt.savefig('figure/test_dataset Loss.tif', dpi = 300, bbox_inches = 'tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5105acbe",
   "metadata": {},
   "source": [
    "### `test_dataset` Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "474f8207",
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(222)\n",
    "colors = ['tab:pink', 'tab:gray', 'tab:olive', 'tab:cyan']\n",
    "linestyle = ['--', '-.', '-']\n",
    "\n",
    "def get_line_arg():\n",
    "    line_arg = {}\n",
    "    line_arg['color'] = random.choice(colors)\n",
    "    line_arg['linestyle'] = random.choice(linestyle)\n",
    "    line_arg['linewidth'] = 2\n",
    "    return line_arg\n",
    "\n",
    "metrics = ['test_accuracy', 'test_precision', 'test_recall', 'test_f1-score']\n",
    "\n",
    "plt.figure(figsize = (15, 10))\n",
    "\n",
    "x = df_test['epoch']\n",
    "for y in metrics:\n",
    "    plt.plot(x, df_test[y], label = y, **get_line_arg())\n",
    "\n",
    "plt.tick_params(labelsize = 25)\n",
    "\n",
    "plt.ylim([0, 1.05])\n",
    "plt.xlabel('Epoch', fontsize = 25)\n",
    "plt.ylabel('Test_Metrics', fontsize = 25)\n",
    "plt.legend(loc = 4, fontsize = 20, frameon = False)\n",
    "plt.savefig('figure/test_dataset performance.tif', dpi = 300, bbox_inches = 'tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3542c0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe7bb008",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ImageClassification",
   "language": "python",
   "name": "imageclassification"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}