{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa3c47de",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install -U keras==2.2.4\n",
    "!pip install matplotlib\n",
    "!pip install dlr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7beb2b5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from sagemaker.tensorflow import TensorFlow\n",
    "from tensorflow.keras.datasets import mnist\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import keras\n",
    "import sagemaker, os\n",
    "from time import sleep\n",
    "from uuid import uuid4\n",
    "import dlr\n",
    "print(keras.__version__)\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed8b4f89",
   "metadata": {},
   "outputs": [],
   "source": [
    "work_dir = './work'\n",
    "!mkdir -p {work_dir}\n",
    "\n",
    "TRAIN_X_PATH = os.path.join(work_dir,'train_X.npy')\n",
    "TEST_X_PATH = os.path.join(work_dir,'test_X.npy')\n",
    "TRAIN_Y_PATH = os.path.join(work_dir,'train_y.npy')\n",
    "TEST_Y_PATH = os.path.join(work_dir,'test_y.npy')\n",
    "(train_X, train_y), (test_X, test_y) = mnist.load_data()\n",
    "train_X = (train_X-127.5)/127.5\n",
    "test_X = (test_X-127.5)/127.5\n",
    "# channel last\n",
    "train_X = train_X.reshape((train_X.shape[0],train_X.shape[1],train_X.shape[2],1))\n",
    "test_X = test_X.reshape((test_X.shape[0],test_X.shape[1],test_X.shape[2],1))\n",
    "# one-hot\n",
    "train_y = np.eye(10)[train_y]\n",
    "test_y = np.eye(10)[test_y]\n",
    "np.save(TRAIN_X_PATH,train_X)\n",
    "np.save(TEST_X_PATH,test_X)\n",
    "np.save(TRAIN_Y_PATH,train_y)\n",
    "np.save(TEST_Y_PATH,test_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ece5c71",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check shapes\n",
    "print(train_X.shape,train_y.shape,test_X.shape,test_y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7ded7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "role = sagemaker.get_execution_role()\n",
    "sess = sagemaker.session.Session()\n",
    "train_X_uri = sess.upload_data(path=TRAIN_X_PATH, bucket=sess.default_bucket(), key_prefix='sagemaker/mnist')\n",
    "train_y_uri = sess.upload_data(path=TRAIN_Y_PATH, bucket=sess.default_bucket(), key_prefix='sagemaker/mnist')\n",
    "test_X_uri = sess.upload_data(path=TEST_X_PATH, bucket=sess.default_bucket(), key_prefix='sagemaker/mnist')\n",
    "test_y_uri = sess.upload_data(path=TEST_Y_PATH, bucket=sess.default_bucket(), key_prefix='sagemaker/mnist')\n",
    "print(train_X_uri)\n",
    "print(train_y_uri)\n",
    "print(test_X_uri)\n",
    "print(test_y_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64e7334d",
   "metadata": {},
   "source": [
    "## Classifier Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8e1801",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = TensorFlow(\n",
    "    entry_point='./src/classifier_train.py',\n",
    "    role=role,\n",
    "    instance_count=1,\n",
    "    instance_type='ml.g4dn.xlarge',\n",
    "    framework_version='1.13', # SageMaker Neo を使うために keras 2.2.4 に合わせる\n",
    "    py_version='py3',\n",
    "    hyperparameters={\n",
    "        'epochs':30,\n",
    "        'increment':'False'\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10d32375",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "print(train_X_uri[:-11]) # dir以下全て\n",
    "estimator.fit({\n",
    "    'train': train_X_uri[:-11],\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cec5729b",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier_model_uri = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts']\n",
    "print(classifier_model_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba576c32",
   "metadata": {},
   "source": [
    "## Neo Compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1c9e0c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "sm_client = boto3.client('sagemaker')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "149f7fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier_output_s3_location = f's3://{sagemaker.session.Session().default_bucket()}/sagemaker/mnist/classifier/compilied_model'\n",
    "print(classifier_output_s3_location)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2088d048",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier_compile_jobname = f'classifier-{str(uuid4())}'\n",
    "print(classifier_compile_jobname)\n",
    "response = sm_client.create_compilation_job(\n",
    "    CompilationJobName=classifier_compile_jobname,\n",
    "    RoleArn=role,\n",
    "    InputConfig={\n",
    "        'S3Uri': classifier_model_uri,\n",
    "        'DataInputConfig': '{\"input_1\":[1,1,28,28]}',\n",
    "        'Framework': 'KERAS',\n",
    "    },\n",
    "    OutputConfig={\n",
    "        'S3OutputLocation': classifier_output_s3_location,\n",
    "        'TargetDevice': 'ml_c5',\n",
    "    },\n",
    "    StoppingCondition={\n",
    "        'MaxRuntimeInSeconds': 900,\n",
    "        'MaxWaitTimeInSeconds': 900\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfcb46c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "while True:\n",
    "    response = sm_client.describe_compilation_job(CompilationJobName=classifier_compile_jobname)\n",
    "    status = response['CompilationJobStatus']\n",
    "    if status in ['COMPLETED','FAILED','STOPPED'] : \n",
    "        print('!')\n",
    "        print(status)\n",
    "        classifier_neo_model_uri = response['ModelArtifacts']['S3ModelArtifacts']\n",
    "        break\n",
    "    else:\n",
    "        print('.',end='')\n",
    "    sleep(5)\n",
    "print(classifier_neo_model_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e0cc87",
   "metadata": {},
   "source": [
    "### ファイルの配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5c92f5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "publisher_dir = './src/ggv2/components/artifacts/com.example.Publisher/1.0.0'\n",
    "subscriber_dir = './src/ggv2/components/artifacts/com.example.Subscriber/1.0.0'\n",
    "iotpublisher_classifier_dir = './src/ggv2/components/artifacts/com.example.IoTPublisher/1.0.1/classifier'\n",
    "\n",
    "# test_x の配置\n",
    "!mkdir -p {publisher_dir}\n",
    "!mv {work_dir}/test_X.npy {publisher_dir}/test_X.npy\n",
    "\n",
    "# # classifier1.h5 のダウンロード・解凍・gz削除\n",
    "!aws s3 cp {classifier_model_uri} {work_dir}/\n",
    "!mkdir -p {subscriber_dir}\n",
    "!tar zxvf {work_dir}/model.tar.gz -C {subscriber_dir}\n",
    "!rm {work_dir}/model.tar.gz \n",
    "\n",
    "# Neo コンパイル済モデルのダウンロード・解凍・gz削除\n",
    "\n",
    "!aws s3 cp {classifier_neo_model_uri} {work_dir}/\n",
    "!mkdir -p {iotpublisher_classifier_dir}\n",
    "!tar zxvf {work_dir}/model-ml_c5.tar.gz -C {iotpublisher_classifier_dir}\n",
    "!rm {work_dir}/model-ml_c5.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9941581",
   "metadata": {},
   "source": [
    "## 動作確認\n",
    "### keras model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95afe394",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "classifier = keras.models.load_model(f'{subscriber_dir}/classifier.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f55ee3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "classifier.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b4ec53",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_y = classifier.predict(test_X[0:1,:,:,:])\n",
    "plt.imshow(test_X[0,:,:,0],'gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "326a8d73",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.argmax(pred_y[0,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45260695",
   "metadata": {},
   "source": [
    "## Neo Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8521a19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dlr\n",
    "import numpy as np\n",
    "\n",
    "classifier_neo = dlr.DLRModel(iotpublisher_classifier_dir, 'cpu', 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d485c27e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_y = classifier_neo.run(test_X[0,:,:,:].reshape(1,1,28,28))[0]\n",
    "np.argmax(pred_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59f37874",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.c5.large",
  "kernelspec": {
   "display_name": "conda_amazonei_tensorflow_p36",
   "language": "python",
   "name": "conda_amazonei_tensorflow_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}