train_mapillary_AWS_run_2_LR_e-5.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/matplotlib/__init__.py:962: UserWarning: Duplicate key in file \"/home/ubuntu/.config/matplotlib/matplotlibrc\", line #2\n",
      "  (fname, cnt))\n",
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/matplotlib/__init__.py:962: UserWarning: Duplicate key in file \"/home/ubuntu/.config/matplotlib/matplotlibrc\", line #3\n",
      "  (fname, cnt))\n",
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:34: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import coco\n",
    "import sys\n",
    "import random\n",
    "import math\n",
    "import re\n",
    "import time\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import skimage.io\n",
    "import json\n",
    "from PIL import Image\n",
    "from config import Config\n",
    "import utils\n",
    "import model as modellib\n",
    "import visualize\n",
    "from model import log\n",
    "\n",
    "# Root directory of the project\n",
    "ROOT_DIR = os.getcwd()\n",
    "\n",
    "# Directory to save logs and trained model\n",
    "MODEL_DIR = os.path.join(ROOT_DIR, \"logs\")\n",
    "\n",
    "# Local path to trained weights file\n",
    "COCO_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_coco.h5\")\n",
    "\n",
    "# local path to mapillary dataset\n",
    "DATASET_DIR = os.path.join(ROOT_DIR, \"mapillary_dataset\")\n",
    "\n",
    "# Download COCO trained weights from Releases if needed\n",
    "if not os.path.exists(COCO_MODEL_PATH):\n",
    "    utils.download_trained_weights(COCO_MODEL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Configurations:\n",
      "BACKBONE_SHAPES                [[256 256]\n",
      " [128 128]\n",
      " [ 64  64]\n",
      " [ 32  32]\n",
      " [ 16  16]]\n",
      "BACKBONE_STRIDES               [4, 8, 16, 32, 64]\n",
      "BATCH_SIZE                     8\n",
      "BBOX_STD_DEV                   [0.1 0.1 0.2 0.2]\n",
      "DETECTION_MAX_INSTANCES        100\n",
      "DETECTION_MIN_CONFIDENCE       0.7\n",
      "DETECTION_NMS_THRESHOLD        0.3\n",
      "GPU_COUNT                      1\n",
      "IMAGES_PER_GPU                 8\n",
      "IMAGE_MAX_DIM                  1024\n",
      "IMAGE_MIN_DIM                  1024\n",
      "IMAGE_PADDING                  True\n",
      "IMAGE_SHAPE                    [1024 1024    3]\n",
      "LEARNING_MOMENTUM              0.9\n",
      "LEARNING_RATE                  1e-05\n",
      "MASK_POOL_SIZE                 14\n",
      "MASK_SHAPE                     [28, 28]\n",
      "MAX_GT_INSTANCES               100\n",
      "MEAN_PIXEL                     [123.7 116.8 103.9]\n",
      "MINI_MASK_SHAPE                (56, 56)\n",
      "NAME                           mapillary\n",
      "NUM_CLASSES                    12\n",
      "POOL_SIZE                      7\n",
      "POST_NMS_ROIS_INFERENCE        1000\n",
      "POST_NMS_ROIS_TRAINING         2000\n",
      "ROI_POSITIVE_RATIO             0.33\n",
      "RPN_ANCHOR_RATIOS              [0.5, 1, 2]\n",
      "RPN_ANCHOR_SCALES              (32, 64, 128, 256, 512)\n",
      "RPN_ANCHOR_STRIDE              1\n",
      "RPN_BBOX_STD_DEV               [0.1 0.1 0.2 0.2]\n",
      "RPN_NMS_THRESHOLD              0.7\n",
      "RPN_TRAIN_ANCHORS_PER_IMAGE    256\n",
      "STEPS_PER_EPOCH                128\n",
      "TRAIN_ROIS_PER_IMAGE           32\n",
      "USE_MINI_MASK                  True\n",
      "USE_RPN_ROIS                   True\n",
      "VALIDATION_STEPS               50\n",
      "WEIGHT_DECAY                   0.0005\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "class MapillaryConfig(coco.CocoConfig):\n",
    "    \"\"\"Configuration for training on the mapillary dataset.\n",
    "    Derives from the base Config class and overrides values specific\n",
    "    to the mapillary dataset.\n",
    "    \"\"\"\n",
    "    # Give the configuration a recognizable name\n",
    "    NAME = \"mapillary\"\n",
    "\n",
    "    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each\n",
    "    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).\n",
    "    GPU_COUNT = 1\n",
    "    IMAGES_PER_GPU = 8\n",
    "\n",
    "    # Number of classes (including background)\n",
    "    # this MUST be explicitly defined, or will run into index out of bound error\n",
    "    NUM_CLASSES = 1 + 11  # background + 11 objects\n",
    "\n",
    "    # Use small images for faster training. Set the limits of the small side\n",
    "    # the large side, and that determines the image shape.\n",
    "    IMAGE_MIN_DIM = 1024\n",
    "    IMAGE_MAX_DIM = 1024\n",
    "\n",
    "    # Use smaller anchors because our image and objects are small\n",
    "    # RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # anchor side in pixels\n",
    "\n",
    "    # Reduce training ROIs per image because the images are small and have\n",
    "    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.\n",
    "    TRAIN_ROIS_PER_IMAGE = 32\n",
    "\n",
    "    # Use a small epoch since the data is simple\n",
    "    STEPS_PER_EPOCH = 128\n",
    "\n",
    "    # use small validation steps since the epoch is small\n",
    "    VALIDATION_STEPS = 50\n",
    "    \n",
    "    LEARNING_RATE = 0.00001\n",
    "    WEIGHT_DECAY = 0.0005\n",
    "    \n",
    "config = MapillaryConfig()\n",
    "config.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MapillaryDataset(utils.Dataset):\n",
    "    \"\"\"Generates the shapes synthetic dataset. The dataset consists of simple\n",
    "    shapes (triangles, squares, circles) placed randomly on a blank surface.\n",
    "    The images are generated on the fly. No file access required.\n",
    "    \"\"\"\n",
    "    \n",
    "    DEBUG = False\n",
    "    CLASS_MAP = {}\n",
    "    \n",
    "    CLASSES = [\"Bird\", \"Person\", \"Bicyclist\", \"Motorcyclist\", \"Bench\", \\\n",
    "         \"Fire Hydrant\", \"Traffic Light\", \"Bus\", \"Car\", \"Motorcycle\", \"Truck\"]\n",
    "    \n",
    "    # local path to image folder, choose 'dev', 'training', or 'testing'\n",
    "    SUBSET_DIR = \"\"\n",
    "\n",
    "    # local path to images inside development folder\n",
    "    IMG_DIR = \"\"\n",
    "\n",
    "    # local path to instance annotations inside development folder\n",
    "    INS_DIR = \"\"\n",
    "\n",
    "\n",
    "    def load_mapillary(self, dataset_dir, subset, class_ids=None,\n",
    "                  class_map=None):\n",
    "        self.SUBSET_DIR = os.path.join(dataset_dir, subset)\n",
    "        self.IMG_DIR = os.path.join(self.SUBSET_DIR, 'images')\n",
    "        self.INS_DIR = os.path.join(self.SUBSET_DIR, 'instances')\n",
    "        # load classes, start with id = 1 to account for background \"BG\"\n",
    "        class_id = 1\n",
    "        for label_id, label in enumerate(class_ids):\n",
    "            if label[\"instances\"] == True and label[\"readable\"] in self.CLASSES:\n",
    "                self.CLASS_MAP[label_id] = class_id\n",
    "                \n",
    "                if (self.DEBUG):\n",
    "                    print(\"{}: Class {} {} added\".format(label_id, class_id, label[\"readable\"]))\n",
    "                    \n",
    "                self.add_class(\"mapillary\", class_id, label[\"readable\"])\n",
    "                class_id = class_id + 1\n",
    "                \n",
    "        # add images \n",
    "        file_names = next(os.walk(self.IMG_DIR))[2]\n",
    "        if (self.DEBUG):\n",
    "            print(\"{} images in directory {}\".format(len(file_names), self.IMG_DIR))\n",
    "        for i in range(len(file_names)):\n",
    "            if file_names[i] != '.DS_Store':\n",
    "                \n",
    "                image_path = os.path.join(self.IMG_DIR, file_names[i])\n",
    "                base_image = Image.open(image_path)\n",
    "                w, h = base_image.size\n",
    "\n",
    "                if (self.DEBUG):\n",
    "                    print(\"Image {} {} x {} added\".format(file_names[i], w, h))\n",
    "\n",
    "                self.add_image(\"mapillary\", image_id = i,\n",
    "                              path = file_names[i],\n",
    "                               width = w,\n",
    "                               height = h\n",
    "                              )\n",
    "                \n",
    "    def load_image(self, image_id):\n",
    "        \"\"\"Generate an image from the specs of the given image ID.\n",
    "        This function loads the image from a file.\n",
    "        \"\"\"\n",
    "        info = self.image_info[image_id]\n",
    "        img_path = os.path.join(self.IMG_DIR, info[\"path\"])\n",
    "        image = Image.open(img_path)\n",
    "        image_array = np.array(image)\n",
    "        return image_array\n",
    "\n",
    "    def image_reference(self, image_id):\n",
    "        \"\"\"Return the local directory path of the image.\"\"\"\n",
    "        info = self.image_info[image_id]\n",
    "        img_path = os.path.join(self.IMG_DIR, info[\"path\"])\n",
    "        return img_path\n",
    "\n",
    "    def load_mask(self, image_id):\n",
    "        \"\"\"Generate instance masks for shapes of the given image ID.\n",
    "        \"\"\"\n",
    "        info = self.image_info[image_id]\n",
    "        instance_path = os.path.join(self.INS_DIR, info[\"path\"])\n",
    "        instance_image = Image.open(instance_path.rsplit(\".\", 1)[0] + \".png\")\n",
    "        \n",
    "        # convert labeled data to numpy arrays for better handling\n",
    "        instance_array = np.array(instance_image, dtype=np.uint16)\n",
    "\n",
    "        instances = np.unique(instance_array)\n",
    "        instaces_count = instances.shape[0]\n",
    "        \n",
    "        label_ids = instances // 256\n",
    "        label_id_count = np.unique(label_ids).shape[0]\n",
    "        \n",
    "        if (self.DEBUG):\n",
    "            print(\"There are {} instances, {} classes labelled instances in the image {}.\"\\\n",
    "                  .format(instaces_count, label_id_count, info[\"path\"]))\n",
    "            \n",
    "        mask = np.zeros([instance_array.shape[0], instance_array.shape[1], instaces_count], dtype=np.uint8)\n",
    "        mask_count = 0\n",
    "        loaded_class_ids = []\n",
    "        for instance in instances:\n",
    "            label_id = instance // 256\n",
    "            if (label_id in self.CLASS_MAP):\n",
    "                m = np.zeros((instance_array.shape[0], instance_array.shape[1]), dtype=np.uint8)\n",
    "                m[instance_array == instance] = 1\n",
    "                m_size = np.count_nonzero(m == 1)\n",
    "                \n",
    "                # only load mask greater than threshold size, \n",
    "                # otherwise bounding box with area zero causes program to crash\n",
    "                if m_size > 32 ** 2:\n",
    "                    mask[:, :, mask_count] = m\n",
    "                    loaded_class_ids.append(self.CLASS_MAP[label_id])\n",
    "                    mask_count = mask_count + 1\n",
    "                    if (self.DEBUG):\n",
    "                        print(\"Non-zero: {}\".format(m_size))\n",
    "                        print(\"Mask {} created for instance {} of class {} {}\"\\\n",
    "                              .format(mask_count, instance, self.CLASS_MAP[label_id], \\\n",
    "                                      self.class_names[self.CLASS_MAP[label_id]]))\n",
    "        mask = mask[:, :, 0:mask_count]\n",
    "        return mask, np.array(loaded_class_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mapping:  public\n",
      "version:  1.0\n",
      "folder_structure: {split}/{content}/{key:.{22}}.{ext}\n",
      "There are 66 classes in the config file\n",
      "There are 12 classes in the model\n",
      "    Class 0: BG\n",
      "    Class 1: Bird\n",
      "    Class 2: Person\n",
      "    Class 3: Bicyclist\n",
      "    Class 4: Motorcyclist\n",
      "    Class 5: Bench\n",
      "    Class 6: Fire Hydrant\n",
      "    Class 7: Traffic Light\n",
      "    Class 8: Bus\n",
      "    Class 9: Car\n",
      "    Class 10: Motorcycle\n",
      "    Class 11: Truck\n"
     ]
    }
   ],
   "source": [
    "# read in config file\n",
    "with open(os.path.join(DATASET_DIR, 'config.json')) as config_file:\n",
    "    class_config = json.load(config_file)\n",
    "# in this example we are only interested in the labels\n",
    "labels = class_config['labels']\n",
    "\n",
    "AWS_run_2_dir = os.path.join(DATASET_DIR, \"AWS_run_2\")\n",
    "\n",
    "# Training dataset\n",
    "dataset_train = MapillaryDataset()\n",
    "dataset_train.load_mapillary(AWS_run_2_dir, \"train_4096/accepted\", class_ids = labels)\n",
    "dataset_train.prepare()\n",
    "\n",
    "# Validation dataset\n",
    "dataset_val = MapillaryDataset()\n",
    "dataset_val.load_mapillary(AWS_run_2_dir, \"dev_512/accepted\", class_ids = labels)\n",
    "dataset_val.prepare()\n",
    "\n",
    "print(\"mapping: \", class_config[\"mapping\"])\n",
    "print(\"version: \", class_config[\"version\"])\n",
    "print(\"folder_structure:\", class_config[\"folder_structure\"])\n",
    "print(\"There are {} classes in the config file\".format(len(labels)))\n",
    "print(\"There are {} classes in the model\".format(len(dataset_train.class_names)))\n",
    "for i in range(len(dataset_train.class_names)):\n",
    "    print(\"    Class {}: {}\".format(i, dataset_train.class_names[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Create model in training mode\n",
    "model = modellib.MaskRCNN(mode=\"training\", config=config,\n",
    "                          model_dir=MODEL_DIR)\n",
    "\n",
    "# Which weights to start with?\n",
    "init_with = \"coco\"  # imagenet, coco, or last\n",
    "\n",
    "if init_with == \"imagenet\":\n",
    "    model.load_weights(model.get_imagenet_weights(), by_name=True)\n",
    "     \n",
    "elif init_with == \"coco\":\n",
    "    # Load weights trained on MS COCO, but skip layers that\n",
    "    # are different due to the different number of classes\n",
    "    # See README for instructions to download the COCO weights\n",
    "    model.load_weights(COCO_MODEL_PATH, by_name=True,\n",
    "                       exclude=[\"mrcnn_class_logits\", \"mrcnn_bbox_fc\", \n",
    "                                \"mrcnn_bbox\", \"mrcnn_mask\"])\n",
    "elif init_with == \"last\":\n",
    "    # Load the last model you trained and continue training\n",
    "    model.load_weights(model.find_last()[1], by_name=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Starting at epoch 0. LR=1e-05\n",
      "\n",
      "Checkpoint Path: /home/ubuntu/DreamNet/logs/mapillary20180313T2209/mask_rcnn_mapillary_{epoch:04d}.h5\n",
      "Selecting layers to train\n",
      "fpn_c5p5               (Conv2D)\n",
      "fpn_c4p4               (Conv2D)\n",
      "fpn_c3p3               (Conv2D)\n",
      "fpn_c2p2               (Conv2D)\n",
      "fpn_p5                 (Conv2D)\n",
      "fpn_p2                 (Conv2D)\n",
      "fpn_p3                 (Conv2D)\n",
      "fpn_p4                 (Conv2D)\n",
      "In model:  rpn_model\n",
      "    rpn_conv_shared        (Conv2D)\n",
      "    rpn_class_raw          (Conv2D)\n",
      "    rpn_bbox_pred          (Conv2D)\n",
      "mrcnn_mask_conv1       (TimeDistributed)\n",
      "mrcnn_mask_bn1         (TimeDistributed)\n",
      "mrcnn_mask_conv2       (TimeDistributed)\n",
      "mrcnn_mask_bn2         (TimeDistributed)\n",
      "mrcnn_class_conv1      (TimeDistributed)\n",
      "mrcnn_class_bn1        (TimeDistributed)\n",
      "mrcnn_mask_conv3       (TimeDistributed)\n",
      "mrcnn_mask_bn3         (TimeDistributed)\n",
      "mrcnn_class_conv2      (TimeDistributed)\n",
      "mrcnn_class_bn2        (TimeDistributed)\n",
      "mrcnn_mask_conv4       (TimeDistributed)\n",
      "mrcnn_mask_bn4         (TimeDistributed)\n",
      "mrcnn_bbox_fc          (TimeDistributed)\n",
      "mrcnn_mask_deconv      (TimeDistributed)\n",
      "mrcnn_class_logits     (TimeDistributed)\n",
      "mrcnn_mask             (TimeDistributed)\n",
      "WARNING:tensorflow:From /home/ubuntu/DreamNet/model.py:2160: calling reduce_mean (from tensorflow.python.ops.math_ops) with keep_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "keep_dims is deprecated, use keepdims instead\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/scipy/misc/pilutil.py:482: FutureWarning: Conversion of the second argument of issubdtype from `int` to `np.signedinteger` is deprecated. In future, it will be treated as `np.int64 == np.dtype(int).type`.\n",
      "  if issubdtype(ts, int):\n",
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/scipy/misc/pilutil.py:485: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  elif issubdtype(type(size), float):\n",
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/scipy/ndimage/interpolation.py:600: UserWarning: From scipy 0.13.0, the output shape of zoom() is calculated with round() instead of int() - for these inputs the size of the returned array has changed.\n",
      "  \"the returned array has changed.\", UserWarning)\n",
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py:98: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  \"Converting sparse IndexedSlices to a dense Tensor of unknown shape. \"\n",
      "/home/ubuntu/anaconda3/lib/python3.6/site-packages/keras/engine/training.py:2087: UserWarning: Using a generator with `use_multiprocessing=True` and multiple workers may duplicate your data. Please consider using the`keras.utils.Sequence class.\n",
      "  UserWarning('Using a generator with `use_multiprocessing=True`'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/8\n",
      "128/128 [==============================] - 1000s 8s/step - loss: 3.5280 - rpn_class_loss: 0.0161 - rpn_bbox_loss: 0.1529 - mrcnn_class_loss: 1.6410 - mrcnn_bbox_loss: 0.9702 - mrcnn_mask_loss: 0.7478 - val_loss: 2.6778 - val_rpn_class_loss: 0.0099 - val_rpn_bbox_loss: 0.0619 - val_mrcnn_class_loss: 1.0494 - val_mrcnn_bbox_loss: 0.8295 - val_mrcnn_mask_loss: 0.7271\n",
      "\n",
      "Epoch 00001: saving model to /home/ubuntu/DreamNet/logs/mapillary20180313T2209/mask_rcnn_mapillary_0001.h5\n",
      "Epoch 2/8\n",
      "128/128 [==============================] - 937s 7s/step - loss: 2.2791 - rpn_class_loss: 0.0156 - rpn_bbox_loss: 0.1546 - mrcnn_class_loss: 0.7718 - mrcnn_bbox_loss: 0.6791 - mrcnn_mask_loss: 0.6580 - val_loss: 1.9209 - val_rpn_class_loss: 0.0094 - val_rpn_bbox_loss: 0.0641 - val_mrcnn_class_loss: 0.6153 - val_mrcnn_bbox_loss: 0.6027 - val_mrcnn_mask_loss: 0.6294\n",
      "\n",
      "Epoch 00002: saving model to /home/ubuntu/DreamNet/logs/mapillary20180313T2209/mask_rcnn_mapillary_0002.h5\n",
      "Epoch 3/8\n",
      "128/128 [==============================] - 941s 7s/step - loss: 1.9389 - rpn_class_loss: 0.0171 - rpn_bbox_loss: 0.1659 - mrcnn_class_loss: 0.5819 - mrcnn_bbox_loss: 0.5709 - mrcnn_mask_loss: 0.6031 - val_loss: 1.4311 - val_rpn_class_loss: 0.0089 - val_rpn_bbox_loss: 0.0654 - val_mrcnn_class_loss: 0.3891 - val_mrcnn_bbox_loss: 0.4242 - val_mrcnn_mask_loss: 0.5436\n",
      "\n",
      "Epoch 00003: saving model to /home/ubuntu/DreamNet/logs/mapillary20180313T2209/mask_rcnn_mapillary_0003.h5\n",
      "Epoch 4/8\n",
      "128/128 [==============================] - 943s 7s/step - loss: 1.7374 - rpn_class_loss: 0.0164 - rpn_bbox_loss: 0.1559 - mrcnn_class_loss: 0.5022 - mrcnn_bbox_loss: 0.5026 - mrcnn_mask_loss: 0.5603 - val_loss: 1.1799 - val_rpn_class_loss: 0.0085 - val_rpn_bbox_loss: 0.0660 - val_mrcnn_class_loss: 0.2936 - val_mrcnn_bbox_loss: 0.3481 - val_mrcnn_mask_loss: 0.4637\n",
      "\n",
      "Epoch 00004: saving model to /home/ubuntu/DreamNet/logs/mapillary20180313T2209/mask_rcnn_mapillary_0004.h5\n",
      "Epoch 5/8\n",
      "122/128 [===========================>..] - ETA: 45s - loss: 1.6080 - rpn_class_loss: 0.0159 - rpn_bbox_loss: 0.1618 - mrcnn_class_loss: 0.4483 - mrcnn_bbox_loss: 0.4574 - mrcnn_mask_loss: 0.5245"
     ]
    }
   ],
   "source": [
    "# Train the head branches\n",
    "# Passing layers=\"heads\" freezes all layers except the head\n",
    "# layers. You can also pass a regular expression to select\n",
    "# which layers to train by name pattern.\n",
    "model.train(dataset_train, dataset_val, \n",
    "            learning_rate=config.LEARNING_RATE, \n",
    "            epochs=8, \n",
    "            layers='heads')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done\n"
     ]
    }
   ],
   "source": [
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}