win10图像迁移学习.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: cython in d:\\anaconda\\lib\\site-packages (0.29.21)\n",
      "Collecting git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI\n",
      "  Cloning https://github.com/philferriere/cocoapi.git to c:\\users\\12499\\appdata\\local\\temp\\pip-req-build-gztsnpjn\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: Command errored out with exit status 128: git clone -q https://github.com/philferriere/cocoapi.git 'C:\\Users\\12499\\AppData\\Local\\Temp\\pip-req-build-gztsnpjn' Check the logs for full command output.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/pytorch/vision.git\n",
      "  Cloning https://github.com/pytorch/vision.git to c:\\users\\12499\\appdata\\local\\temp\\pip-req-build-jav86j59\n",
      "Requirement already satisfied: numpy in d:\\anaconda\\lib\\site-packages (from torchvision==0.11.0a0+3b8d51e) (1.19.2)\n",
      "Requirement already satisfied: torch in d:\\anaconda\\lib\\site-packages (from torchvision==0.11.0a0+3b8d51e) (1.9.0)\n",
      "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in d:\\anaconda\\lib\\site-packages (from torchvision==0.11.0a0+3b8d51e) (8.0.1)\n",
      "Requirement already satisfied: typing_extensions in d:\\anaconda\\lib\\site-packages (from torch->torchvision==0.11.0a0+3b8d51e) (3.7.4.3)\n",
      "Building wheels for collected packages: torchvision\n",
      "  Building wheel for torchvision (setup.py): started\n",
      "  Building wheel for torchvision (setup.py): finished with status 'done'\n",
      "  Created wheel for torchvision: filename=torchvision-0.11.0a0+3b8d51e-cp38-cp38-win_amd64.whl size=680524 sha256=5d55920de9829f5e5d689736f0d9dc5d03a5effc3d7173b516d4972d6dbf1473\n",
      "  Stored in directory: C:\\Users\\12499\\AppData\\Local\\Temp\\pip-ephem-wheel-cache-rhohm6if\\wheels\\14\\ed\\60\\36a87c8a0b750f166c6113d23c64082c0c1ca83e1427ca4172\n",
      "Successfully built torchvision\n",
      "Installing collected packages: torchvision\n",
      "  Attempting uninstall: torchvision\n",
      "    Found existing installation: torchvision 0.11.0a0+16b9a40\n",
      "    Uninstalling torchvision-0.11.0a0+16b9a40:\n",
      "      Successfully uninstalled torchvision-0.11.0a0+16b9a40\n",
      "Successfully installed torchvision-0.11.0a0+3b8d51e\n"
     ]
    }
   ],
   "source": [
    "!pip install cython\n",
    "!pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI\n",
    "!pip install git+https://github.com/pytorch/vision.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "img=Image.open('C:/Users/12499/Documents/PennFudanPed/PNGImages/FudanPed00001.png')\n",
    "img.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "class PennFudanDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, root, transforms):\n",
    "        self.root = root\n",
    "        self.transforms = transforms\n",
    "        # load all image files, sorting them to\n",
    "        # ensure that they are aligned\n",
    "        self.imgs = list(sorted(os.listdir(os.path.join(root, \"PNGImages\"))))\n",
    "        self.masks = list(sorted(os.listdir(os.path.join(root, \"PedMasks\"))))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # load images and masks\n",
    "        img_path = os.path.join(self.root, \"PNGImages\", self.imgs[idx])\n",
    "        mask_path = os.path.join(self.root, \"PedMasks\", self.masks[idx])\n",
    "        img = Image.open(img_path).convert(\"RGB\")\n",
    "        # note that we haven't converted the mask to RGB,\n",
    "        # because each color corresponds to a different instance\n",
    "        # with 0 being background\n",
    "        mask = Image.open(mask_path)\n",
    "        # convert the PIL Image into a numpy array\n",
    "        mask = np.array(mask)\n",
    "        # instances are encoded as different colors\n",
    "        obj_ids = np.unique(mask)\n",
    "        # first id is the background, so remove it\n",
    "        obj_ids = obj_ids[1:]\n",
    "\n",
    "        # split the color-encoded mask into a set\n",
    "        # of binary masks\n",
    "        masks = mask == obj_ids[:, None, None]\n",
    "\n",
    "        # get bounding box coordinates for each mask\n",
    "        num_objs = len(obj_ids)\n",
    "        boxes = []\n",
    "        for i in range(num_objs):\n",
    "            pos = np.where(masks[i])\n",
    "            xmin = np.min(pos[1])\n",
    "            xmax = np.max(pos[1])\n",
    "            ymin = np.min(pos[0])\n",
    "            ymax = np.max(pos[0])\n",
    "            boxes.append([xmin, ymin, xmax, ymax])\n",
    "\n",
    "        # convert everything into a torch.Tensor\n",
    "        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n",
    "        # there is only one class\n",
    "        labels = torch.ones((num_objs,), dtype=torch.int64)\n",
    "        masks = torch.as_tensor(masks, dtype=torch.uint8)\n",
    "\n",
    "        image_id = torch.tensor([idx])\n",
    "        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n",
    "        # suppose all instances are not crowd\n",
    "        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)\n",
    "\n",
    "        target = {}\n",
    "        target[\"boxes\"] = boxes\n",
    "        target[\"labels\"] = labels\n",
    "        target[\"masks\"] = masks\n",
    "        target[\"image_id\"] = image_id\n",
    "        target[\"area\"] = area\n",
    "        target[\"iscrowd\"] = iscrowd\n",
    "\n",
    "        if self.transforms is not None:\n",
    "            img, target = self.transforms(img, target)\n",
    "\n",
    "        return img, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "\n",
    "# load a model pre-trained pre-trained on COCO\n",
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n",
    "\n",
    "# replace the classifier with a new one, that has\n",
    "# num_classes which is user-defined\n",
    "num_classes = 2  # 1 class (person) + background\n",
    "# get number of input features for the classifier\n",
    "in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "# replace the pre-trained head with a new one\n",
    "model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision.models.detection import FasterRCNN\n",
    "from torchvision.models.detection.rpn import AnchorGenerator\n",
    "\n",
    "# load a pre-trained model for classification and return\n",
    "# only the features\n",
    "backbone = torchvision.models.mobilenet_v2(pretrained=True).features\n",
    "# FasterRCNN needs to know the number of\n",
    "# output channels in a backbone. For mobilenet_v2, it's 1280\n",
    "# so we need to add it here\n",
    "backbone.out_channels = 1280\n",
    "\n",
    "# let's make the RPN generate 5 x 3 anchors per spatial\n",
    "# location, with 5 different sizes and 3 different aspect\n",
    "# ratios. We have a Tuple[Tuple[int]] because each feature\n",
    "# map could potentially have different sizes and\n",
    "# aspect ratios\n",
    "anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),\n",
    "                                   aspect_ratios=((0.5, 1.0, 2.0),))\n",
    "\n",
    "# let's define what are the feature maps that we will\n",
    "# use to perform the region of interest cropping, as well as\n",
    "# the size of the crop after rescaling.\n",
    "# if your backbone returns a Tensor, featmap_names is expected to\n",
    "# be [0]. More generally, the backbone should return an\n",
    "# OrderedDict[Tensor], and in featmap_names you can choose which\n",
    "# feature maps to use.\n",
    "roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],\n",
    "                                                output_size=7,\n",
    "                                                sampling_ratio=2)\n",
    "\n",
    "# put the pieces together inside a FasterRCNN model\n",
    "model = FasterRCNN(backbone,\n",
    "                   num_classes=2,\n",
    "                   rpn_anchor_generator=anchor_generator,\n",
    "                   box_roi_pool=roi_pooler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n",
    "\n",
    "\n",
    "def get_model_instance_segmentation(num_classes):\n",
    "    # load an instance segmentation model pre-trained pre-trained on COCO\n",
    "    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)\n",
    "\n",
    "    # get number of input features for the classifier\n",
    "    in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "    # replace the pre-trained head with a new one\n",
    "    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "\n",
    "    # now get the number of input features for the mask classifier\n",
    "    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n",
    "    hidden_layer = 256\n",
    "    # and replace the mask predictor with a new one\n",
    "    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,\n",
    "                                                       hidden_layer,\n",
    "                                                       num_classes)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'escape' from 'cgi' (D:\\anaconda\\lib\\cgi.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-2-8f09a7ff1c98>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mtransforms\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mT\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      2\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;32mdef\u001b[0m \u001b[0mget_transform\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m     \u001b[0mtransforms\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m     \u001b[0mtransforms\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mT\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mToTensor\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\anaconda\\lib\\site-packages\\transforms\\__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[1;33m.\u001b[0m\u001b[0msafe_html\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0msafe_html\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbodyfinder\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      2\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;31m# modules = [\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;31m#     'st',             # zopish\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;31m#     'rest',           # docutils\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\anaconda\\lib\\site-packages\\transforms\\safe_html.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mscrapy\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlinkextractors\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mLinkExtractor\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mre\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 3\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[0mcgi\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mescape\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      4\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[1;33m.\u001b[0m\u001b[0mutils\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0msafeToInt\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mImportError\u001b[0m: cannot import name 'escape' from 'cgi' (D:\\anaconda\\lib\\cgi.py)"
     ]
    }
   ],
   "source": [
    "import transforms as T\n",
    "\n",
    "def get_transform(train):\n",
    "    transforms = []\n",
    "    transforms.append(T.ToTensor())\n",
    "    if train:\n",
    "        transforms.append(T.RandomHorizontalFlip(0.5))\n",
    "    return T.Compose(transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'engine'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-11-7ad5399db9b4>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mpyttsx\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      2\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mengine\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtrain_one_epoch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mevaluate\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mutils\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mD:\\anaconda\\lib\\site-packages\\pyttsx\\__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     16\u001b[0m \u001b[0mOR\u001b[0m \u001b[0mIN\u001b[0m \u001b[0mCONNECTION\u001b[0m \u001b[0mWITH\u001b[0m \u001b[0mTHE\u001b[0m \u001b[0mUSE\u001b[0m \u001b[0mOR\u001b[0m \u001b[0mPERFORMANCE\u001b[0m \u001b[0mOF\u001b[0m \u001b[0mTHIS\u001b[0m \u001b[0mSOFTWARE\u001b[0m\u001b[1;33m.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     17\u001b[0m '''\n\u001b[1;32m---> 18\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[0mengine\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mEngine\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     19\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mweakref\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     20\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'engine'"
     ]
    }
   ],
   "source": [
    "import pyttsx\n",
    "from engine import train_one_epoch, evaluate\n",
    "import utils\n",
    "\n",
    "\n",
    "def main():\n",
    "    # train on the GPU or on the CPU, if a GPU is not available\n",
    "    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "    # our dataset has two classes only - background and person\n",
    "    num_classes = 2\n",
    "    # use our dataset and defined transformations\n",
    "    dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))\n",
    "    dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))\n",
    "\n",
    "    # split the dataset in train and test set\n",
    "    indices = torch.randperm(len(dataset)).tolist()\n",
    "    dataset = torch.utils.data.Subset(dataset, indices[:-50])\n",
    "    dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])\n",
    "\n",
    "    # define training and validation data loaders\n",
    "    data_loader = torch.utils.data.DataLoader(\n",
    "        dataset, batch_size=2, shuffle=True, num_workers=4,\n",
    "        collate_fn=utils.collate_fn)\n",
    "\n",
    "    data_loader_test = torch.utils.data.DataLoader(\n",
    "        dataset_test, batch_size=1, shuffle=False, num_workers=4,\n",
    "        collate_fn=utils.collate_fn)\n",
    "\n",
    "    # get the model using our helper function\n",
    "    model = get_model_instance_segmentation(num_classes)\n",
    "\n",
    "    # move model to the right device\n",
    "    model.to(device)\n",
    "\n",
    "    # construct an optimizer\n",
    "    params = [p for p in model.parameters() if p.requires_grad]\n",
    "    optimizer = torch.optim.SGD(params, lr=0.005,\n",
    "                                momentum=0.9, weight_decay=0.0005)\n",
    "    # and a learning rate scheduler\n",
    "    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,\n",
    "                                                   step_size=3,\n",
    "                                                   gamma=0.1)\n",
    "\n",
    "    # let's train it for 10 epochs\n",
    "    num_epochs = 10\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        # train for one epoch, printing every 10 iterations\n",
    "        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)\n",
    "        # update the learning rate\n",
    "        lr_scheduler.step()\n",
    "        # evaluate on the test dataset\n",
    "        evaluate(model, data_loader_test, device=device)\n",
    "\n",
    "    print(\"That's it!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n",
    "dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    " dataset, batch_size=2, shuffle=True, num_workers=4,\n",
    " collate_fn=utils.collate_fn)\n",
    "# For Training\n",
    "images,targets = next(iter(data_loader))\n",
    "images = list(image for image in images)\n",
    "targets = [{k: v for k, v in t.items()} for t in targets]\n",
    "output = model(images,targets)   # Returns losses and detections\n",
    "# For inference\n",
    "model.eval()\n",
    "x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]\n",
    "predictions = model(x)           # Returns predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from engine import train_one_epoch, evaluate\n",
    "import utils\n",
    "\n",
    "\n",
    "def main():\n",
    "    # train on the GPU or on the CPU, if a GPU is not available\n",
    "    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "    # our dataset has two classes only - background and person\n",
    "    num_classes = 2\n",
    "    # use our dataset and defined transformations\n",
    "    dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))\n",
    "    dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))\n",
    "\n",
    "    # split the dataset in train and test set\n",
    "    indices = torch.randperm(len(dataset)).tolist()\n",
    "    dataset = torch.utils.data.Subset(dataset, indices[:-50])\n",
    "    dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])\n",
    "\n",
    "    # define training and validation data loaders\n",
    "    data_loader = torch.utils.data.DataLoader(\n",
    "        dataset, batch_size=2, shuffle=True, num_workers=4,\n",
    "        collate_fn=utils.collate_fn)\n",
    "\n",
    "    data_loader_test = torch.utils.data.DataLoader(\n",
    "        dataset_test, batch_size=1, shuffle=False, num_workers=4,\n",
    "        collate_fn=utils.collate_fn)\n",
    "\n",
    "    # get the model using our helper function\n",
    "    model = get_model_instance_segmentation(num_classes)\n",
    "\n",
    "    # move model to the right device\n",
    "    model.to(device)\n",
    "\n",
    "    # construct an optimizer\n",
    "    params = [p for p in model.parameters() if p.requires_grad]\n",
    "    optimizer = torch.optim.SGD(params, lr=0.005,\n",
    "                                momentum=0.9, weight_decay=0.0005)\n",
    "    # and a learning rate scheduler\n",
    "    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,\n",
    "                                                   step_size=3,\n",
    "                                                   gamma=0.1)\n",
    "\n",
    "    # let's train it for 10 epochs\n",
    "    num_epochs = 10\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        # train for one epoch, printing every 10 iterations\n",
    "        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)\n",
    "        # update the learning rate\n",
    "        lr_scheduler.step()\n",
    "        # evaluate on the test dataset\n",
    "        evaluate(model, data_loader_test, device=device)\n",
    "\n",
    "    print(\"That's it!\")"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}