cli_yaml/exercise04_train_remote.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise04 : Train on Remote GPU Virtual Machine\n",
    "\n",
    "Now we run our previous sample (see \"[Exercise03 : Just Train in Your Working Machine](./exercise03_train_simple.ipynb)\") on remote virtual machine with GPU utilized.\n",
    "\n",
    "> Note : If you don't have GPU quota, you can also run this example on CPU.\n",
    "\n",
    "*back to [index](https://github.com/tsmatz/azureml-tutorial/)*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variable's Setting\n",
    "\n",
    "Replace below's branket's string and set the required variables.\n",
    "\n",
    "> Note : By the following ```az configure --defaults```, you can skip setting for ```--resource-group``` and ```--workspace-name``` options in each ```az ml``` command.<br>\n",
    "> ```az configure --defaults group=$resource_group workspace=$aml_workspace```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_resource_group = \"{AML-RESOURCE-GROUP-NAME}\"\n",
    "my_workspace = \"{AML-WORSPACE-NAME}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save your training script as file (train.py)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create ```scirpt``` directory and save Python script as ```./script/train.py```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "script_folder = './script'\n",
    "os.makedirs(script_folder, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing script/train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile script/train.py\n",
    "import os\n",
    "import argparse\n",
    "import tensorflow as tf\n",
    "\n",
    "# device test\n",
    "print(\"##### List of available GPU #####\")\n",
    "print(tf.config.list_physical_devices(\"GPU\"))\n",
    "\n",
    "# parse arguments\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\n",
    "    \"--data_folder\",\n",
    "    type=str,\n",
    "    default=\"./data/train\",\n",
    "    help=\"Folder path for input data\")\n",
    "parser.add_argument(\n",
    "    \"--model_folder\",\n",
    "    type=str,\n",
    "    default=\"./outputs\",  # AML experiments outputs folder\n",
    "    help=\"Folder path for model output\")\n",
    "parser.add_argument(\n",
    "    \"--learning_rate\",\n",
    "    type=float,\n",
    "    default=\"0.001\",\n",
    "    help=\"Learning Rate\")\n",
    "parser.add_argument(\n",
    "    \"--first_layer\",\n",
    "    type=int,\n",
    "    default=\"128\",\n",
    "    help=\"Neuron number for the first hidden layer\")\n",
    "parser.add_argument(\n",
    "    \"--second_layer\",\n",
    "    type=int,\n",
    "    default=\"64\",\n",
    "    help=\"Neuron number for the second hidden layer\")\n",
    "parser.add_argument(\n",
    "    \"--epochs_num\",\n",
    "    type=int,\n",
    "    default=\"6\",\n",
    "    help=\"Number of epochs\")\n",
    "FLAGS, unparsed = parser.parse_known_args()\n",
    "\n",
    "# build model\n",
    "model = tf.keras.models.Sequential([\n",
    "    tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "    tf.keras.layers.Dense(FLAGS.first_layer, activation=\"relu\"),\n",
    "    tf.keras.layers.Dense(FLAGS.second_layer, activation=\"relu\"),\n",
    "    tf.keras.layers.Dense(10)\n",
    "])\n",
    "model.compile(\n",
    "    optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate),\n",
    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],\n",
    ")\n",
    "\n",
    "# run training\n",
    "train_data = tf.data.experimental.load(FLAGS.data_folder)\n",
    "model.fit(\n",
    "    train_data.shuffle(1000).batch(128).prefetch(tf.data.AUTOTUNE),\n",
    "    epochs=FLAGS.epochs_num\n",
    ")\n",
    "\n",
    "# save model and variables\n",
    "model_path = os.path.join(FLAGS.model_folder, \"mnist_tf_model\")\n",
    "model.save(model_path)\n",
    "print(\"current working directory : \", os.getcwd())\n",
    "print(\"model folder : \", model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on remote VM\n",
    "\n",
    "Now let's start to integrate with AML and automate training on remote virtual machine."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1 : Create new remote virtual machine\n",
    "\n",
    "Create your new reomte virtual machine with GPU.<br>\n",
    "Before starting, **please check the following**.\n",
    "\n",
    "- Make sure that the following size (in the following script, ```Standard_NC4as_T4_v3```) is supported in the location (in which AML workspace resides).\n",
    "- You should have quota for ML GPU VM in your Azure subscription. If you don't have, please request quota in Azure Portal.\n",
    "\n",
    "**If you don't have any quota for GPU, please use CPU VM (such as, Standard_D2_v2).**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[K{\\ Finished ..\n",
      "  \"id\": \"/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/computes/myvm01\",\n",
      "  \"idle_time_before_scale_down\": 120,\n",
      "  \"location\": \"eastus\",\n",
      "  \"max_instances\": 1,\n",
      "  \"min_instances\": 0,\n",
      "  \"name\": \"myvm01\",\n",
      "  \"network_settings\": {},\n",
      "  \"provisioning_state\": \"Succeeded\",\n",
      "  \"resourceGroup\": \"rg-AML\",\n",
      "  \"size\": \"STANDARD_NC6\",\n",
      "  \"ssh_public_access_enabled\": true,\n",
      "  \"tier\": \"dedicated\",\n",
      "  \"type\": \"amlcompute\"\n",
      "}\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml compute create --name myvm01 \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace \\\n",
    "  --type amlcompute \\\n",
    "  --min-instances 0 \\\n",
    "  --max-instances 1 \\\n",
    "  --size Standard_NC4as_T4_v3  # change such as Standard_NC6 or Standard_D2_v2 if needed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By setting 0 in ```--min-instances```, the node will be terminated if it's inactive. (You can save money.)    \n",
    "\n",
    "> Note : By setting appropriate time duration in ```--idle-time-before-scale-down``` option, you can prevent scaling-down when the training has finished. (Otherwise, it will scale down in 120 seconds after the training has finished, and the next training will slow to start because of cluster resizing.)\n",
    "\n",
    "> Note : You can also attach an existing virtual machine (bring your own compute resource) as a compute target."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2 : Submit training job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Submit a training job.\n",
    "\n",
    "In this example, I use the registered data asset named ```mnist_data``` and mount this data in my compute target. (Run \"[Exercise02 : Prepare Data](./exercise02_prepare_data.ipynb)\" for data preparation.)<br>\n",
    "In order to use data asset in AML, set ```azureml:{DATA_NAME}:{DATA_VERSION}``` or ```azureml:{DATA_NAME}@latest``` for the latest version of assets as follows.\n",
    "\n",
    "> Note : Here I use AML built-in environment (```AzureML-tensorflow-2.7-ubuntu20.04-py38-cuda11-gpu```), but you can build and use your own environment.<br>\n",
    "> In the later example in this notebook, I'll run the same script with my own environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing 04_mnist_train_job.yml\n"
     ]
    }
   ],
   "source": [
    "%%writefile 04_mnist_train_job.yml\n",
    "$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json\n",
    "code: script\n",
    "command: >-\n",
    "  python train.py\n",
    "  --data_folder ${{inputs.mnist_tf}}/train\n",
    "inputs:\n",
    "  mnist_tf:\n",
    "    type: uri_folder\n",
    "    path: azureml:mnist_data@latest\n",
    "environment: azureml:AzureML-tensorflow-2.7-ubuntu20.04-py38-cuda11-gpu@latest\n",
    "compute: azureml:myvm01\n",
    "display_name: tf_remote_experiment\n",
    "experiment_name: tf_remote_experiment\n",
    "description: This is example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's submit a job with AML CLI.<br>\n",
    "See the progress and results in job view on [AML Studio](https://ml.azure.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mUploading script (0.0 MBs): 100%|███████| 2005/2005 [00:00<00:00, 232952.34it/s]\u001b[0m\n",
      "\u001b[39m\n",
      "\n",
      "{\n",
      "  \"code\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/codes/ccb97fff-ecb4-4097-8353-2677db8c6bf2/versions/1\",\n",
      "  \"command\": \"python train.py --data_folder ${{inputs.mnist_tf}}/train\",\n",
      "  \"compute\": \"azureml:myvm01\",\n",
      "  \"creation_context\": {\n",
      "    \"created_at\": \"2022-10-04T03:24:09.641642+00:00\",\n",
      "    \"created_by\": \"Tsuyoshi Matsuzaki\",\n",
      "    \"created_by_type\": \"User\"\n",
      "  },\n",
      "  \"description\": \"This is example\",\n",
      "  \"display_name\": \"tf_remote_experiment\",\n",
      "  \"environment\": \"azureml:AzureML-tensorflow-2.7-ubuntu20.04-py38-cuda11-gpu:23\",\n",
      "  \"environment_variables\": {},\n",
      "  \"experiment_name\": \"tf_remote_experiment\",\n",
      "  \"id\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/jobs/tidy_helmet_tn621n5rqw\",\n",
      "  \"inputs\": {\n",
      "    \"mnist_tf\": {\n",
      "      \"mode\": \"ro_mount\",\n",
      "      \"path\": \"azureml:mnist_data:1\",\n",
      "      \"type\": \"uri_folder\"\n",
      "    }\n",
      "  },\n",
      "  \"name\": \"tidy_helmet_tn621n5rqw\",\n",
      "  \"outputs\": {\n",
      "    \"default\": {\n",
      "      \"mode\": \"rw_mount\",\n",
      "      \"path\": \"azureml://datastores/workspaceartifactstore/ExperimentRun/dcid.tidy_helmet_tn621n5rqw\",\n",
      "      \"type\": \"uri_folder\"\n",
      "    }\n",
      "  },\n",
      "  \"parameters\": {},\n",
      "  \"properties\": {\n",
      "    \"ContentSnapshotId\": \"9cadeda6-0230-4d16-8b6e-c4a105d3ba61\",\n",
      "    \"_azureml.ComputeTargetType\": \"amlctrain\"\n",
      "  },\n",
      "  \"resourceGroup\": \"rg-AML\",\n",
      "  \"resources\": {\n",
      "    \"instance_count\": 1,\n",
      "    \"properties\": {},\n",
      "    \"shm_size\": \"2g\"\n",
      "  },\n",
      "  \"services\": {\n",
      "    \"Studio\": {\n",
      "      \"endpoint\": \"https://ml.azure.com/runs/tidy_helmet_tn621n5rqw?wsid=/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourcegroups/rg-AML/workspaces/ws01&tid=72f988bf-86f1-41af-91ab-2d7cd011db47\",\n",
      "      \"job_service_type\": \"Studio\"\n",
      "    },\n",
      "    \"Tracking\": {\n",
      "      \"endpoint\": \"azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01?\",\n",
      "      \"job_service_type\": \"Tracking\"\n",
      "    }\n",
      "  },\n",
      "  \"status\": \"Starting\",\n",
      "  \"tags\": {},\n",
      "  \"type\": \"command\"\n",
      "}\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml job create --file 04_mnist_train_job.yml \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please wait until the job is completed.\n",
    "\n",
    "You can show the progress and result with [AML studio UI](https://ml.azure.com/) (see \"Jobs\" pane) or with the following CLI command.<br>\n",
    "(**Replace ```tidy_helmet_tn621n5rqw``` with your generated job name**. For getting job name, see above output.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = \"tidy_helmet_tn621n5rqw\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\r\n",
      "  \"code\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/codes/ccb97fff-ecb4-4097-8353-2677db8c6bf2/versions/1\",\r\n",
      "  \"command\": \"python train.py --data_folder ${{inputs.mnist_tf}}/train\",\r\n",
      "  \"compute\": \"azureml:myvm01\",\r\n",
      "  \"creation_context\": {\r\n",
      "    \"created_at\": \"2022-10-04T03:24:09.641642+00:00\",\r\n",
      "    \"created_by\": \"Tsuyoshi Matsuzaki\",\r\n",
      "    \"created_by_type\": \"User\"\r\n",
      "  },\r\n",
      "  \"description\": \"This is example\",\r\n",
      "  \"display_name\": \"tf_remote_experiment\",\r\n",
      "  \"environment\": \"azureml:AzureML-tensorflow-2.7-ubuntu20.04-py38-cuda11-gpu:23\",\r\n",
      "  \"environment_variables\": {},\r\n",
      "  \"experiment_name\": \"tf_remote_experiment\",\r\n",
      "  \"id\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/jobs/tidy_helmet_tn621n5rqw\",\r\n",
      "  \"inputs\": {\r\n",
      "    \"mnist_tf\": {\r\n",
      "      \"mode\": \"ro_mount\",\r\n",
      "      \"path\": \"azureml:mnist_data:1\",\r\n",
      "      \"type\": \"uri_folder\"\r\n",
      "    }\r\n",
      "  },\r\n",
      "  \"name\": \"tidy_helmet_tn621n5rqw\",\r\n",
      "  \"outputs\": {\r\n",
      "    \"default\": {\r\n",
      "      \"mode\": \"rw_mount\",\r\n",
      "      \"path\": \"azureml://datastores/workspaceartifactstore/ExperimentRun/dcid.tidy_helmet_tn621n5rqw\",\r\n",
      "      \"type\": \"uri_folder\"\r\n",
      "    }\r\n",
      "  },\r\n",
      "  \"parameters\": {},\r\n",
      "  \"properties\": {\r\n",
      "    \"ContentSnapshotId\": \"9cadeda6-0230-4d16-8b6e-c4a105d3ba61\",\r\n",
      "    \"ProcessInfoFile\": \"azureml-logs/process_info.json\",\r\n",
      "    \"ProcessStatusFile\": \"azureml-logs/process_status.json\",\r\n",
      "    \"_azureml.ComputeTargetType\": \"amlctrain\"\r\n",
      "  },\r\n",
      "  \"resourceGroup\": \"rg-AML\",\r\n",
      "  \"resources\": {\r\n",
      "    \"instance_count\": 1,\r\n",
      "    \"properties\": {},\r\n",
      "    \"shm_size\": \"2g\"\r\n",
      "  },\r\n",
      "  \"services\": {\r\n",
      "    \"Studio\": {\r\n",
      "      \"endpoint\": \"https://ml.azure.com/runs/tidy_helmet_tn621n5rqw?wsid=/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourcegroups/rg-AML/workspaces/ws01&tid=72f988bf-86f1-41af-91ab-2d7cd011db47\",\r\n",
      "      \"job_service_type\": \"Studio\"\r\n",
      "    },\r\n",
      "    \"Tracking\": {\r\n",
      "      \"endpoint\": \"azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01?\",\r\n",
      "      \"job_service_type\": \"Tracking\"\r\n",
      "    }\r\n",
      "  },\r\n",
      "  \"status\": \"Completed\",\r\n",
      "  \"tags\": {},\r\n",
      "  \"type\": \"command\"\r\n",
      "}\r\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml job show --name $job_name \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3 : Download results and evaluate\n",
    "\n",
    "Now let's check the generated model in local computer.\n",
    "\n",
    "Go to [Azure ML studio UI](https://ml.azure.com/).<br>\n",
    "You can then see the saved model in outputs directory.\n",
    "\n",
    "![Saved Outputs](https://tsmatz.github.io/images/github/azure-ml-tensorflow-complete-sample/20220225_Experiment_Outputs.jpg)\n",
    "\n",
    "By running the following ```az ml job download``` command, the logs and outputs are downloaded in local computer.<br>\n",
    "The logs are saved in ```artifacts/logs``` and outputs are in ```artifacts/outputs```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading artifact azureml://datastores/workspaceartifactstore/ExperimentRun/dcid.tidy_helmet_tn621n5rqw to /home/tsmatsuz/cli_yaml/artifacts\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml job download --name $job_name \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now check the downloaded result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-10-04 03:49:14.331348: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-10-04 03:49:14.513554: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2022-10-04 03:49:14.513669: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "2022-10-04 03:49:14.555072: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2022-10-04 03:49:15.588518: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory\n",
      "2022-10-04 03:49:15.588669: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory\n",
      "2022-10-04 03:49:15.588688: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "2022-10-04 03:49:16.552451: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory\n",
      "2022-10-04 03:49:16.552507: W tensorflow/stream_executor/cuda/cuda_driver.cc:263] failed call to cuInit: UNKNOWN ERROR (303)\n",
      "2022-10-04 03:49:16.552553: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (client1004): /proc/driver/nvidia/version does not exist\n",
      "2022-10-04 03:49:16.552855: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted 7, True 7\n",
      "Predicted 2, True 2\n",
      "Predicted 1, True 1\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "test_data = tf.data.Dataset.load(\"./data/test\")\n",
    "\n",
    "loaded_model = tf.keras.models.load_model(\"./artifacts/outputs/mnist_tf_model\")\n",
    "for image, true_value in test_data.take(3):\n",
    "    pred_output = loaded_model(tf.expand_dims(image, axis=0))\n",
    "    pred_value = tf.math.argmax(pred_output, axis=-1).numpy().item()\n",
    "    print(\"Predicted {}, True {}\".format(pred_value, true_value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4 : Register Model\n",
    "\n",
    "Now upload (register) the downloaded model into AML model management.\n",
    "\n",
    "> Note : You can directly register model file from run history as follows.<br>\n",
    "> ```az ml model create --name mnist_model_test --version 1 --path runs:/{RUN ID}/outputs/mnist_tf_model/ ...```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mUploading mnist_tf_model (1.42 MBs): 100%|█| 1418402/1418402 [00:00<00:00, 13276\u001b[0m\n",
      "\u001b[39m\n",
      "\n",
      "{\n",
      "  \"creation_context\": {\n",
      "    \"created_at\": \"2022-10-04T03:50:02.387431+00:00\",\n",
      "    \"created_by\": \"Tsuyoshi Matsuzaki\",\n",
      "    \"created_by_type\": \"User\",\n",
      "    \"last_modified_at\": \"2022-10-04T03:50:02.387431+00:00\",\n",
      "    \"last_modified_by\": \"Tsuyoshi Matsuzaki\",\n",
      "    \"last_modified_by_type\": \"User\"\n",
      "  },\n",
      "  \"id\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/models/mnist_model_test/versions/1\",\n",
      "  \"name\": \"mnist_model_test\",\n",
      "  \"path\": \"azureml://subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/workspaces/ws01/datastores/workspaceblobstore/paths/LocalUpload/9daf16eed2faa4a869ce2130d35121ba/mnist_tf_model\",\n",
      "  \"properties\": {},\n",
      "  \"resourceGroup\": \"rg-AML\",\n",
      "  \"tags\": {},\n",
      "  \"type\": \"custom_model\",\n",
      "  \"version\": \"1\"\n",
      "}\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml model create --name mnist_model_test \\\n",
    "  --version 1 \\\n",
    "  --path ./artifacts/outputs/mnist_tf_model \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Optional] Step 5 : Train with your own environment\n",
    "\n",
    "**This is not mandatory. (You can skip this section.)**\n",
    "\n",
    "You can also build your own environment with custom docker image.<br>\n",
    "Here we create a new docker environments for running scripts, and run the same training with this environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First I create conda dependancies yaml and save as ```04_conda_pydata.yml```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing 04_conda_pydata.yml\n"
     ]
    }
   ],
   "source": [
    "%%writefile 04_conda_pydata.yml\n",
    "name: project_environment\n",
    "dependencies:\n",
    "- python=3.8\n",
    "- pip:\n",
    "  - tensorflow-gpu==2.10.0\n",
    "channels:\n",
    "- anaconda\n",
    "- conda-forge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Register custom environment (named ```test-remote-gpu-env```) in AML with previous conda configuration.\n",
    "\n",
    "> Note : For CPU image, you can use ```mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing 04_env_register.yml\n"
     ]
    }
   ],
   "source": [
    "%%writefile 04_env_register.yml\n",
    "$schema: https://azuremlschemas.azureedge.net/latest/environment.schema.json\n",
    "name: test-remote-gpu-env\n",
    "image: mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.2-cudnn8-ubuntu20.04\n",
    "conda_file: 04_conda_pydata.yml\n",
    "description: This is example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\r\n",
      "  \"conda_file\": {\r\n",
      "    \"channels\": [\r\n",
      "      \"anaconda\",\r\n",
      "      \"conda-forge\"\r\n",
      "    ],\r\n",
      "    \"dependencies\": [\r\n",
      "      \"python=3.8\",\r\n",
      "      {\r\n",
      "        \"pip\": [\r\n",
      "          \"tensorflow-gpu==2.10.0\"\r\n",
      "        ]\r\n",
      "      }\r\n",
      "    ],\r\n",
      "    \"name\": \"project_environment\"\r\n",
      "  },\r\n",
      "  \"creation_context\": {\r\n",
      "    \"created_at\": \"2022-10-04T03:50:38.374638+00:00\",\r\n",
      "    \"created_by\": \"Tsuyoshi Matsuzaki\",\r\n",
      "    \"created_by_type\": \"User\",\r\n",
      "    \"last_modified_at\": \"2022-10-04T03:50:38.374638+00:00\",\r\n",
      "    \"last_modified_by\": \"Tsuyoshi Matsuzaki\",\r\n",
      "    \"last_modified_by_type\": \"User\"\r\n",
      "  },\r\n",
      "  \"description\": \"This is example\",\r\n",
      "  \"id\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/environments/test-remote-gpu-env/versions/1\",\r\n",
      "  \"image\": \"mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.2-cudnn8-ubuntu20.04\",\r\n",
      "  \"name\": \"test-remote-gpu-env\",\r\n",
      "  \"os_type\": \"linux\",\r\n",
      "  \"resourceGroup\": \"rg-AML\",\r\n",
      "  \"tags\": {},\r\n",
      "  \"version\": \"1\"\r\n",
      "}\r\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml environment create --file 04_env_register.yml \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Go to [AML Studio UI](https://ml.azure.com/) and click \"Environments\". Next, click \"Custom environments\" tab and select the above environment.<br>\n",
    "Please wait until the environment image build status is succeeded.\n",
    "\n",
    "![Environment status](https://tsmatz.github.io/images/github/azure-ml-tensorflow-complete-sample/20221220_Environment_Status.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train script with above custom environment.\n",
    "\n",
    "It will take a long time (over 30 minutes) for the first time run, because it'll pull base image, generate new image (custom environment), start nodes in cluster, and run scripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing 04_mnist_train_job_with_custom_image.yml\n"
     ]
    }
   ],
   "source": [
    "%%writefile 04_mnist_train_job_with_custom_image.yml\n",
    "$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json\n",
    "code: script\n",
    "command: >-\n",
    "  python train.py\n",
    "  --data_folder ${{inputs.mnist_tf}}/train\n",
    "inputs:\n",
    "  mnist_tf:\n",
    "    type: uri_folder\n",
    "    path: azureml:mnist_data@latest\n",
    "environment: azureml:test-remote-gpu-env@latest\n",
    "compute: azureml:myvm01\n",
    "display_name: tf_remote_experiment\n",
    "experiment_name: tf_remote_experiment\n",
    "description: This is example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"code\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/codes/ccb97fff-ecb4-4097-8353-2677db8c6bf2/versions/1\",\n",
      "  \"command\": \"python train.py --data_folder ${{inputs.mnist_tf}}/train\",\n",
      "  \"compute\": \"azureml:myvm01\",\n",
      "  \"creation_context\": {\n",
      "    \"created_at\": \"2022-10-04T03:51:19.432458+00:00\",\n",
      "    \"created_by\": \"Tsuyoshi Matsuzaki\",\n",
      "    \"created_by_type\": \"User\"\n",
      "  },\n",
      "  \"description\": \"This is example\",\n",
      "  \"display_name\": \"tf_remote_experiment\",\n",
      "  \"environment\": \"azureml:test-remote-gpu-env:1\",\n",
      "  \"environment_variables\": {},\n",
      "  \"experiment_name\": \"tf_remote_experiment\",\n",
      "  \"id\": \"azureml:/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01/jobs/boring_answer_jdm40htxhm\",\n",
      "  \"inputs\": {\n",
      "    \"mnist_tf\": {\n",
      "      \"mode\": \"ro_mount\",\n",
      "      \"path\": \"azureml:mnist_data:1\",\n",
      "      \"type\": \"uri_folder\"\n",
      "    }\n",
      "  },\n",
      "  \"name\": \"boring_answer_jdm40htxhm\",\n",
      "  \"outputs\": {\n",
      "    \"default\": {\n",
      "      \"mode\": \"rw_mount\",\n",
      "      \"path\": \"azureml://datastores/workspaceartifactstore/ExperimentRun/dcid.boring_answer_jdm40htxhm\",\n",
      "      \"type\": \"uri_folder\"\n",
      "    }\n",
      "  },\n",
      "  \"parameters\": {},\n",
      "  \"properties\": {\n",
      "    \"ContentSnapshotId\": \"9cadeda6-0230-4d16-8b6e-c4a105d3ba61\",\n",
      "    \"_azureml.ComputeTargetType\": \"amlctrain\"\n",
      "  },\n",
      "  \"resourceGroup\": \"rg-AML\",\n",
      "  \"resources\": {\n",
      "    \"instance_count\": 1,\n",
      "    \"properties\": {},\n",
      "    \"shm_size\": \"2g\"\n",
      "  },\n",
      "  \"services\": {\n",
      "    \"Studio\": {\n",
      "      \"endpoint\": \"https://ml.azure.com/runs/boring_answer_jdm40htxhm?wsid=/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourcegroups/rg-AML/workspaces/ws01&tid=72f988bf-86f1-41af-91ab-2d7cd011db47\",\n",
      "      \"job_service_type\": \"Studio\"\n",
      "    },\n",
      "    \"Tracking\": {\n",
      "      \"endpoint\": \"azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/b3ae1c15-4fef-4362-8c3a-5d804cdeb18d/resourceGroups/rg-AML/providers/Microsoft.MachineLearningServices/workspaces/ws01?\",\n",
      "      \"job_service_type\": \"Tracking\"\n",
      "    }\n",
      "  },\n",
      "  \"status\": \"Starting\",\n",
      "  \"tags\": {},\n",
      "  \"type\": \"command\"\n",
      "}\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml job create --file 04_mnist_train_job_with_custom_image.yml \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6 : Remove AML compute\n",
    "\n",
    "**You don't need to remove your AML compute** for saving money, because the nodes will be automatically terminated, when it's inactive.<br>\n",
    "But if you want to clean up, please run as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleting compute myvm01 \n",
      ".......Done.\n",
      "(0m 36s)\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!az ml compute delete --name myvm01 \\\n",
    "  --resource-group $my_resource_group \\\n",
    "  --workspace-name $my_workspace \\\n",
    "  --yes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}