notebooks/hls-fm-finteuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8fb85522-53a4-43fb-95bc-b37d953662c4",
   "metadata": {},
   "source": [
    "# HLS Foundation Model Finetuning notebook\n",
    "\n",
    "This notebook demonstrates the steps to fintune the HLS foundation model (A.K.A Prithvi) which is trained using HLSL30 and HLSS30 datasets. \n",
    "\n",
    "Note: Entierty of this notebook is desigend to work well within the AWS sagemaker environment. AWS sagemaker environment access for your account can be found using https://creds-workshop.nasa-impact.net/.\n",
    "\n",
    "![HLS Training](../images/HLS-training.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77b9afa3-5424-42dc-bfcb-800df0435f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "!pip install -r requirements.txt\n",
    "\n",
    "# Create directories needed for data, model, and config preparations\n",
    "!mkdir datasets\n",
    "!mkdir models\n",
    "!mkdir configs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c4f1ad",
   "metadata": {},
   "source": [
    "### Download HLS Burn Scars dataset from Huggingface: https://huggingface.co/datasets/ibm-nasa-geospatial/hls_burn_scars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c0c4e10-15cd-4f52-8dfe-ec04074efe43",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! cd datasets; git clone https://huggingface.co/datasets/ibm-nasa-geospatial/hls_burn_scars; tar -xvzf hls_burn_scars/hls_burn_scars.tar.gz "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ef75557",
   "metadata": {},
   "source": [
    "## Download config and Pre-trained model\n",
    "\n",
    "The HLS Foundation Model (pre-trained model), and configuration for Burn Scars downstream task are available in Huggingface. We use `huggingface_hub` python package to download the files locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe11e274",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define constants\n",
    "BUCKET_NAME = '<your-bucket-name>' # Replace this with the bucket name available from https://creds-workshop.nasa-impact.net/\n",
    "CONFIG_PATH = './configs'\n",
    "DATASET_PATH = './datasets'\n",
    "MODEL_PATH = './models'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8f74b4d",
   "metadata": {},
   "source": [
    "Note: The configuration file in Huggingface have place holders and won't be directly usable for fine-tuning. Placeholder values need to be updated for individual usecases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb5d6b5d-d58d-4b90-a952-6179c255280e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download config file from huggingface\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "hf_hub_download(repo_id=\"ibm-nasa-geospatial/Prithvi-100M-burn-scar\", filename=\"burn_scars_Prithvi_100M.py\", local_dir='./configs')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90acf449",
   "metadata": {},
   "source": [
    "*Warning: * Before running the remaining cells please update the details in the configuration file as mentioned below:\n",
    "\n",
    "1. Update line number 13 from `data_root = '<path to data root>'` to `data_root = '/opt/ml/data/'`. This is the base of our data inside of sagemaker.\n",
    "2. Update line number 41 from `pretrained_weights_path = '<path to pretrained weights>'` to `pretrained_weights_path = f\"{data_root}/models/prithvi-global-300M.pt\"`. This provides the pre-trained model path to the train script.\n",
    "3. Update line number 52 from `experiment = '<experiment name>'` to `experiment = 'burn_scars'` or your choice of experiment name.\n",
    "4. Update line number 53 from `project_dir = '<project directory name>'` to `project_dir = 'v1'` or your choice of project directory name. \n",
    "5. Save the config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ff4ddab-fd48-4015-9822-17047d3a4039",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare sagemaker session with files uploaded to s3 bucket\n",
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "train_images = sagemaker_session.upload_data(path='datasets/training', bucket=BUCKET_NAME, key_prefix='data/training')\n",
    "val_images = sagemaker_session.upload_data(path='datasets/validation', bucket=BUCKET_NAME, key_prefix='data/validation')\n",
    "test_images = sagemaker_session.upload_data(path='datasets/validation', bucket=BUCKET_NAME, key_prefix='data/test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01d77ad-728e-430c-bf59-f3cb36e19592",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rename configuration file name to user specific filename\n",
    "import os\n",
    "\n",
    "identifier = '<your choice of identifier>' # Please update this with an identifier\n",
    "\n",
    "config_filename = 'configs/burn_scars_Prithvi_100M.py'\n",
    "new_config_filename = f\"configs/{identifier}-burn_scars_Prithvi_100M.py\"\n",
    "os.rename(config_filename, new_config_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a2384e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload config files to s3 bucket\n",
    "configs = sagemaker_session.upload_data(path=new_config_filename, bucket=BUCKET_NAME, key_prefix='data/configs')\n",
    "models = sagemaker_session.upload_data(path='../models/prithvi-global-300M.pt', bucket=BUCKET_NAME, key_prefix='data/models')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e69e63",
   "metadata": {},
   "source": [
    "\n",
    "Note: For HLS Foundation Model, MMCV and MMSEG were used. These libraries use pytorch underneath them for training, data distribution etc. However, these packages are not available in sagemaker by default. Thus, custom script training is required. Sagemaker utilizes Docker for custom training scripts. If interested, the code included in the image we are using for training (637423382292.dkr.ecr.us-west-2.amazonaws.com/sagemaker_hls:latest) is bundled with this repository, and the train script used is `train.py`.\n",
    "\n",
    "The current HLS Foundation model fits in a single NVIDIA Tesla V100 GPU (16GB VRAM). Hence, `ml.p3.2xlarge` instance is used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18302154-eee5-4705-b1c4-11cfd47ef6f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup variables for training using sagemaker\n",
    "from datetime import time\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "\n",
    "name = f'{identifier}-sagemaker'\n",
    "role = get_execution_role()\n",
    "input_s3_uri = f\"s3://{BUCKET_NAME}/data\"\n",
    "\n",
    "environment_variables = {\n",
    "    'CONFIG_FILE': f\"/opt/ml/data/{new_config_filename}\",\n",
    "    'MODEL_DIR': \"/opt/ml/models/\",\n",
    "    'MODEL_NAME': f\"{identifier}-workshop.pth\",\n",
    "    'S3_URL': input_s3_uri,\n",
    "    'ROLE_ARN': role,\n",
    "    'ROLE_NAME': role.split('/')[-1],\n",
    "    'EVENT_TYPE': 'burn_scars',\n",
    "    'VERSION': 'v1'\n",
    "}\n",
    "\n",
    "ecr_container_url = '637423382292.dkr.ecr.us-west-2.amazonaws.com/sagemaker_hls:latest'\n",
    "sagemaker_role = 'SageMaker-ExecutionRole-20240206T151814'\n",
    "\n",
    "instance_type = 'ml.p3.2xlarge'\n",
    "\n",
    "instance_count = 1\n",
    "memory_volume = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37bbbf5a-6df2-44d7-a8bc-81026631d41e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Establish an estimator (model) using sagemaker and the configurations from the previous cell.\n",
    "estimator = Estimator(image_uri=ecr_container_url,\n",
    "                      role=get_execution_role(),\n",
    "                      base_job_name=name,\n",
    "                      instance_count=1,\n",
    "                      environment=environment_variables,\n",
    "                      instance_type=instance_type)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82c98a7a-f199-49c6-a655-38c19dd28689",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Start training\n",
    "estimator.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc35c27d-eba9-4c9a-a28f-292848b982b0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "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.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "8987afb44532b2110e1a5e1b229dd281f8440b44477d285826a54acdd52d8797"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}