train_detection.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision.models.alexnet import AlexNet\n",
    "from torchvision import transforms\n",
    "\n",
    "from dataset_models import DetectionDataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 判断可用设备\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"using {device} device.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 指定数据集目录\n",
    "image_path = os.path.abspath('datasets/JPEGImages/')\n",
    "if not os.path.exists(image_path):\n",
    "    raise Exception(f\"{image_path} path does not exist.\")\n",
    "\n",
    "prop_path = os.path.abspath('datasets/Proposals/')\n",
    "if not os.path.exists(prop_path):\n",
    "    raise Exception(f\"{prop_path} path does not exist.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据预处理与增强\n",
    "train_transform = transforms.Compose([transforms.Resize((224, 224)),\n",
    "                                     transforms.RandomHorizontalFlip(p=0.5),\n",
    "                                     transforms.RandomVerticalFlip(p=0.5),\n",
    "                                     transforms.ColorJitter(brightness=0.2, contrast=0.2, hue=0.2),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),])\n",
    "val_transform = transforms.Compose([transforms.Resize((224, 224)),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = DetectionDataset(image_path, prop_path, train_transform)\n",
    "validate_dataset = DetectionDataset(image_path, prop_path, val_transform)\n",
    "\n",
    "item_list = train_dataset.item_list\n",
    "train_num = int(0.9 * len(item_list))\n",
    "val_num = len(item_list) - train_num\n",
    "random.shuffle(item_list)\n",
    "train_dataset.item_list = item_list[:train_num]\n",
    "validate_dataset.item_list = item_list[train_num:]\n",
    "\n",
    "train_num, val_num\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "cpu_count = os.cpu_count()\n",
    "num_workers = cpu_count - 1 if cpu_count and cpu_count > 1 else 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)\n",
    "validate_loader = DataLoader(validate_dataset, batch_size=4, shuffle=False,  num_workers=num_workers)\n",
    "print(f\"using {train_num} images for training, {val_num} images for validation.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 实例化模型，并送进设备\n",
    "net = AlexNet(num_classes=2)\n",
    "net.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 指定损失函数用于计算损失；指定优化器用于更新模型参数；指定训练迭代的轮数，训练权重的存储地址\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "optimizer = Adam(net.parameters(), lr=0.0002)\n",
    "epochs = 10\n",
    "save_path = os.path.abspath('./weights')\n",
    "if not os.path.exists(save_path):\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "\n",
    "best_acc = 0.0  # 初始化验证集上最好的准确率，以便后面用该指标筛选模型最优参数。\n",
    "train_steps = len(train_loader)  # rain_steps = len(dataset) / batch_size\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    net.train()\n",
    "    train_loss = torch.zeros(1).to(device) # 初始化，用于计算训练损失           torch.zeros(2, 3)————这就是torch.zeros的用法，括号内是size\n",
    "                                           # tensor([[ 0.,  0.,  0.],\n",
    "                                           #         [ 0.,  0.,  0.]])\n",
    "    acc_num = torch.zeros(1).to(device)    # 初始化，用于计算训练过程中预测正确的数量\n",
    "    sample_num = 0                         # 初始化，用于记录当前迭代中，已经计算了多少个样本\n",
    "    # tqdm是一个进度条显示器，可以在终端打印出现在的训练进度\n",
    "    train_bar = tqdm(train_loader, file=sys.stdout, ncols=100)\n",
    "    for step, data in enumerate(train_bar):\n",
    "        images, labels = data\n",
    "        sample_num += images.shape[0]\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(images.to(device))          # output_shape: [batch_size, num_classes]   这里的images应该是前向传播forward中的x\n",
    "        pred_class = torch.max(outputs, dim=1)[1] # torch.max 返回值是一个tuple，第一个元素是max值，第二个元素是max值的索引。  \n",
    "        #这里dim表示要降维的维度，pred_class范围的是分类号\n",
    "        acc_num += torch.eq(pred_class, labels.to(device)).sum() #torch.eq()判断后面两个数组对应元素是否相等，相等为true，不等为flase，这里对bool数组求和啥意思呢\n",
    "\n",
    "        loss = loss_function(outputs, labels.to(device)) # 求损失，  ？？？为什么上下同时对labels做计算，但是计算的维度信息并不匹配啊？？？\n",
    "        loss.backward() # 自动求导\n",
    "        optimizer.step() # 梯度下降\n",
    "\n",
    "        # print statistics\n",
    "        train_loss += loss.detach()  / (step + 1)\n",
    "        train_acc = acc_num.item() / sample_num \n",
    "        # .desc是进度条tqdm中的成员变量，作用是描述信息\n",
    "        train_bar.desc = f\"Epoch {epoch + 1}/{epochs}\"\n",
    "\n",
    "    # validate\n",
    "    net.eval()\n",
    "    acc_num = 0.0  # accumulate accurate number per epoch\n",
    "    with torch.no_grad(): \n",
    "        for val_data in validate_loader:\n",
    "            val_images, val_labels = val_data\n",
    "            outputs = net(val_images.to(device))\n",
    "            predict_y = torch.max(outputs, dim=1)[1]\n",
    "            acc_num += torch.eq(predict_y, val_labels.to(device)).sum().item() \n",
    "\n",
    "    val_acc = acc_num / val_num\n",
    "    print(f'Epoch {epoch + 1}/{epochs}: train_loss={float(train_loss / train_steps):.3f} train_acc={float(train_acc):.3f} val_accuracy={float(val_acc):.3f}')\n",
    "    # 判断当前验证集的准确率是否是最大的，如果是，则更新之前保存的权重\n",
    "    if val_acc > best_acc:\n",
    "        best_acc = val_acc\n",
    "        torch.save(net.state_dict(), os.path.join(save_path, \"Detection.pth\"))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "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.11.3"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "26de051ba29f2982a8de78e945f0abaf191376122a1563185a90213a26c5da77"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}