 archived/Amazon_JumpStart_NLP_Regression_Free_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b6a4199a",
   "metadata": {},
   "source": [
    "# Backward Compatibility During Data Updates by Weight Interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f18b446a",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58ef3f69",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa319f72",
   "metadata": {},
   "source": [
    "Backward compatibility of model predictions is a desired property when updating a machine learning driven application. It allows to seamlessly improve the underlying model without introducing regression bugs. In classification tasks these bugs occur in the form of negative flips: A new model incorrectly predicts the output for a test sample that was correctly classified by the old (reference) model.\n",
    "\n",
    "A common reason to update the model is when new training data becomes available and needs to be incorporated. Simply retraining the model with the updated data introduces the unwanted negative flips. \n",
    "\n",
    "In this notebook we introduce and implement Backward Compatible Weight Interpolation (BCWI) that reduces reduce regression during data updates. This method interpolates between the weights of the old and new model.\n",
    "\n",
    "As demonstrated in the notebook implementation below, our BCWI technique provides signigicant 40% reduction in negative flip rate. Further details of the method will be released in the upcoming paper.\n",
    "\n",
    "If you find this notebook useful please consider citing our prepring\n",
    "```\n",
    "@article{Schumann2023BCWI,\n",
    "title={Backward Compatibility During Data Updates by Weight Interpolation},\n",
    "author={Raphael Schumann and Elman Mansimov and Yi-An Lai and Nikolaos Pappas and Xibin Gao and Yi Zhang},\n",
    "journal={ArXiv},\n",
    "year={2023}}\n",
    "```\n",
    "\n",
    "Also, we suggest you checkout out related notebook on improving backward-compatibility during model achitecture updates [LINK](https://github.com/aws/amazon-sagemaker-examples/blob/main/introduction_to_amazon_algorithms/jumpstart_regression_free_training/Amazon_JumpStart_Regression_Free_Training.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e75356b",
   "metadata": {},
   "source": [
    "## Install required python libraries (transformers and datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc5bd3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install huggingface transformers and datasets\n",
    "!pip install transformers\n",
    "!pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cedeb6ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import json\n",
    "import os\n",
    "import torch\n",
    "import transformers\n",
    "import datasets\n",
    "import requests\n",
    "import copy\n",
    "\n",
    "from transformers import RobertaTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c87166f",
   "metadata": {},
   "source": [
    "## Shared variables pointing to the data url and model name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa298a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Shared Variables\n",
    "DATA_URL = \"https://raw.githubusercontent.com/amazon-science/regression-constraint-model-upgrade/main/nlp/data/MASSIVE/\"\n",
    "PT_MODEL_NAME = \"roberta-base\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ac09f0a",
   "metadata": {},
   "source": [
    "## Setup tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "243c7ad1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load tokenizer\n",
    "tokenizer = []\n",
    "\n",
    "\n",
    "def load_tokenizer():\n",
    "    if len(tokenizer) == 0:\n",
    "        tokenizer.append(RobertaTokenizer.from_pretrained(PT_MODEL_NAME))\n",
    "\n",
    "\n",
    "load_tokenizer()\n",
    "\n",
    "print(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3b1f795",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_function(examples):\n",
    "    \"\"\"Helper tokenize function\"\"\"\n",
    "    return tokenizer[0](\n",
    "        examples[\"text\"], padding=False, truncation=True, return_attention_mask=False\n",
    "    )\n",
    "\n",
    "\n",
    "def load_dataset(splits, file_template):\n",
    "    \"\"\"Helper function to load MASSIVE dataset\"\"\"\n",
    "    data_files = dict()\n",
    "    for split in splits:\n",
    "        data_files[split] = file_template.format(split)\n",
    "\n",
    "    print(data_files)\n",
    "    dataset = datasets.load_dataset(\"json\", data_files=data_files)\n",
    "    dataset_info = json.loads(requests.get(file_template.format(\"info\")[:-1]).content)\n",
    "\n",
    "    new_label_ids = [dataset_info[\"labels\"].index(c) for c in dataset_info[\"add_classes\"]]\n",
    "    old_label_ids = [\n",
    "        i for i, c in enumerate(dataset_info[\"labels\"]) if c not in dataset_info[\"add_classes\"]\n",
    "    ]\n",
    "    dataset_info[\"new_label_ids\"] = new_label_ids\n",
    "    dataset_info[\"old_label_ids\"] = old_label_ids\n",
    "\n",
    "    tokenized_dataset = dataset.map(tokenize_function, batched=False)\n",
    "\n",
    "    return tokenized_dataset, dataset_info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59689bf5",
   "metadata": {},
   "source": [
    "## Download the initial dataset used before new data is introduced (old dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf5089b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPLETE: Pull the data before update\n",
    "old_dataset_files = os.path.join(DATA_URL, \"add_data\", \"old\", \"{}.jsonl\")\n",
    "old_dataset, old_dataset_info = load_dataset([\"train\", \"dev\", \"test\"], old_dataset_files)\n",
    "\n",
    "# should contain 1000 lines in train, 333 lines in dev, and 4000 lines in test\n",
    "print(\"Old dataset before update\")\n",
    "print(old_dataset)\n",
    "\n",
    "print(\"Old dataset before update info\")\n",
    "print(old_dataset_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a162296",
   "metadata": {},
   "source": [
    "## Set up SageMaker and SageMaker Training Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7407800",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Setup Sagemaker environment\n",
    "import sagemaker\n",
    "from sagemaker.huggingface import HuggingFace\n",
    "import boto3\n",
    "\n",
    "# Use remote mode\n",
    "sagemaker_region = boto3.Session().region_name\n",
    "sagemaker_session = sagemaker.Session()\n",
    "bucket_name = sagemaker_session.default_bucket()\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "instance_type = \"ml.p3.2xlarge\"\n",
    "\n",
    "print(f\"sagemaker region: {sagemaker_region}\")\n",
    "print(f\"sagemaker session: {sagemaker_session}\")\n",
    "print(f\"bucket name: {bucket_name}\")\n",
    "print(f\"role: {role}\")\n",
    "print(f\"instance type: {instance_type}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44225ea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPLETE: Train the model using the data before update\n",
    "# Something along these lines\n",
    "\n",
    "# git configuration to download regression-free training script\n",
    "git_config = {\n",
    "    \"repo\": \"https://github.com/amazon-science/regression-constraint-model-upgrade.git\",\n",
    "    \"branch\": \"main\",\n",
    "}\n",
    "\n",
    "huggingface_estimator = HuggingFace(\n",
    "    entry_point=\"train.py\",\n",
    "    source_dir=\"nlp\",\n",
    "    git_config=git_config,\n",
    "    instance_type=instance_type,\n",
    "    instance_count=1,\n",
    "    role=role,\n",
    "    transformers_version=\"4.17.0\",\n",
    "    pytorch_version=\"1.10.2\",\n",
    "    py_version=\"py38\",\n",
    "    hyperparameters={\n",
    "        \"dataset\": \"MASSIVE\",\n",
    "        \"scenario\": \"add_data\",\n",
    "        \"data_type\": \"old\",\n",
    "        \"bucket_name\": bucket_name,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bf103f7",
   "metadata": {},
   "source": [
    "## Train the initial model before data update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb97e7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "huggingface_estimator.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e670f365",
   "metadata": {},
   "source": [
    "## Pull the trained model from S3 in order to initialize the new model and inference "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef60f780",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPLETE: Load old model from S3\n",
    "# Used for inference and calculating negative flip rate\n",
    "import boto3\n",
    "\n",
    "s3 = boto3.client(\"s3\")\n",
    "\n",
    "load_from_s3 = f\"./bcwi_nlp_outputs/v1/MASSIVE/1111/add_data/old_model/model\"\n",
    "\n",
    "print(f\"Loading from S3 {load_from_s3}\")\n",
    "os.makedirs(\"old_model_dir\", exist_ok=True)\n",
    "# load stuff from s3\n",
    "with open(\"old_model_dir/config.json\", \"wb\") as f:\n",
    "    s3.download_fileobj(bucket_name, os.path.join(load_from_s3, \"config.json\"), f)\n",
    "    print(\"Downloaded old_model_dir/config.json\")\n",
    "with open(\"old_model_dir/hparams.json\", \"wb\") as f:\n",
    "    s3.download_fileobj(bucket_name, os.path.join(load_from_s3, \"hparams.json\"), f)\n",
    "    print(\"Downloaded old_model_dir/hparams.json\")\n",
    "with open(\"old_model_dir/pytorch_model.bin\", \"wb\") as f:\n",
    "    s3.download_fileobj(bucket_name, os.path.join(load_from_s3, \"pytorch_model.bin\"), f)\n",
    "    print(\"Downloaded old_model_dir/pytorch_model.bin\")\n",
    "\n",
    "from transformers import RobertaForSequenceClassification\n",
    "\n",
    "old_model = RobertaForSequenceClassification.from_pretrained(\"old_model_dir\")\n",
    "\n",
    "print(old_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa351b60",
   "metadata": {},
   "source": [
    "## Download the updated dataset with new data added to it (new dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0c0e5a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPLETE: Pull the data after update\n",
    "new_dataset_files = os.path.join(DATA_URL, \"add_data\", \"updated\", \"{}.jsonl\")\n",
    "new_dataset, new_dataset_info = load_dataset([\"train\", \"dev\", \"test\"], new_dataset_files)\n",
    "\n",
    "# should contain 1500 lines in train, 500 lines in dev, and 4000 lines in test\n",
    "print(\"New dataset after update\")\n",
    "print(new_dataset)\n",
    "\n",
    "print(\"New dataset after update info\")\n",
    "print(new_dataset_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23f1812d",
   "metadata": {},
   "source": [
    "## Train the new model using updated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4842c478",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPLETE: Get the old data checkpoint and continue training on the new data\n",
    "\n",
    "# git configuration to download regression-free training script\n",
    "git_config = {\n",
    "    \"repo\": \"https://github.com/amazon-science/regression-constraint-model-upgrade.git\",\n",
    "    \"branch\": \"main\",\n",
    "}\n",
    "\n",
    "huggingface_estimator_new = HuggingFace(\n",
    "    entry_point=\"train.py\",\n",
    "    source_dir=\"nlp\",\n",
    "    git_config=git_config,\n",
    "    instance_type=instance_type,\n",
    "    instance_count=1,\n",
    "    role=role,\n",
    "    transformers_version=\"4.17.0\",\n",
    "    pytorch_version=\"1.10.2\",\n",
    "    py_version=\"py38\",\n",
    "    hyperparameters={\n",
    "        \"dataset\": \"MASSIVE\",\n",
    "        \"scenario\": \"add_data\",\n",
    "        \"data_type\": \"updated\",\n",
    "        \"load_from_s3\": f\"./bcwi_nlp_outputs/v1/MASSIVE/1111/add_data/old_model/model\",\n",
    "        \"bucket_name\": bucket_name,\n",
    "        \"output_dir\": \"add_data\",\n",
    "        \"num_epochs\": 3,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3467d721",
   "metadata": {},
   "outputs": [],
   "source": [
    "huggingface_estimator_new.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "706aa625",
   "metadata": {},
   "source": [
    "## Pull the newly trained model from S3 for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d09b6f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPLETE: Load new model from S3\n",
    "# Used for inference and calculating negative flip rate\n",
    "\n",
    "load_from_s3 = f\"./bcwi_nlp_outputs/add_data/MASSIVE/1111/add_data/old_model/model\"\n",
    "\n",
    "print(f\"Loading from S3 {load_from_s3}\")\n",
    "os.makedirs(\"new_model_dir\", exist_ok=True)\n",
    "# load stuff from s3\n",
    "with open(\"new_model_dir/config.json\", \"wb\") as f:\n",
    "    s3.download_fileobj(bucket_name, os.path.join(load_from_s3, \"config.json\"), f)\n",
    "    print(\"Downloaded new_model_dir/config.json\")\n",
    "with open(\"new_model_dir/hparams.json\", \"wb\") as f:\n",
    "    s3.download_fileobj(bucket_name, os.path.join(load_from_s3, \"hparams.json\"), f)\n",
    "    print(\"Downloaded new_model_dir/hparams.json\")\n",
    "with open(\"new_model_dir/pytorch_model.bin\", \"wb\") as f:\n",
    "    s3.download_fileobj(bucket_name, os.path.join(load_from_s3, \"pytorch_model.bin\"), f)\n",
    "    print(\"Downloaded new_model_dir/pytorch_model.bin\")\n",
    "\n",
    "from transformers import RobertaForSequenceClassification\n",
    "\n",
    "new_model = RobertaForSequenceClassification.from_pretrained(\"new_model_dir\")\n",
    "\n",
    "print(new_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a46448b4",
   "metadata": {},
   "source": [
    "## Helper function to calculate accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d24909b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare for calculating negative flip rate and accuracy\n",
    "test_data = new_dataset[\"test\"]  # test sets in old and new datasets are the same\n",
    "\n",
    "\n",
    "def calculate_accuracy(data, model, batch_size=40):\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    for i in range(0, len(data), batch_size):\n",
    "        if int(i / batch_size) % 10 == 0:\n",
    "            print(f\"working on {int(i/batch_size)} out of {int(len(data)/batch_size)}\")\n",
    "        torch.cuda.empty_cache()  # clear memory\n",
    "        # process examples in the batch\n",
    "        examples = data[i : (i + batch_size)]\n",
    "        text = examples[\"text\"]\n",
    "        label = examples[\"label\"]\n",
    "        text_tokenizer = tokenizer[0](\n",
    "            text,\n",
    "            padding=\"max_length\",\n",
    "            truncation=True,\n",
    "            max_length=max([len(t.split()) for t in text]),\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**text_tokenizer)\n",
    "        preds = outputs.logits.argmax(-1).tolist()\n",
    "\n",
    "        # merge them into the list that combines all labels and predictions\n",
    "        all_preds.extend(preds)\n",
    "        all_labels.extend(label)\n",
    "    return all_preds, all_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa70e45a",
   "metadata": {},
   "source": [
    "## Calculate accuracy of old and new models together with a negative flip rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1366a7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get old model prediction\n",
    "print(\"Getting old model predictions\")\n",
    "old_model_preds, labels = calculate_accuracy(test_data, old_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8868a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get new model prediction\n",
    "print(\"Getting new model predictions\")\n",
    "new_model_preds, _ = calculate_accuracy(test_data, new_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b517ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get old model accuracy\n",
    "old_acc = (\n",
    "    100 * sum([old_pred == l for old_pred, l in zip(old_model_preds, labels)]) / float(len(labels))\n",
    ")\n",
    "print(f\"Old model accuracy {old_acc}%\")\n",
    "\n",
    "# Get new model accuracy\n",
    "new_acc = (\n",
    "    100 * sum([new_pred == l for new_pred, l in zip(new_model_preds, labels)]) / float(len(labels))\n",
    ")\n",
    "print(f\"New model accuracy {new_acc}%\")\n",
    "\n",
    "# Calculate negative flip rate\n",
    "nfr = (\n",
    "    100\n",
    "    * sum(\n",
    "        [\n",
    "            old_pred == l and new_pred != l\n",
    "            for old_pred, new_pred, l in zip(old_model_preds, new_model_preds, labels)\n",
    "        ]\n",
    "    )\n",
    "    / float(len(labels))\n",
    ")\n",
    "print(f\"Negative Flip Rate {nfr}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "681b9f7d",
   "metadata": {},
   "source": [
    "## Backward compatible weight interpolation (our core method to reduce regression)\n",
    "\n",
    "Backward compatible weight interpolation (BCWI) is defined as \n",
    "\n",
    "\\begin{equation}\n",
    "    \\theta_{\\mathrm{BCWI}} = \\alpha \\theta_{old} + (1-\\alpha) \\theta_{new},\n",
    "\\end{equation}\n",
    "\n",
    "Where $\\alpha \\in [0.0, 1.0]$, $\\theta_{old}$ are old model parameters and $\\theta_{new}$ are new model parameters.\n",
    "\n",
    "We use $\\alpha = 0.3$ in this notebook "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a39e7099",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interpolate between old and new model\n",
    "# More details can be found in the Github repo\n",
    "\n",
    "\n",
    "def interpolate_weights(old_model, new_models, alpha, new_label_ids=None, weighted=None):\n",
    "    # Form soup ensemble of new models\n",
    "    new_state_dicts = [new_model.state_dict() for new_model in new_models]\n",
    "    new_model_state_dict = dict()\n",
    "    for key in new_models[0].state_dict():\n",
    "        if not (key.endswith(\"bias\") or key.endswith(\"weight\")):\n",
    "            continue\n",
    "\n",
    "        new_model_state_dict[key] = torch.mean(\n",
    "            torch.stack([s[key] for s in new_state_dicts]), dim=0\n",
    "        )\n",
    "\n",
    "    print(\"alpha\", alpha)\n",
    "    metrics = dict()\n",
    "\n",
    "    # Use the old model as the basis of the interpolated model weights\n",
    "    model = copy.deepcopy(old_model)\n",
    "    # All weights of a model can be accessed by its state_dict\n",
    "    state_dict = model.state_dict()\n",
    "    for key in state_dict:\n",
    "        # Be sure to only interpolate weight matrices; includes e.g. layer norm matrices\n",
    "        if not (key.endswith(\"bias\") or key.endswith(\"weight\")):\n",
    "            continue\n",
    "\n",
    "        if weighted is not None:\n",
    "            # when alpha = 1.0, there can be NaN values due to numerical instabilities when values in the weight\n",
    "            # matrix are too small. In this case we replace the NaNs with the weights of the old model.\n",
    "            if alpha == 1.0:\n",
    "                c = state_dict[key].detach().clone()\n",
    "\n",
    "            # Inplace operations to modify the weights of the model.\n",
    "            # State_dict initially holds the weights of the old model.\n",
    "            state_dict[key] *= alpha * weighted[key]\n",
    "            state_dict[key] += (1 - alpha) * new_model_state_dict[key]\n",
    "            state_dict[key] /= alpha * weighted[key] + (1 - alpha)\n",
    "\n",
    "            # Three lines above as one-liner\n",
    "            # state_dict[key].data.copy_(((alpha * weighted[key] * state_dict[key]) + ((1-alpha) * new_model_state_dict[key])) / (alpha * weighted[key] + (1-alpha)))\n",
    "\n",
    "            if alpha == 1.0:\n",
    "                nans = state_dict[key] != state_dict[key]\n",
    "                state_dict[key][nans] = c[nans]\n",
    "        else:\n",
    "            # Simple linear interpolation with parameter alpha.\n",
    "            # State_dict initially holds the weights of the old model.\n",
    "            state_dict[key] *= alpha\n",
    "            state_dict[key] += (1 - alpha) * new_model_state_dict[key]\n",
    "\n",
    "        # Copy classifier weights of new classes from the new model. The old model was not trained on those classes.\n",
    "        if new_label_ids:\n",
    "            if key == \"classifier.out_proj.weight\":\n",
    "                state_dict[key][new_label_ids, :] = new_model_state_dict[key][new_label_ids, :]\n",
    "            if key == \"classifier.out_proj.bias\":\n",
    "                state_dict[key][new_label_ids] = new_model_state_dict[key][new_label_ids]\n",
    "    return state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a7fb52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interpolated model\n",
    "interpolated_model = RobertaForSequenceClassification.from_pretrained(\"old_model_dir\")\n",
    "\n",
    "# Initialize interpolated model with old model weights\n",
    "interpolated_state_dict = interpolate_weights(interpolated_model, [new_model], alpha=0.3)\n",
    "# load interpolated_state_dict into new model\n",
    "interpolated_model.load_state_dict(interpolated_state_dict, strict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68945504",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get interpolated model prediction\n",
    "print(\"Getting interpolated model predictions\")\n",
    "interpolated_model_preds, _ = calculate_accuracy(test_data, interpolated_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f550633",
   "metadata": {},
   "source": [
    "## Calculate accuracy and negative flip rate of interpolated model. Compare it to the baselines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b46f81e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get old model accuracy\n",
    "old_acc = (\n",
    "    100 * sum([old_pred == l for old_pred, l in zip(old_model_preds, labels)]) / float(len(labels))\n",
    ")\n",
    "print(f\"Old model accuracy {old_acc}%\")\n",
    "\n",
    "# Get new model accuracy\n",
    "new_acc = (\n",
    "    100 * sum([new_pred == l for new_pred, l in zip(new_model_preds, labels)]) / float(len(labels))\n",
    ")\n",
    "print(f\"New model accuracy {new_acc}%\")\n",
    "\n",
    "# Get interpolated model accuracy\n",
    "interpolated_acc = (\n",
    "    100\n",
    "    * sum(\n",
    "        [interpolated_pred == l for interpolated_pred, l in zip(interpolated_model_preds, labels)]\n",
    "    )\n",
    "    / float(len(labels))\n",
    ")\n",
    "print(f\"Interpolated model accuracy {interpolated_acc}%\")\n",
    "\n",
    "\n",
    "# Calculate negative flip rate\n",
    "nfr = (\n",
    "    100\n",
    "    * sum(\n",
    "        [\n",
    "            old_pred == l and new_pred != l\n",
    "            for old_pred, new_pred, l in zip(old_model_preds, new_model_preds, labels)\n",
    "        ]\n",
    "    )\n",
    "    / float(len(labels))\n",
    ")\n",
    "print(f\"Negative Flip Rate {nfr}%\")\n",
    "\n",
    "# Calculate negative flip rate of old model and interpolated models\n",
    "interpolate_nfr = (\n",
    "    100\n",
    "    * sum(\n",
    "        [\n",
    "            old_pred == l and interpolate_pred != l\n",
    "            for old_pred, interpolate_pred, l in zip(\n",
    "                old_model_preds, interpolated_model_preds, labels\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "    / float(len(labels))\n",
    ")\n",
    "print(f\"Negative Flip Rate Old and Interpolated (Ours) {interpolate_nfr}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04398680",
   "metadata": {},
   "source": [
    "At the running the last cell of the notebook you should see the following outputs:\n",
    "\n",
    "```\n",
    "Old model accuracy 81.7%\n",
    "New model accuracy 82.55%\n",
    "Interpolated model accuracy 82.65%\n",
    "Negative Flip Rate 3.25%\n",
    "Negative Flip Rate Old and Interpolated (Ours) 2.225%\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09d43128",
   "metadata": {},
   "source": [
    "Hope you found the implementation helpful!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5218db12",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/introduction_to_amazon_algorithms|jumpstart_regression_free_training|Amazon_JumpStart_NLP_Regression_Free_Training.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   }
  ],
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.13 Python 3.9 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/pytorch-1.13-cpu-py39"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}