COVID-19.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import torch\n",
    "\n",
    "print(sys.version)\n",
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi -L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下载Kaggle数据集\n",
    "!cd ~\n",
    "!pwd -P\n",
    "\n",
    "import json\n",
    "token = {\"username\":\"xxxxxxxx\",\"key\":\"xxxxxxxx\"}\n",
    "with open('./content/kaggle.json', 'w') as file:\n",
    "    json.dump(token, file)\n",
    "    \n",
    "!pip install kaggle\n",
    "!mkdir -p ~/.kaggle\n",
    "!cp ./content/kaggle.json ~/.kaggle/\n",
    "!chmod 600 ~/.kaggle/kaggle.json\n",
    "!kaggle config set -n path -v ./content\n",
    "!kaggle datasets download -d hgunraj/covidxct\n",
    "!unzip ./content/datasets/hgunraj/covidxct/covidxct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import copy\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import torch,torchvision\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets,models,transforms\n",
    "import torch.optim as optim\n",
    "# from torchsummary import summary\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.utils import resample\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from torch.nn import Parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install torchsummary\n",
    "path='./densenet201/densenet201_3.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Iterable\n",
    "\n",
    "def set_freeze_by_idxs(model, idxs, freeze=True): #默认freeze=True是冻结层\n",
    "    if not isinstance(idxs, Iterable):\n",
    "        idxs = [idxs]\n",
    "    num_child = len(list(model.children()))\n",
    "    idxs = tuple(map(lambda idx: num_child + idx if idx < 0 else idx, idxs))\n",
    "    for idx, child in enumerate(model.children()):\n",
    "        if idx not in idxs:\n",
    "            continue\n",
    "        for param in child.parameters():\n",
    "            param.requires_grad = not freeze\n",
    "    return model\n",
    "            \n",
    "def freeze_by_idxs(model, idxs):\n",
    "    return set_freeze_by_idxs(model, idxs, True)\n",
    "\n",
    "def unfreeze_by_idxs(model, idxs):\n",
    "    return set_freeze_by_idxs(model, idxs, False)\n",
    "def set_parameter_requires_grad(model):#冻结特征提取层\n",
    "    for param in model.parameters():\n",
    "        param.requires_grad=False\n",
    "    return model\n",
    "\n",
    "def initialize_model(model_name,num_classes,use_pretrained,unfreeze_num):\n",
    "    if model_name=='vgg16':\n",
    "        model_pre=models.vgg16(pretrained=use_pretrained) # True：torchversion 已经预训练好的模型，False：随机初始化的模型\n",
    "#         for i in model_pre.named_children():\n",
    "#             print(i)\n",
    "        model_pre.features[0].in_channels=1\n",
    "        model_pre.features[0].weight=Parameter(model_pre.features[0].weight[:,1:2,:,:])\n",
    "        model_pre=set_parameter_requires_grad(model_pre) # 冻结全部层\n",
    "        num_ftrs=model_pre.classifier[6].in_features # feature_map 的大小\n",
    "        model_pre.classifier[6]=nn.Linear(num_ftrs,num_classes) #重新设计全连接层\n",
    "        if unfreeze_num==1:\n",
    "            unfreeze=[-1]\n",
    "        elif unfreeze_num==2:\n",
    "            unfreeze=[-1,-3]\n",
    "        elif unfreeze_num==3:\n",
    "            unfreeze=[-1,-3,-5]\n",
    "        else:\n",
    "            unfreeze=[-1,-3,-5,-7]\n",
    "        model_pre.features=unfreeze_by_idxs(model_pre.features,unfreeze)\n",
    "        for param in model_pre.classifier.parameters():\n",
    "            param.requires_grad=True\n",
    "        input_size=224\n",
    "    elif model_name=='vgg19':\n",
    "        model_pre=models.vgg19(pretrained=use_pretrained) # True：torchversion 已经预训练好的模型，False：随机初始化的模型\n",
    "        model_pre.features[0].in_channels=1\n",
    "        model_pre.features[0].weight=Parameter(model_pre.features[0].weight[:,1:2,:,:])\n",
    "        model_pre=set_parameter_requires_grad(model_pre) # 冻结全部层\n",
    "        num_ftrs=model_pre.classifier[6].in_features # feature_map 的大小\n",
    "        model_pre.classifier[6]=nn.Linear(num_ftrs,num_classes) #重新设计全连接层\n",
    "        if unfreeze_num==1:\n",
    "            unfreeze=[-1]\n",
    "        elif unfreeze_num==2:\n",
    "            unfreeze=[-1,-3]\n",
    "        elif unfreeze_num==3:\n",
    "            unfreeze=[-1,-3,-5]\n",
    "        else:\n",
    "            unfreeze=[-1,-3,-5,-7]\n",
    "        model_pre.features=unfreeze_by_idxs(model_pre.features,unfreeze)\n",
    "        for param in model_pre.classifier.parameters():\n",
    "            param.requires_grad=True\n",
    "        input_size=224\n",
    "    elif model_name=='resnet101':\n",
    "        model_pre=models.resnet101(pretrained=use_pretrained) # True：torchversion 已经预训练好的模型，False：随机初始化的模型\n",
    "        model_pre.conv1.in_channels=1\n",
    "        model_pre.conv1.weight=Parameter(model_pre.conv1.weight[:,1:2,:,:])\n",
    "        model_pre=set_parameter_requires_grad(model_pre)\n",
    "        num_ftrs=model_pre.fc.in_features # feature_map 的大小\n",
    "        model_pre.fc=nn.Linear(num_ftrs,num_classes)\n",
    "        \n",
    "        for i in range(unfreeze_num):\n",
    "            model_pre.layer4=unfreeze_by_idxs(model_pre.layer4,-i)\n",
    "        for param in model_pre.fc.parameters():\n",
    "            param.requires_grad=True\n",
    "        input_size=224\n",
    "    elif model_name=='resnet152':\n",
    "        model_pre=models.resnet152(pretrained=use_pretrained) #True：torchversion 已经预训练好的模型，False：随机初始化的模型\n",
    "        model_pre.conv1.in_channels=1\n",
    "        model_pre.conv1.weight=Parameter(model_pre.conv1.weight[:,1:2,:,:])\n",
    "        model_pre=set_parameter_requires_grad(model_pre)\n",
    "        num_ftrs=model_pre.fc.in_features #feature_map 的大小\n",
    "        model_pre.fc=nn.Linear(num_ftrs,num_classes)\n",
    "\n",
    "        for i in range(unfreeze_num):\n",
    "            model_pre.layer4=unfreeze_by_idxs(model_pre.layer4,-i)\n",
    "        for param in model_pre.fc.parameters():\n",
    "            param.requires_grad=True\n",
    "        input_size=224\n",
    "    elif model_name=='densenet161':\n",
    "        model_pre=models.densenet161(pretrained=use_pretrained) #True：torchversion 已经预训练好的模型，False：随机初始化的模型\n",
    "        model_pre.features[0].in_channels=1\n",
    "        model_pre.features[0].weight=Parameter(model_pre.features[0].weight[:,1:2,:,:])\n",
    "        model_pre=set_parameter_requires_grad(model_pre)\n",
    "        num_ftrs=model_pre.classifier.in_features #feature_map 的大小\n",
    "        model_pre.classifier=nn.Linear(num_ftrs,num_classes)\n",
    "        \n",
    "        for i in range(unfreeze_num):\n",
    "            model_pre.features.denseblock4=unfreeze_by_idxs(model_pre.features.denseblock4,-i)\n",
    "        for param in model_pre.classifier.parameters():\n",
    "            param.requires_grad=True\n",
    "        input_size=224\n",
    "    elif model_name=='densenet201':\n",
    "        model_pre=models.densenet201(pretrained=use_pretrained) #True：torchversion 已经预训练好的模型，False：随机初始化的模型\n",
    "        model_pre.features[0].in_channels=1\n",
    "        model_pre.features[0].weight=Parameter(model_pre.features[0].weight[:,1:2,:,:])\n",
    "        model_pre=set_parameter_requires_grad(model_pre)\n",
    "        num_ftrs=model_pre.classifier.in_features #feature_map 的大小\n",
    "        model_pre.classifier=nn.Linear(num_ftrs,num_classes)\n",
    "        \n",
    "        for i in range(unfreeze_num):\n",
    "            model_pre.features.denseblock4=unfreeze_by_idxs(model_pre.features.denseblock4,-i)\n",
    "        for param in model_pre.classifier.parameters():\n",
    "            param.requires_grad=True\n",
    "        input_size=224\n",
    "    else:\n",
    "        print('model not implemented')\n",
    "        return None,None\n",
    "    return model_pre,input_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "model_all=[]\n",
    "dir='.'\n",
    "def auto_net(model_name,num_classes,use_pretrained,unfreeze_num):\n",
    "    \n",
    "    for k in range(unfreeze_num):\n",
    "        model,input_size=initialize_model(model_name,num_classes,use_pretrained,k+1)\n",
    "        my_path=Path(dir+'/{}'.format(model_name))\n",
    "        if not my_path.is_dir():    # 若指定的目录不存在，创建文件夹\n",
    "            os.mkdir(my_path)\n",
    "        torch.save(model,dir+'/{}/{}_{}.pth'.format(model_name,model_name,k)) # 0 1 2 3\n",
    "            \n",
    "        model_all.append(model)\n",
    "#         print('------------------------删除{}层,添加{}层,倒数{}层求梯度------------------------'.format(i+1,j+1,k+1))\n",
    "#         for m in model.named_children():\n",
    "#             print(m)\n",
    "#         print('----------------')\n",
    "#         for n in model.named_parameters():\n",
    "#             print(n)\n",
    "    return model_all\n",
    "model_name=['vgg16','vgg19','resnet101','resnet152','densenet161','densenet201'] # 'inception_v3','resnext101'\n",
    "# for name in model_name:\n",
    "#     model_all=auto_net(name,num_classes=2,use_pretrained=True,unfreeze_num=4)\n",
    "model_all=auto_net(model_name[5],num_classes=3,use_pretrained=True,unfreeze_num=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取train.txt\n",
    "train_df = pd.read_csv('./train_COVIDx_CT-2A.txt', sep=\" \", header=None)\n",
    "train_df.columns=['filename', 'label', 'xmin','ymin','xmax','ymax']\n",
    "train_df=train_df.drop(['xmin', 'ymin','xmax', 'ymax'], axis=1 )\n",
    "# 读取test.txt\n",
    "val_df = pd.read_csv('./val_COVIDx_CT-2A.txt', sep=\" \", header=None)\n",
    "val_df.columns=['filename', 'label', 'xmin','ymin','xmax','ymax']\n",
    "val_df=val_df.drop(['xmin', 'ymin','xmax', 'ymax'], axis=1 )\n",
    "\n",
    "test_df = pd.read_csv('./test_COVIDx_CT-2A.txt', sep=\" \", header=None)\n",
    "test_df.columns=['filename', 'label', 'xmin','ymin','xmax','ymax']\n",
    "test_df=test_df.drop(['xmin', 'ymin','xmax', 'ymax'], axis=1 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# labels={0:'Normal',1:'Pneumonia',2:'COVID-19'}\n",
    "train_df.head()\n",
    "train_df.label.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = './2A_images/'  #directory path\n",
    "train_df['filename'] = image_path+train_df['filename']\n",
    "val_df['filename'] = image_path+val_df['filename']\n",
    "test_df['filename'] = image_path + test_df['filename']\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N  = train_df[train_df['label']==0]\n",
    "P = train_df[train_df['label']==1]\n",
    "C = train_df[train_df['label']==2]\n",
    "from sklearn.utils import resample\n",
    "N_download = resample(N, replace = True, n_samples = 25496,random_state=0)\n",
    "C_download = resample(C, replace = True, n_samples = 25496,random_state=0)\n",
    "train_df = pd.concat([N_download, P, C_download])\n",
    "train_df.label.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_v  = val_df[val_df['label']==0]\n",
    "P_v = val_df[val_df['label']==1]\n",
    "C_v = val_df[val_df['label']==2]\n",
    "from sklearn.utils import resample\n",
    "N_v_download = resample(N_v, replace = True, n_samples = 6244,random_state=0)\n",
    "P_v_download = resample(P_v, replace = True, n_samples = 6244,random_state=0)\n",
    "val_df = pd.concat([N_v_download, P_v_download, C_v])\n",
    "val_df.label.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = shuffle(train_df) # 打乱顺序\n",
    "val_df = shuffle(val_df)\n",
    "test_df = shuffle(test_df)\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels={0:'Normal',1:'Pneumonia',2:'COVID-19'}\n",
    "class_names=['Normal','Pneumonia','COVID-19']\n",
    "\n",
    "train_df['label_n']=[labels[b] for b in train_df['label']]\n",
    "val_df['label_n']=[labels[b] for b in val_df['label']]\n",
    "test_df['label_n']=[labels[b] for b in test_df['label']]\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Negative and positive values of train: \\n{train_df['label_n'].value_counts()}\")\n",
    "print(f\"Negative and positive values of validation: \\n{val_df['label_n'].value_counts()}\")\n",
    "print(f\"Negative and positive values of test: \\n{test_df['label_n'].value_counts()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df=train_df.reset_index()\n",
    "val_df=val_df.reset_index()\n",
    "test_df=test_df.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CovidDataset(Dataset):\n",
    "    def __init__(self, dataset_df, transform=None):\n",
    "        self.dataset_df = dataset_df\n",
    "        self.transform = transform\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.dataset_df.shape[0]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        image_name = self.dataset_df['filename'][idx]\n",
    "        img = Image.open(image_name)\n",
    "        label = self.dataset_df['label'][idx]\n",
    "        \n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "        return img, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "input_channel = 1\n",
    "input_size = (224,224)\n",
    "crop_size=(340,380)\n",
    "num_classes=3\n",
    "num_epochs = 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = {\n",
    "    'train':transforms.Compose([\n",
    "        transforms.RandomHorizontalFlip(p=0.5),\n",
    "        transforms.RandomVerticalFlip(p=0.5),\n",
    "        transforms.RandomRotation(30),\n",
    "        transforms.CenterCrop(crop_size),\n",
    "        transforms.Resize(input_size),\n",
    "        transforms.Grayscale(input_channel),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.6349431],[0.32605055])\n",
    "    ]),\n",
    "    'test':transforms.Compose([\n",
    "        transforms.CenterCrop(crop_size),\n",
    "        transforms.Resize(input_size),\n",
    "        transforms.Grayscale(input_channel),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.63507175],[0.3278614])\n",
    "    ])\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_names=['train','val','test']\n",
    "image_transforms = {'train':transform['train'], 'val':transform['test'],'test':transform['test']}\n",
    "\n",
    "train_dataset = CovidDataset(train_df, transform=image_transforms['train'])\n",
    "val_dataset = CovidDataset(val_df, transform=image_transforms['val'])\n",
    "test_dataset = CovidDataset(test_df, transform=image_transforms['test'])\n",
    "\n",
    "image_dataset = {'train':train_dataset, 'val':val_dataset,'test':test_dataset}\n",
    "\n",
    "dataloaders = {x:DataLoader(image_dataset[x],batch_size=batch_size,shuffle=True,num_workers=4) for x in dataset_names}\n",
    "\n",
    "dataset_sizes = {x:len(image_dataset[x]) for x in dataset_names}\n",
    "\n",
    "# class_names = {}\n",
    "# print(dataset_sizes)\n",
    "# print(class_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "def show_tensor_img(tensor_img):\n",
    "#     to_pil = transforms.ToPILImage()\n",
    "#     img = tensor_img.cpu().clone()\n",
    "#     img = to_pil(img)\n",
    "    img=transforms.ToPILImage()(tensor_img)\n",
    "    plt.figure()\n",
    "    plt.imshow(img,plt.cm.gray)\n",
    "    plt.show()\n",
    "\n",
    "def show_img(idx):\n",
    "    show_tensor_img(train_dataset[idx][0])\n",
    "for i in range(4):\n",
    "    show_img(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "# 绘制混淆矩阵\n",
    "def plot_confusion_matrix(cm, classes, normalize=False, title='Confusion matrix', cmap=plt.cm.Blues):\n",
    "    \"\"\"\n",
    "    This function prints and plots the confusion matrix.\n",
    "    Normalization can be applied by setting `normalize=True`.\n",
    "    Input\n",
    "    - cm : 计算出的混淆矩阵的值\n",
    "    - classes : 混淆矩阵中每一行每一列对应的列\n",
    "    - normalize : True:显示百分比, False:显示个数\n",
    "    \"\"\"\n",
    "    cm=cm.numpy()\n",
    "    if normalize:\n",
    "        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]\n",
    "        print(\"Normalized confusion matrix\")\n",
    "    else:\n",
    "        cm=cm.astype('int')\n",
    "        print('Confusion matrix, without normalization')\n",
    "#     print(cm)\n",
    "    plt.imshow(cm, interpolation='nearest', cmap=cmap)\n",
    "    plt.title(title)\n",
    "    plt.colorbar()\n",
    "    tick_marks = np.arange(len(classes))\n",
    "    plt.xticks(tick_marks, classes, rotation=45)\n",
    "    plt.yticks(tick_marks, classes)\n",
    "    fmt = '{:.2f}' if normalize else '{}'\n",
    "    thresh = cm.max() / 2.\n",
    "    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):\n",
    "      plt.text(i, j, fmt.format(cm[i, j]),horizontalalignment=\"center\",color=\"white\" if cm[i, j] > thresh else \"black\")\n",
    "    plt.tight_layout()\n",
    "    plt.ylabel('True label')\n",
    "    plt.xlabel('Predicted label')\n",
    "    plt.savefig('./confusion_matrix/' + title + '.png')\n",
    "    plt.show()\n",
    "\n",
    "# 生成混淆矩阵\n",
    "def confusion_matrix(preds, labels, conf_matrix):\n",
    "    preds = torch.argmax(preds, 1)\n",
    "    for p, t in zip(preds, labels):\n",
    "        conf_matrix[t, p] += 1\n",
    "    return conf_matrix\n",
    "\n",
    "def calculate_all_prediction(conf_matrix):\n",
    "    '''\n",
    "    计算总精度：对角线上所有值除以总数\n",
    "    '''\n",
    "    total_sum = conf_matrix.sum()\n",
    "    correct_sum = (np.diag(conf_matrix)).sum()\n",
    "    prediction = round(100*float(correct_sum)/float(total_sum),2)\n",
    "    return prediction\n",
    " \n",
    "def calculate_label_prediction(conf_matrix,labelidx):\n",
    "    '''\n",
    "    计算某一个类标预测精度：该类被预测正确的数除以该类的总数\n",
    "    '''\n",
    "    label_total_sum = conf_matrix.sum(axis=0)[labelidx]\n",
    "    label_correct_sum = conf_matrix[labelidx][labelidx]\n",
    "    prediction = 0\n",
    "    if label_total_sum != 0:\n",
    "        prediction = round(100*float(label_correct_sum)/float(label_total_sum),2)\n",
    "    return prediction\n",
    " \n",
    "def calculate_label_recall(conf_matrix,labelidx):\n",
    "    '''\n",
    "    计算某一个类标的召回率：\n",
    "    '''\n",
    "    label_total_sum = conf_matrix.sum(axis=1)[labelidx]\n",
    "    label_correct_sum = conf_matrix[labelidx][labelidx]\n",
    "    recall = 0\n",
    "    if label_total_sum != 0:\n",
    "        recall = round(100*float(label_correct_sum)/float(label_total_sum),2)\n",
    "    return recall\n",
    " \n",
    "def calculate_f1(prediction,recall):\n",
    "    if (prediction+recall)==0:\n",
    "        return 0\n",
    "    return round(2*prediction*recall/(prediction+recall),2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# path='./vgg16/vgg16_1.pth'\n",
    "model=torch.load(path)\n",
    "model=model.to(device)\n",
    "# print(model)\n",
    "\n",
    "# model=models.vgg16(pretrained=True)\n",
    "# # 将所有参数都设置为不计算梯度\n",
    "# for param in model.parameters():\n",
    "#     param.requires_grad=False\n",
    "# num_ftrs=model.classifier[6].in_features # feature_map 的大小\n",
    "# model.classifier[6]=nn.Linear(num_ftrs,num_classes) #重新设计全连接层\n",
    "# model=model.to(device)\n",
    "\n",
    "criterion=nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(),lr=0.0001,betas=(0.9, 0.999)) #优化函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model,epoch,num_epochs,criterion,optimizer):\n",
    "    model.train()\n",
    "    print('-' * 100)\n",
    "    print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "    running_loss = 0.0\n",
    "    running_corrects = 0\n",
    "    for idx, (inputs, labels) in enumerate(dataloaders['train']):# 对dataloader进行遍历，dataloader时包含数据及标签的元组\n",
    "        inputs,labels=inputs.to(device),labels.to(device)\n",
    "        outputs = model(inputs) # output接受结果\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "        loss = criterion(outputs, labels)  # 默认平均，计算损失值\n",
    "\n",
    "        #反向传播及更新\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if idx % 100 == 99:\n",
    "            print('train iteration:{},loss:{},acc:{}%'.format( idx, loss.item(),torch.sum(preds == labels.data)/batch_size*100))\n",
    "        running_loss += loss.item() * inputs.size(0)\n",
    "        running_corrects += torch.sum(preds == labels.data)\n",
    "\n",
    "    epoch_loss = running_loss / dataset_sizes['train']\n",
    "    epoch_acc = running_corrects.double() / dataset_sizes['train']\n",
    "    print('train_total Loss: {:.4f} Acc: {:.4f}%'.format( epoch_loss, epoch_acc*100))\n",
    "    return epoch_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model,epoch,num_epochs,criterion,optimizer,best_acc):\n",
    "    model.eval()\n",
    "    running_loss = 0.0\n",
    "    running_corrects = 0\n",
    "    best_acc=best_acc\n",
    "    best_model_wts=copy.deepcopy(model.state_dict())\n",
    "    conf_matrix = torch.zeros(num_classes, num_classes) # 混淆矩阵初始化\n",
    "    with torch.no_grad():\n",
    "        for idx, (inputs, labels) in enumerate(dataloaders['val']):\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            outputs = model(inputs)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            loss = criterion(outputs, labels)\n",
    "            conf_matrix = confusion_matrix(outputs, labels, conf_matrix) # 生成混淆矩阵\n",
    "\n",
    "            running_loss += loss.item() * inputs.size(0)\n",
    "            running_corrects += torch.sum(preds == labels.data) #preds是tensor,labels.data是tensor\n",
    "\n",
    "        plot_confusion_matrix(conf_matrix, classes=class_names, normalize=False, title='confusion matrix epoch ' + str(epoch)) # 混淆矩阵的可视化\n",
    "\n",
    "    epoch_loss = running_loss / dataset_sizes['val'] \n",
    "    epoch_acc = running_corrects.double() / dataset_sizes['val'] #type为tensor\n",
    "    print('val_total Loss: {:.4f} Acc: {:.4f}%'.format( epoch_loss, epoch_acc*100))\n",
    "\n",
    "    all_prediction = calculate_all_prediction(conf_matrix) # 总精度=准确率\n",
    "    print('all_prediction:{}'.format(all_prediction))\n",
    "    label_prediction = [] # 存放每个类的精确率\n",
    "    label_recall = [] # 存放每个类的召回率\n",
    "    for i in range(num_classes):\n",
    "        label_prediction.append(calculate_label_prediction(conf_matrix,i))\n",
    "        label_recall.append(calculate_label_recall(conf_matrix,i))\n",
    "\n",
    "    keys=class_names\n",
    "    values=list(range(num_classes))\n",
    "    dictionary = dict(zip(keys, values))\n",
    "    for ei,i in enumerate(dictionary):\n",
    "        print(ei,'\\t',i,'\\t','prediction=',label_prediction[ei],'%,\\trecall=',label_recall[ei],'%,\\tf1=',calculate_f1(label_prediction[ei],label_recall[ei])) # 输出每个类的，精确率，召回率，F1\n",
    "    p = round(np.array(label_prediction).sum()/len(label_prediction),2) # 总精确率\n",
    "    r = round(np.array(label_recall).sum()/len(label_prediction),2) # 总召回率\n",
    "    print('MACRO-averaged:\\nprediction=',p,'%,recall=',r,'%,f1=',calculate_f1(p,r)) #输出总精确率和召回率\n",
    "\n",
    "#     print(epoch_acc.tpye)\n",
    "#     print(best_acc.type)\n",
    "    if epoch_acc > best_acc:# 获取最好的模型和准确率\n",
    "        best_acc=epoch_acc.item()\n",
    "        best_model_wts=copy.deepcopy(model.state_dict())\n",
    "#     model.load_state_dict(best_model_wts)\n",
    "\n",
    "    return best_model_wts,best_acc,epoch_acc.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if __name__ == '__main__':\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "    acc=[]\n",
    "    loss=[]\n",
    "    for epoch in range(num_epochs):\n",
    "        epoch_loss=train(model,epoch,num_epochs,criterion,optimizer)\n",
    "        best_model_wts,best_acc,epoch_acc=test(model,epoch,num_epochs,criterion,optimizer,best_acc)\n",
    "        acc.append(epoch_acc)\n",
    "        loss.append(epoch_loss)\n",
    "    print('*' * 100)\n",
    "    print('best_acc:{}'.format(best_acc))\n",
    "    print('*' * 100)\n",
    "    torch.save(best_model_wts, 'densenet201_3_model_best_acc.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=range(len(acc))\n",
    "y=acc\n",
    "plt.figure()\n",
    "plt.title('densenet201_3_acc')\n",
    "plt.plot(x,y)\n",
    "plt.savefig('mini64_lr0.0001_e20_densenet201_3_acc.jpg')\n",
    "plt.savefig('densenet201_3_acc.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=range(len(loss))\n",
    "y=loss\n",
    "plt.figure()\n",
    "plt.title('densenet201_3_loss')\n",
    "plt.plot(x,y)\n",
    "plt.savefig('mini64_lr0.0001_e20_densenet201_3_loss.jpg')\n",
    "plt.savefig('densenet201_3_loss.png')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}