notebook-xview3-blog-walkthrough.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# xView3 SAR - Object detection on giga-pixel satellite imagery\n",
    "\n",
    "## Background\n",
    "\n",
    "In this notebook, we demonstrate how to train, evaluate, and deploy a custom object detection model for xView3 SAR."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [Optional] Configure docker \n",
    "Below we configure Docker in our Amazon SageMaker environment to increase the shared memory size and specify a root directory located on the instance EBS volume."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile /home/ec2-user/SageMaker/daemon.json\n",
    "{\n",
    "    \"runtimes\": {\n",
    "        \"nvidia\": {\n",
    "            \"path\": \"nvidia-container-runtime\",\n",
    "            \"runtimeArgs\": []\n",
    "        }\n",
    "    },\n",
    "    \"default-shm-size\": \"200G\",\n",
    "    \"data-root\": \"/home/ec2-user/SageMaker/docker\"\n",
    "} "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the shell script below to make changes to Docker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "sudo service docker stop\n",
    "mkdir -p /home/ec2-user/SageMaker/docker\n",
    "sudo rsync -aqxP /var/lib/docker/ /home/ec2-user/SageMaker/docker\n",
    "sudo mv /var/lib/docker /var/lib/docker.old\n",
    "sudo mv /home/ec2-user/SageMaker/daemon.json /etc/docker/\n",
    "sudo service docker start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "\n",
    "import boto3\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.estimator import Estimator\n",
    "from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor, ScriptProcessor\n",
    "\n",
    "sys.path.append('tools/')\n",
    "from docker_utils import build_and_push_docker_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define execution role, S3 bucket in account, and SageMaker session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "role = get_execution_role()\n",
    "region = boto3.Session().region_name\n",
    "bucket = 'xview3-blog-sagemaker'\n",
    "sagemaker_session = sagemaker.Session(default_bucket=bucket)\n",
    "account = sagemaker_session.account_id()\n",
    "tags =[{'Key': 'project', 'Value': 'xview3-blog'}]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the boolean flag `USE_TINY` to run the notebook using a tiny subset of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_TINY = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Creation with SageMaker Processing\n",
    "In this section we create the following from the xView3 challenge dataset:\n",
    "1. a new `train` and `valid`, after merging the train and validation set provided by the challenge. \n",
    "2. Detectron2 compatible dataset dicts to be used for training. \n",
    "\n",
    "\n",
    "#### Merge & split data labels. \n",
    "The xView3 challenge provided detection labels for each scene in `train.csv`, `validation.csv`, and `public.csv`. \n",
    "We will merge the `train.csv` and `validation.csv` and create a new `train` and `validation` set for training. The `public` leaderboard set will remain fixed.\n",
    "\n",
    "#### Create Detectron2 Datasets\n",
    "Here we create the Detectron2-compatible [dataset dicts](https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html) used for training models in Detectron2. The format of the dataset is a list of dictionaries with each dict containing information for one image with at least the following fields:\n",
    "- `filename`:str\n",
    "- `height`:int\n",
    "- `width`: int\n",
    "- `image_id`:str or int\n",
    "- `annotations`: list[dict]\n",
    "\n",
    "For more information on how to generate the dataset dict, see [Detectron2 docs] (https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html#standard-dataset-dicts).\n",
    "\n",
    "Our dataset dict is generated from the information provided in the label `csv` files used in the previous section. Depending on whether we train our models with inputs originating from image chips (tiles) or from the full scene we will use one of two functions in the `xview3_d2` pacakage: `create_xview3_chipped_scene_annotations` or `create_xview3_full_scene_annotations`, respectively. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build base image for SM Processing tasks.\n",
    "For convenience, we build a base processing container which handles package installations. We can build a subsequent image from this base container to include the code we want to run. Here is what the base processing container looks like.\n",
    "\n",
    "For building and pushing the containers, we use helper function `build_and_push_docker_image` in `tools/docker_utils.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize -l docker docker/processing/base.Dockerfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's build and push the base processing container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "processing_base_name = 'xview3-processing:base'\n",
    "base_image = build_and_push_docker_image(processing_base_name, \n",
    "                                         dockerfile='docker/processing/base.Dockerfile')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the main processing container, which copies the `.py` scripts in `tools/`. In each processing job, to follow, we can specify which the entrypoint `.py` script to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize -l docker docker/processing/main.Dockerfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build and push main processing container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processing_base_image = f'{account}.dkr.ecr.{region}.amazonaws.com/xview3-processing:base'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "processing_main_name = 'xview3-processing:main'\n",
    "processing_main_image = build_and_push_docker_image(processing_main_name, \n",
    "                                                    dockerfile='docker/processing/main.Dockerfile', \n",
    "                                                    base_image=base_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processing_main_image = f'{account}.dkr.ecr.{region}.amazonaws.com/xview3-processing:main'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Launch SageMaker Processing job for dataset preparation. \n",
    "The SageMaker Processing task will run `tools/create_xview3_dataset_dict.py`. This script creates a detectron2-compatible dataset dict for full scene imagery or chipped scenes. Optionally, this script will merge train and validation csvs and create a new split. \n",
    "\n",
    "Let's see the arguments required this script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pygmentize -l python tools/create_xview3_dataset_dict.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize SM Processing job. \n",
    "We only need 1 instance for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_type = 'ml.t3.xlarge'\n",
    "volume_size_in_gb = 30 \n",
    "instance_count = 1\n",
    "base_job_name = 'xview3-dataset-prep'\n",
    "                      \n",
    "dataset_processor = Processor(image_uri=processing_image_name,\n",
    "                              role=role,\n",
    "                              instance_count=instance_count,\n",
    "                              base_job_name=base_job_name,\n",
    "                              instance_type=instance_type, \n",
    "                              volume_size_in_gb=volume_size_in_gb, \n",
    "                              entrypoint=['python3', 'create_xview3_dataset_dict.py'],\n",
    "                              sagemaker_session=sagemaker_session, \n",
    "                              tags=tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Specify inputs and run processing job. \n",
    "\n",
    "`tools/create_xview3_dataset_dict.py` has several defaults, which can be overridden by providing the relevant argument in the processor `arugments`.  The cell below will launch a processor job that creates a new data split and creates a dataset dict for full scenes. To create a dataset dict for chipped scenes, change `dataset-type` to `chipped` and provide additional inputs and/or arguments such as `--shoreline_dir`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "override = False\n",
    "current_timestamp = '202207250702'\n",
    "SEED = 46998886\n",
    "\n",
    "if override:\n",
    "    current_timestamp = datetime.now().strftime(\"%Y%m%d%M%S\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_labels = ProcessingInput(source='data/labels/', \n",
    "                               destination='/opt/ml/processing/input/labels',\n",
    "                              input_name='labels')\n",
    "input_stats = ProcessingInput(source='data/scene-stats.csv', \n",
    "                              destination='/opt/ml/processing/input/scene-stats',\n",
    "                             input_name='stats')\n",
    "\n",
    "job_output = ProcessingOutput(source='/opt/ml/processing/output/prepared/',  \n",
    "                              destination=f's3://xview3-blog/data/processing/{current_timestamp}',\n",
    "                              output_name='prepared-dataset')\n",
    "\n",
    "dataset_processor.run(inputs=[input_labels, input_stats], \n",
    "              outputs=[job_output],\n",
    "              arguments=[\"--dataset-type\", \"full\", \n",
    "                         \"--train-labels-csv\", f\"{input_labels.destination}/train.csv\",\n",
    "                         \"--valid-labels-csv\", f\"{input_labels.destination}/validation.csv\",\n",
    "                         \"--tiny-labels-csv\", f\"{input_labels.destination}/tiny.csv\",\n",
    "                         \"--scene-stats-csv\", f\"{input_stats.destination}/scene-stats.csv\",\n",
    "                         \"--seed\", str(SEED), \n",
    "                         \"--output-dir\", job_output.source,\n",
    "              ],\n",
    "              wait=True,\n",
    "              logs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [Optional] Run processing job to created dataset dict for chipped scenes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processing_image_name = f'{account}.dkr.ecr.{region}.amazonaws.com/xview3-processing:main'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_type = 'ml.t3.xlarge'\n",
    "volume_size_in_gb = 30 \n",
    "instance_count = 1\n",
    "base_job_name = 'xview3-dataset-prep'\n",
    "                      \n",
    "dataset_processor = Processor(image_uri=processing_image_name,\n",
    "                              role=role,\n",
    "                              instance_count=instance_count,\n",
    "                              base_job_name=base_job_name,\n",
    "                              instance_type=instance_type, \n",
    "                              volume_size_in_gb=volume_size_in_gb, \n",
    "                              entrypoint=['python3', 'create_xview3_dataset_dict.py'],\n",
    "                              sagemaker_session=sagemaker_session, \n",
    "                              tags=tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_destination_uri = f's3://xview3-blog/data/processing/{current_timestamp}'\n",
    "\n",
    "input_stats = ProcessingInput(source='data/scene-stats.csv', \n",
    "                              destination='/opt/ml/processing/input/scene-stats',\n",
    "                              input_name='stats')\n",
    "input_label_trn = ProcessingInput(source=f'{s3_destination_uri}/labels/train.csv',\n",
    "                                  destination='/opt/ml/processing/input/labels/train',\n",
    "                                  input_name='trn-labels')\n",
    "input_labels_tiny = ProcessingInput(source=f'{s3_destination_uri}/labels/tiny-train.csv',\n",
    "                                    destination='/opt/ml/processing/input/labels/tiny',\n",
    "                                    input_name='tiny-labels')\n",
    "inputs_shoreline = ProcessingInput(source='s3://xview3-blog/data/shoreline/trainval/', \n",
    "                                  destination='/opt/ml/processing/input/shoreline/')\n",
    "\n",
    "job_output = ProcessingOutput(source='/opt/ml/processing/output/prepared/',  \n",
    "                              destination=s3_destination_uri,\n",
    "                              output_name='prepared-dataset')\n",
    "\n",
    "dataset_processor.run(inputs=[input_label_trn, input_labels_tiny, input_stats, inputs_shoreline], \n",
    "                      outputs=[job_output],\n",
    "                      arguments=[\"--dataset-type\", \"chipped\", \n",
    "                                 \"--scene-stats-csv\", f\"{input_stats.destination}/scene-stats.csv\",\n",
    "                                 \"--seed\", str(SEED), \n",
    "                                 \"--output-dir\", job_output.source, \n",
    "                                 \"--shoreline-dir\", inputs_shoreline.destination,\n",
    "                                 \"--gt-labels-dir\", str(Path(input_label_trn.destination).parent)],\n",
    "                      wait=True,\n",
    "                      logs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imagery Preparation with SageMaker Processing\n",
    "We use SageMaker Processing to prepare our imagery for training. \n",
    "The imagery data will be uploaded to the SageMaker session S3 bucket under `imagery`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### a. Save native scene imagery in file storage/\n",
    "For dynamically sampling from full scene imagery, we observed that we can speed up training and evaluation by a factor of 10 if the scene imagery was stored in `hdf5` format, compared to loading the provided GeoTIFF (Geostationary Earth Orbit Tagged Image File Format) imagery data with `rasterio`. This is also useful during inference for evaluation.\n",
    "\n",
    "Let's kick of SageMaker Processsing job to convert imagery to `hdf5`. This only needs to be done once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "instance_type = 'ml.t3.xlarge'\n",
    "volume_size_in_gb = 300 \n",
    "instance_count = 1 if USE_TINY else 75\n",
    "                      \n",
    "s3_uri_source = 's3://xview3-blog/data/raw'\n",
    "s3_uri_imagery = f'{s3_destination_uri}/imagery'\n",
    "\n",
    "storage_processor = Processor(image_uri=processing_image_name,\n",
    "                              role=role,\n",
    "                              instance_count=instance_count, \n",
    "                              base_job_name='xview3-storage',\n",
    "                              instance_type=instance_type, \n",
    "                              volume_size_in_gb=volume_size_in_gb,\n",
    "                              entrypoint=['python3', 'store_xview3_imagery.py'],\n",
    "                              sagemaker_session=sagemaker_session,\n",
    "                              tags=tags,)\n",
    "\n",
    "storage_processor.run(inputs=[ProcessingInput(source=s3_uri_source, \n",
    "                                              destination='/opt/ml/processing/input/',\n",
    "                                              s3_data_distribution_type='ShardedByS3Key')], \n",
    "                      outputs=[ProcessingOutput(source='/opt/ml/processing/output/imagery/', \n",
    "                                                destination=s3_uri_imagery,\n",
    "                                                output_name='imagery',\n",
    "                                                s3_upload_mode=\"Continuous\")],\n",
    "                      arguments=[\"--store-format\", \"hdf5\"],\n",
    "                      wait=False,\n",
    "                      logs=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### b. [Optional] Image chipping \n",
    "If we decide to train with image chips, we can also use SageMaker Processing to generate image chips using the dataset dict created in the previous section.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_destination_uri = f's3://xview3-blog/data/processing/{current_timestamp}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_uri_imagery = f'{s3_destination_uri}/imagery'\n",
    "s3_uri_imagery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processing_image_name = f'{account}.dkr.ecr.{region}.amazonaws.com/xview3-processing:main'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_uri_destination_base = f\"{s3_uri_imagery}/chipped-scenes\"\n",
    "s3_uri_source_base = \"s3://xview3-blog/data/raw\"\n",
    "s3_uri_d2_datasets = f'{s3_destination_uri}/detectron2_dataset/'\n",
    "\n",
    "\n",
    "d2_dataset_fn = f\"xview3-chipped_2560x2560-{'tiny' if USE_TINY else 'train'}.dataset\"\n",
    "num_instances = 2 if USE_TINY else 50 \n",
    "s3_uri_imagery_source = f\"{s3_uri_source_base}/{'tiny' if USE_TINY else 'trainval'}\"\n",
    "s3_uri_destination = f\"{s3_uri_destination_base}/{'tiny' if USE_TINY else 'train'}\"\n",
    "\n",
    "# specify local input data for SageMaker Processing job.\n",
    "input_scenes = ProcessingInput(source=s3_uri_imagery_source, \n",
    "                               destination='/opt/ml/processing/input/scenes/', \n",
    "                               s3_data_distribution_type='ShardedByS3Key')\n",
    "\n",
    "input_d2_dataset = ProcessingInput(source=s3_uri_d2_datasets, \n",
    "                                   destination='/opt/ml/processing/input/datasets/',)\n",
    "                                                \n",
    "job_output = ProcessingOutput(source='/opt/ml/processing/output/', \n",
    "                              destination=s3_uri_destination, \n",
    "                              s3_upload_mode=\"Continuous\",)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Need at least 32GB CPU instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chip_processor = Processor(image_uri=processing_image_name,\n",
    "                           role=role,\n",
    "                           instance_count=num_instances, \n",
    "                           base_job_name=f\"xview3-chip-scenes-{'tiny' if USE_TINY else 'train'}\", \n",
    "                           instance_type='ml.t3.2xlarge',#'ml.r5.xlarge', \n",
    "                           volume_size_in_gb=1024, \n",
    "                           entrypoint=['python3', 'chip_scenes_from_annotations.py'],\n",
    "                           sagemaker_session=sagemaker_session, \n",
    "                           tags=tags)\n",
    "\n",
    "chip_processor.run(inputs=[input_scenes, input_d2_dataset], \n",
    "                   outputs=[job_output],\n",
    "                   arguments=['--scenes-input-dir', input_scenes.destination,\n",
    "                              '--d2-dataset', f\"{input_d2_dataset.destination}/{d2_dataset_fn}\",],\n",
    "                   wait=USE_TINY,\n",
    "                   logs=USE_TINY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "from sagemaker.inputs import TrainingInput"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_CHIPPED = False\n",
    "LOCAL = False "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_train_dockerfile = str(Path(\"docker/training/base.Dockerfile\").resolve())\n",
    "train_dockerfile = str(Path(\"docker/training/main.Dockerfile\").resolve())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize -l docker {base_train_dockerfile}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build Base Training Container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "training_base_name = 'xview3-training:base'\n",
    "\n",
    "base_image_uri = build_and_push_docker_image(training_base_name,  \n",
    "                                             dockerfile=str(base_train_dockerfile),)\n",
    "print(f'Base image: {base_image_uri}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build Training Container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize -l docker {train_dockerfile}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_base_name = 'xview3-training:base'\n",
    "base_image_uri = f'{account}.dkr.ecr.{region}.amazonaws.com/{training_base_name}'\n",
    "training_main_name = 'xview3-training:train'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "training_image_uri = build_and_push_docker_image(training_main_name, \n",
    "                                                 dockerfile=str(train_dockerfile),\n",
    "                                                 base_image=base_image_uri)\n",
    "print(f'Training image: {training_image_uri}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_image_uri = f'{account}.dkr.ecr.{region}.amazonaws.com/xview3-training:train'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir='/opt/ml/model/FRCNN/auto'\n",
    "shoreline_dir  = '/opt/ml/input/data/shoreline/'\n",
    "\n",
    "metrics = [\n",
    "    {\"Name\": \"training:loss\", \"Regex\": \"total_loss: ([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"training:loss_cls\", \"Regex\": \"loss_cls: ([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"training:loss_box_reg\", \"Regex\": \"loss_box_reg: ([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"training:loss_rpn_cls\", \"Regex\": \"loss_rpn_cls: ([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"training:loss_rpn_loc\", \"Regex\": \"loss_rpn_loc: ([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"training:loss_length_reg\", \"Regex\": \"loss_length_reg: ([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"training:lr\", \"Regex\": \"lr: ([0-9\\\\.]+)\"},\n",
    "    {\"Name\": \"training:dataloader_time\", \"Regex\": \"data_time: ([0-9\\\\.]+)\"},\n",
    "    {\"Name\": \"training:time\", \"Regex\": \"time: ([0-9\\\\.]+)\"},\n",
    "    {\"Name\": \"validation:aggregate\", \"Regex\": \"aggregate=([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"validation:loc_fscore\", \"Regex\": \"loc_fscore=([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"validation:loc_fscore_shore\", \"Regex\": \"loc_fscore_shore=([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"validation:vessel_fscore\", \"Regex\": \"vessel_fscore=([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"validation:fishing_fscore\", \"Regex\": \"fishing_fscore=([0-9\\\\.]+)\",},\n",
    "    {\"Name\": \"validation:length_acc\", \"Regex\": \"length_acc=([0-9\\\\.]+)\",},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_iterations_from_epochs(epochs, bs, num_annotations, max_evals, warmup_prop, num_gpus=1):\n",
    "    iter_max = int(num_annotations / (num_gpus * bs) * epochs)\n",
    "    eval_period = iter_max//max_evals\n",
    "    iter_warmup = int(iter_max * warmup_prop)\n",
    "    \n",
    "    return iter_max, eval_period, iter_warmup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass(order=True)\n",
    "class Instances:\n",
    "    name: str\n",
    "    num_gpus: int = 1\n",
    "    instance_limit: int = 1\n",
    "    num_workers: int = 4\n",
    "    batch_size: int = 12\n",
    "    volume: int = 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_members = [Instances('local_gpu', num_gpus=4),\n",
    "                    Instances('ml.p3.2xlarge'), \n",
    "                    Instances('ml.p3.8xlarge', 4, 4, 16), \n",
    "                    Instances('ml.p3.16xlarge', 8, 2, 32),\n",
    "                    Instances('ml.p3dn.24xlarge', 8, num_workers=48, batch_size=24, volume=1800)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_ANNOTS = {'tiny': 1679, \n",
    "              'train': 54360}\n",
    "\n",
    "if USE_CHIPPED:\n",
    "    NUM_ANNOTS['tiny'] = 1907\n",
    "    NUM_ANNOTS['train'] = 62766"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance = instance_members[-2]\n",
    "instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 6\n",
    "num_annotations = NUM_ANNOTS['tiny'] if USE_TINY else NUM_ANNOTS['train']\n",
    "bs = instance.batch_size\n",
    "#num_gpus = 1 #instance.num_gpus\n",
    "max_evals = 5\n",
    "max_checkpoints = max_evals * 2\n",
    "warmup_prop = 0.2\n",
    "\n",
    "max_iter, eval_period, warmup_iter = compute_iterations_from_epochs(epochs, bs, num_annotations, num_gpus, max_evals, warmup_prop)\n",
    "checkpoint_period = eval_period // 2\n",
    "print(max_iter, eval_period, warmup_iter, checkpoint_period)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Datasets\n",
    "mode = \"tiny\" if USE_TINY else \"trainval\"\n",
    "imagery_s3_uri = f's3://xview3-blog/data/processing/202207250702/imagery/hdf5/{mode}/'\n",
    "\n",
    "if USE_CHIPPED:\n",
    "    imagery_s3_uri = f's3://xview3-blog/data/processing/202207250702/imagery/chipped-scenes/{mode}/xview3_chipped_2560x2560_{mode.replace(\"val\", \"\")}/'\n",
    "    val_imagery_s3_uri = f's3://xview3-blog/data/processing/202207250702/imagery/hdf5/{mode}/'\n",
    "    s3_channel_valid_imagery = TrainingInput(val_imagery_s3_uri, \n",
    "                                   distribution='FullyReplicated', \n",
    "                                   s3_data_type='S3Prefix',\n",
    "                                   input_mode='FastFile')\n",
    "    \n",
    "shoreline_s3_uri = 's3://xview3-blog/data/shoreline/trainval/'\n",
    "datasets_s3_uri = 's3://xview3-blog/data/processing/202207250702/detectron2_dataset/'\n",
    "\n",
    "s3_channel_imagery = TrainingInput(imagery_s3_uri, \n",
    "                                   distribution='FullyReplicated', \n",
    "                                   s3_data_type='S3Prefix',\n",
    "                                   input_mode='FastFile')\n",
    "s3_channel_shoreline = TrainingInput(shoreline_s3_uri, \n",
    "                                     distribution='FullyReplicated', \n",
    "                                     s3_data_type='S3Prefix', \n",
    "                                     input_mode='FastFile')\n",
    "s3_channel_datasets = TrainingInput(datasets_s3_uri, \n",
    "                                    distribution='FullyReplicated', \n",
    "                                    s3_data_type='S3Prefix',\n",
    "                                    input_mode='FastFile')\n",
    "\n",
    "train_inputs = {'imagery': s3_channel_imagery, \n",
    "                'shoreline': s3_channel_shoreline, \n",
    "                'datasets': s3_channel_datasets}\n",
    "if USE_CHIPPED:\n",
    "    train_inputs['valid_imagery'] = s3_channel_valid_imagery\n",
    "\n",
    "# Use EFS if local\n",
    "if LOCAL:\n",
    "    train_inputs['imagery'] = f'file:////home/ec2-user/SageMaker/xview3-blog/data/imagery/hdf5/tiny/'\n",
    "    train_inputs['shoreline'] = 'file:///home/ec2-user/SageMaker/xview3-blog/data/shoreline/trainval/'\n",
    "    train_inputs['datasets'] = 'file:///home/ec2-user/SageMaker/xview3-blog/data/detectron2_datasets/new/'\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = 'frcnn_X101_32x8d_FPN_full.yaml'#'frcnn_R101_FPN_full.yaml'#'frcnn_R101_FPN_full_VH3.yaml' \n",
    "if USE_CHIPPED:\n",
    "    config_file = 'frcnn_R101_FPN_chipped_histeq.yaml'\n",
    "\n",
    "config_params = [f'OUTPUT_DIR {output_dir}',\n",
    "                 f'TEST.INPUT.SHORELINE_DIR {shoreline_dir}',\n",
    "                 f'INPUT.DATA.SHORELINE_DIR {shoreline_dir}',\n",
    "                 f\"SOLVER.IMS_PER_BATCH {bs}\",\n",
    "                 f\"TEST.EVAL_PERIOD {eval_period}\",\n",
    "                 f\"SOLVER.WARMUP_ITERS {warmup_iter}\",\n",
    "                 f\"SOLVER.MAX_ITER {max_iter}\",\n",
    "                 f\"SOLVER.CHECKPOINT_PERIOD {checkpoint_period}\",\n",
    "                 f\"DATALOADER.NUM_WORKERS {instance.num_workers}\",\n",
    "                 \"SOLVER.LR_SCHEDULER_NAME WarmupCosineLR\",\n",
    "                 \"SOLVER.BASE_LR 0.005\",\n",
    "                ]\n",
    "\n",
    "training_job_hp = {'config-file': f'/opt/ml/code/configs/{config_file}',\n",
    "                   'imagery-dir': '/opt/ml/input/data/imagery',\n",
    "                   'd2-dataset-dir': '/opt/ml/input/data/datasets',\n",
    "                   'zopts': ' '.join(config_params)}\n",
    "\n",
    "if USE_CHIPPED:\n",
    "    training_job_hp['valid-imagery-dir'] = '/opt/ml/input/data/valid_imagery'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#base_job_name = f\"xview3-{'chipped' if USE_CHIPPED else 'full'}-{'tiny' if USE_TINY else 'trainval'}\"\n",
    "base_job_name = f\"xview3-{config_file.split('.')[0].replace('_', '-')}\"\n",
    "\n",
    "training_instance = instance.name\n",
    "num_instances = 1\n",
    "training_session = sagemaker_session\n",
    "\n",
    "\n",
    "if training_instance.startswith(\"local\"):\n",
    "    training_session = sagemaker.LocalSession()\n",
    "    training_session.config = {\"local\": {\"local_code\": True}}\n",
    "    LOCAL = True\n",
    "\n",
    "d2_estimator = Estimator(image_uri=training_image_uri,\n",
    "                         role=role, \n",
    "                         sagemaker_session=training_session, \n",
    "                         instance_count=num_instances, \n",
    "                         instance_type=training_instance, \n",
    "                         volume_size=instance.volume,\n",
    "                         metric_definitions=metrics, \n",
    "                         hyperparameters=training_job_hp,\n",
    "                         base_job_name=base_job_name, \n",
    "                         max_retry_attempts=30, \n",
    "                         max_run=432000,\n",
    "                         checkpoint_local_path=None if LOCAL else '/opt/ml/checkpoints/' ,\n",
    "                         checkpoint_s3_uri=None if LOCAL else 's3://xview3-blog-sagemaker/checkpoints/',\n",
    "                         disable_profiler=True,\n",
    "                         debugger_hook_config=False,\n",
    "                        tags=tags)\n",
    "\n",
    "d2_estimator.fit(inputs=train_inputs, \n",
    "                 wait=True if USE_TINY else False, \n",
    "                 logs=\"All\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`/tmp/tmp7dix_o_f/algo-1-2i620`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.7 ('xview3')",
   "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.9.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "2e86a4bdce09334967ec02703ecb4f73f52314f4bf74047e8184a512364b7029"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}