Optimisation_TensorRT_FP16.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config IPCompleter.greedy=True\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "from tensorflow import keras\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vérification des versions installées"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tf.version.VERSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!dpkg -l | grep nvinfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Création du dataset d'entrées"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On crée les dataset au format suivant :"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://github.com/AlexandreBourrieau/JetsonNano/blob/main/images/datset_regression_heatmap2.png?raw=true\" width=400>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "repertoire_courant = os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TACHE = \"visage\"\n",
    "CATEGORIES = ['nez']\n",
    "batch_size = 1\n",
    "\n",
    "dataset_base = tf.keras.preprocessing.image_dataset_from_directory(\n",
    "    repertoire_courant+\"/projet_regression/\"+TACHE,\n",
    "    validation_split=0.0,\n",
    "    batch_size=batch_size,\n",
    "    label_mode=\"categorical\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "repertoires_images = dataset_base.file_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from imgaug.augmentables.heatmaps import HeatmapsOnImage\n",
    "\n",
    "def CreateHeatmap(x, y, x0, y0, width, height):\n",
    "    alpha = 0.4\n",
    "    sigma_x = alpha*width/6.0\n",
    "    sigma_y = height/6.0\n",
    "    return np.exp(-((x - x0)**2. / (2. * sigma_x**2.) + (y - y0)**2. / (2. * sigma_y**2.)))\n",
    "\n",
    "def CreationDatasetRegression(liste_fichiers,noms_classes,width,height):\n",
    "    heatmap_ = [[] for i in range(len(noms_classes))]\n",
    "    images_= []\n",
    "\n",
    "    width = width\n",
    "    height = height\n",
    "\n",
    "    # Création du heatmap nul\n",
    "    heatmap_zero = np.zeros((height,width)).astype(np.float32)\n",
    "    heatmap_zero = HeatmapsOnImage(heatmap_zero, shape=(height,width), min_value=0.0, max_value=1.0)\n",
    "    heatmap_zero = heatmap_zero.avg_pool(4)\n",
    "\n",
    "    # Création de la grille pour le heatmap\n",
    "    x,y = np.meshgrid(np.linspace(0,width-1,width), np.linspace(0,height-1,height))\n",
    "    \n",
    "    # Définition des constantes pour le heatmap\n",
    "    alpha = 0.4\n",
    "    sigma_x = alpha*width/6.0\n",
    "    sigma_y = alpha*height/6.0\n",
    "    \n",
    "    # Conversion des classe du heatmap zero en tenseur 1D\n",
    "    # (56,56,nbr_classes)\n",
    "    heatmap_zero = tf.convert_to_tensor(heatmap_zero.get_arr())\n",
    "\n",
    "    for fichier in liste_fichiers:\n",
    "        # Chargement de l'image\n",
    "        image = tf.keras.preprocessing.image.load_img(fichier)\n",
    "        image = tf.keras.preprocessing.image.img_to_array(image)\n",
    "\n",
    "        # Extraction des coordonnées (x0,Y0)\n",
    "        # correspondantes aux dimensions de l'image chargée (height,width)\n",
    "        element = tf.strings.split(fichier,sep=\"_image_\")\n",
    "        element = tf.strings.split(element[1],sep=\"_\")\n",
    "        x0 = tf.strings.to_number(element[1],out_type=tf.dtypes.float32)\n",
    "        y0 = tf.strings.split(element[2],sep=\".\")\n",
    "        y0 = tf.strings.to_number(y0[0],out_type=tf.dtypes.float32)\n",
    "        \n",
    "        # Extraction de la classe\n",
    "        classe = tf.strings.split(fichier,sep=\"/\")[-2]\n",
    "        \n",
    "        # Extraction du label\n",
    "        label = tf.cast(tf.strings.regex_full_match(noms_classes, classe),dtype=\"int32\")\n",
    "        \n",
    "        # Extraction de la valeur binaire du label\n",
    "        maxi = tf.math.argmax(label,output_type=tf.dtypes.int32)\n",
    "        \n",
    "        # Création du heatmap de l'image\n",
    "        heatmap = CreateHeatmap(x,y,x0,y0,width,height).astype(np.float32)\n",
    "        heatmap = HeatmapsOnImage(heatmap, (height,width,1), min_value=0.0, max_value=1.0)\n",
    "        heatmap = heatmap.max_pool(4)\n",
    "        \n",
    "        # Sauvegarde du heatmap et de l'image dans les listes\n",
    "        for i in range(len(noms_classes)):\n",
    "            if i == maxi:\n",
    "                heatmap = tf.convert_to_tensor(heatmap.get_arr())\n",
    "                heatmap_[i].append(heatmap)\n",
    "            else:\n",
    "                heatmap_[i].append(heatmap_zero)\n",
    "        images_.append(image)\n",
    "   \n",
    "    # Création du dataset\n",
    "    images_ = tf.convert_to_tensor(images_)                                # (nbr_images,H,W,3)\n",
    "    heatmap_ = tf.convert_to_tensor(heatmap_)                              # (nbr_images,H/4,W/4,1)\n",
    "    \n",
    "    heatmap_ = tf.transpose(heatmap_,perm=[1,2,3,0,4])                     # (nbr_image,H/4,W/4)\n",
    "    heatmap_ = tf.squeeze(heatmap_,-1)                                     # (nbr_image,H/4,W/4)\n",
    "\n",
    "    datasetHeatmap = tf.data.Dataset.from_tensors(heatmap_)                # (nbr_images,3,H/4,W/4)\n",
    "    datasetImg = tf.data.Dataset.from_tensors(images_)                     # (nbr_images,H,W,3) \n",
    "    dataset = tf.data.Dataset.zip((datasetImg,datasetHeatmap))\n",
    "\n",
    "    return (dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_regression = CreationDatasetRegression(repertoires_images,dataset_base.class_names,width=224, height=224)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for image,heatmap in dataset_regression.take(1):\n",
    "    print(image.shape)\n",
    "    print(heatmap.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modifications du réseau"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorRT réalise plusieurs modifications et optimisations sur le réseau. Tout d'abord, les couches n'utilisant pas de sorties sont supprimées afin d'alléger les calculs. Ensuite, lorsque cela est possible, les couches de convolution, les offsets (bias) et les activations sont fusionnées pour ne former qu'une couche unique.  \n",
    "La figure ci-dessous montre un exemple de **réseau de convolution non optimisé** :"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://github.com/AlexandreBourrieau/JetsonNano/blob/main/images/Reseau_NonOptimise2.png?raw=true\" width=600>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "La figure ci-dessous montre le résultat obtenu après optimisation par **fusion des couches de manière verticale** (les couches fusionnées sont nommées **CBR**). Ceci permet d'améliorer les temps de calculs sur les GPUs car les opérations peuvent être faites sur des blocks parallélisés :"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://github.com/AlexandreBourrieau/JetsonNano/blob/main/images/FusionVerticale_All.png?raw=true\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Un autre manière d'optimiser le réseau est de réaliser une **fusion horizontale des couches** qui prennent des entrées identiques :"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://github.com/AlexandreBourrieau/JetsonNano/blob/main/images/FusionHorizontale_All.png?raw=true\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conversion du modèle vers un modèle TensorRT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pour réaliser la converison du modèle, nous utilisons l'outil trtexec, en précisant le répertoire dans lequel le modèle est sauvegardé et les paramètres de conversion :  \n",
    "- **precision_mode** : Format de codage des nombres : FP32, FP16 ou INT8. Les formats inférieurs au FP32 (FP16 et INT8) peuvent améliorer les performances des calculs. Le mode FP16 utilise des coeurs matériels avec des instructions sur des flottants 16bits lorsque cela est possible. Le mode INT8 utilise des coeurs matériels avec des instructions sur des entiers.     \n",
    "- **max_batch_size** : Le batch-size maximum à utiliser pendant l'optimisation. Pendant l'excéution en temps réel, on peut choisir une valeur plus petite mais pas plus grande.  \n",
    "- **minimum_segment_size** : Ce paramètre permet de préciser la valeur minimale de noeuds qu'il faut pour que la conversion du réseau soit exécutée. En conséquence, en général on choisit des valeurs inférieures à 5. Ce paramètre permet également de choisir le nombre minimum de noeuds pendant l'optimisation finale INT8 et donc d'optimiser la précision des résultats finaux.  \n",
    "- **max_workspace_size_byte** : Les opérations d'optimisation de TensorRT ont besoin d'utiliser de l'espace de stockage temporaire. Ce paramètre permet de limiter l'espace maximal utilisé dans le GPU qu'une couche peut utiliser. Si une valeur insuffisante est donnée, TensorRT peut ne pas réussir à optimiser le modèle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!/usr/src/tensorrt/bin/trtexec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversion au format FP16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!/usr/src/tensorrt/bin/trtexec --verbose --fp16 --workspace=256 --maxBatch=1 --onnx=\"/home/alexandre/Regression_Resnet18_saved_model/Regression_Resnet18_saved_model.onnx\" --saveEngine=\"/home/alexandre/Regression_Resnet18_saved_model/Regression_Resnet18_FP16.engine\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utilisation du modèle TensorRT optimisé"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chargement du moteur (engine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorrt as trt\n",
    "\n",
    "# Construction de la class du logger\n",
    "class MyLogger(trt.ILogger):\n",
    "    def __init__(self):\n",
    "        trt.ILogger.__init__(self)\n",
    "\n",
    "    def log(self, severity, msg):\n",
    "        print(\"%s : %s\" %(severity,msg))\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "\n",
    "PRECISION = trt.float32\n",
    "\n",
    "logger = MyLogger()\n",
    "runtime = trt.Runtime(logger)\n",
    "\n",
    "with open(\"Regression_Resnet18_saved_model/Regression_Resnet18_FP16.engine\", \"rb\") as f:\n",
    "    engine = runtime.deserialize_cuda_engine(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### Création du contexte"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Création du contexte nécessaire pour lancer le modèle :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = engine.create_execution_context()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Allocation de l'espace mémoire"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Récupère la mémoire disponible du GPU :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda.mem_get_info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://github.com/AlexandreBourrieau/JetsonNano/blob/main/images/MemoireGPU.png?raw=true\" width=600>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for binding in engine:\n",
    "    print(engine.get_binding_shape(binding))\n",
    "    print(trt.nptype(engine.get_binding_dtype(binding)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allocation des buffers d'entrée / sortie dans la mémoire GPU\n",
    "size_input = trt.volume(engine.get_binding_shape(0))* engine.max_batch_size\n",
    "size_output = trt.volume(engine.get_binding_shape(1))* engine.max_batch_size\n",
    "\n",
    "# Allocation de mémoire de type \"page-locked\" sur l'hôte\n",
    "input_host_mem = cuda.pagelocked_empty(size_input, trt.nptype(PRECISION))\n",
    "output_host_mem = cuda.pagelocked_empty(size_output, trt.nptype(PRECISION))\n",
    "\n",
    "# Allocation de mémoire dans la mémoire GPU\n",
    "input_device_mem = cuda.mem_alloc(input_host_mem.nbytes)\n",
    "output_device_mem = cuda.mem_alloc(output_host_mem.nbytes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Récupère les adresses en GPU des buffers entrées / sorties\n",
    "bindings = [int(input_device_mem), int(output_device_mem)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exécution d'une prédiction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Récupère l'image depuis le dataset :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for image,heatmaps in dataset_regression.unbatch().take(1):\n",
    "    image = np.asarray(image).astype(trt.nptype(PRECISION))       # (224,224,3)\n",
    "    image = np.expand_dims(image,axis=0)                          # (1,224,224,3)\n",
    "    np.copyto(input_host_mem,image.ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transfert les données de l'image vers la mémoire GPU (transfert Host => Device) :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda.memcpy_htod(input_device_mem, input_host_mem)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exécution du modèle :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.execute(batch_size=1,bindings=bindings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda.memcpy_dtoh(output_host_mem, output_device_mem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.reshape(input_host_mem,(224,224,3)).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.reshape(output_host_mem,(56,56))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.reshape(output_host_mem,(56,56)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Libère la mémoire allouée :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_device_mem.free()\n",
    "output_device_mem.free()\n",
    "del input_host_mem\n",
    "del output_host_mem\n",
    "del context\n",
    "del engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyse des performances du modèle optimisé"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyse de la vitesse de traitement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def MesureDesPredictions_TRT(dataset,nbr_run,PRECISION):\n",
    "    delais = []\n",
    "    predictions = []\n",
    "    \n",
    "    # Allocation des buffers d'entrée / sortie dans la mémoire GPU\n",
    "    size_input = trt.volume(engine.get_binding_shape(0))* engine.max_batch_size\n",
    "    size_output = trt.volume(engine.get_binding_shape(1))* engine.max_batch_size\n",
    "\n",
    "    # Allocation de mémoire de type \"page-locked\" sur l'hôte\n",
    "    input_host_mem = cuda.pagelocked_empty(size_input, trt.nptype(PRECISION))\n",
    "    output_host_mem = cuda.pagelocked_empty(size_output, trt.nptype(PRECISION))\n",
    "\n",
    "    # Allocation de mémoire dans la mémoire GPU\n",
    "    input_device_mem = cuda.mem_alloc(input_host_mem.nbytes)\n",
    "    output_device_mem = cuda.mem_alloc(output_host_mem.nbytes)    \n",
    "\n",
    "    # Récupère les adresses en GPU des buffers entrées / sorties\n",
    "    bindings = [int(input_device_mem), int(output_device_mem)]\n",
    "    \n",
    "    # Transfert de l'mimage en mémoire sur le host\n",
    "    for image,heatmaps in dataset.unbatch().take(1):\n",
    "        image = np.asarray(image).astype(trt.nptype(PRECISION))\n",
    "        image = np.expand_dims(image,axis=0)\n",
    "        np.copyto(input_host_mem,image.ravel())\n",
    "    \n",
    "    # Initialisation des calculs\n",
    "    print(\"Initialisation des calculs...\")\n",
    "    for i in range(5):\n",
    "        cuda.memcpy_htod(input_device_mem, input_host_mem)\n",
    "        context.execute(batch_size=1,bindings=bindings)        \n",
    "        cuda.memcpy_dtoh(output_host_mem, output_device_mem)    \n",
    "    \n",
    "    # Lance les inférences\n",
    "    for i in range(nbr_run):\n",
    "        time0 = time.time()\n",
    "        cuda.memcpy_htod(input_device_mem, input_host_mem)\n",
    "        context.execute(batch_size=1,bindings=bindings)        \n",
    "        cuda.memcpy_dtoh(output_host_mem, output_device_mem)    \n",
    "        time_end = time.time()\n",
    "\n",
    "        delais = np.append(delais,time_end - time0)\n",
    "\n",
    "        predictions.append(output_host_mem)\n",
    "        \n",
    "        if i%10 == 0:\n",
    "            print(\"Etape %d-%d moyenne : %4.1f ms\" %(i,i+5,(delais[-10:].mean())*1000))\n",
    "            \n",
    "    # Libère la mémoire GPU\n",
    "    input_device_mem.free()\n",
    "    output_device_mem.free()\n",
    "    \n",
    "    # Libère la mémoire host\n",
    "    del input_host_mem\n",
    "    del output_host_mem\n",
    "\n",
    "    print(\"Débit : %.0f images/s\" %(nbr_run / delais.sum()))\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorrt as trt\n",
    "\n",
    "# Construction de la class du logger\n",
    "class MyLogger(trt.ILogger):\n",
    "    def __init__(self):\n",
    "        trt.ILogger.__init__(self)\n",
    "\n",
    "    def log(self, severity, msg):\n",
    "        print(\"%s : %s\" %(severity,msg))\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "\n",
    "PRECISION = trt.float32\n",
    "\n",
    "logger = MyLogger()\n",
    "runtime = trt.Runtime(logger)\n",
    "\n",
    "with open(\"Regression_Resnet18_saved_model/Regression_Resnet18_FP16.engine\", \"rb\") as f:\n",
    "    engine = runtime.deserialize_cuda_engine(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = engine.create_execution_context()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = MesureDesPredictions_TRT(dataset_regression,100,PRECISION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyse de la précision du modèle optimisé"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorrt as trt\n",
    "\n",
    "# Construction de la class du logger\n",
    "class MyLogger(trt.ILogger):\n",
    "    def __init__(self):\n",
    "        trt.ILogger.__init__(self)\n",
    "\n",
    "    def log(self, severity, msg):\n",
    "        print(\"%s : %s\" %(severity,msg))\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "\n",
    "PRECISION = trt.float32\n",
    "\n",
    "logger = MyLogger()\n",
    "runtime = trt.Runtime(logger)\n",
    "\n",
    "with open(\"Regression_Resnet18_saved_model/Regression_Resnet18_FP16.engine\", \"rb\") as f:\n",
    "    engine = runtime.deserialize_cuda_engine(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = engine.create_execution_context()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "erreurs = []\n",
    "\n",
    "# Allocation des buffers d'entrée / sortie dans la mémoire GPU\n",
    "size_input = trt.volume(engine.get_binding_shape(0))* engine.max_batch_size\n",
    "size_output = trt.volume(engine.get_binding_shape(1))* engine.max_batch_size\n",
    "\n",
    "# Allocation de mémoire de type \"page-locked\" sur l'hôte\n",
    "input_host_mem = cuda.pagelocked_empty(size_input, trt.nptype(PRECISION))\n",
    "output_host_mem = cuda.pagelocked_empty(size_output, trt.nptype(PRECISION))\n",
    "\n",
    "# Allocation de mémoire dans la mémoire GPU\n",
    "input_device_mem = cuda.mem_alloc(input_host_mem.nbytes)\n",
    "output_device_mem = cuda.mem_alloc(output_host_mem.nbytes)    \n",
    "\n",
    "# Récupère les adresses en GPU des buffers entrées / sorties\n",
    "bindings = [int(input_device_mem), int(output_device_mem)]\n",
    "\n",
    "for images,heatmaps in dataset_regression.unbatch().take(10):\n",
    "    image = np.asarray(images).astype(trt.nptype(PRECISION))\n",
    "    image = np.expand_dims(image,axis=0)\n",
    "    np.copyto(input_host_mem,image.ravel())\n",
    "    \n",
    "    cuda.memcpy_htod(input_device_mem, input_host_mem)\n",
    "    context.execute(batch_size=1,bindings=bindings)        \n",
    "    cuda.memcpy_dtoh(output_host_mem, output_device_mem)    \n",
    "\n",
    "    erreur = ((np.reshape(output_host_mem,(56,56,1)) - np.asarray(heatmaps))**2).mean(axis=None)\n",
    "    print(erreur)\n",
    "    erreurs = np.append(erreurs,erreur)\n",
    "\n",
    "# Libère la mémoire GPU\n",
    "input_device_mem.free()\n",
    "output_device_mem.free()\n",
    "    \n",
    "# Libère la mémoire host\n",
    "del input_host_mem\n",
    "del output_host_mem\n",
    "    \n",
    "print(\"Erreur moyenne : %f \" %erreurs.mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Expérimentations avec le modèle optimisé"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Création de la classe Camera"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -ltrh /dev/video*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import traitlets\n",
    "import threading\n",
    "import atexit\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class Camera(traitlets.HasTraits):\n",
    "    type_camera = traitlets.Unicode(\"CSI\")\n",
    "    capture_device = traitlets.Integer(default_value=0)\n",
    "    capture_width = traitlets.Integer(default_value=1280)\n",
    "    capture_height = traitlets.Integer(default_value=720)\n",
    "    display_width = traitlets.Integer(default_value=640)\n",
    "    display_height = traitlets.Integer(default_value=480)\n",
    "    fps = traitlets.Integer(default_value=30)\n",
    "    flip = traitlets.Integer(default_value=0)\n",
    "    image = traitlets.Any()\n",
    "    video_on = traitlets.Bool(default_value=False)\n",
    "    \n",
    "    def __init__(self,*args,**kwargs):\n",
    "        super(Camera, self).__init__(*args, **kwargs)\n",
    "        self._running = False\n",
    "        self.image = np.empty((self.display_height, self.display_width, 3), dtype=np.uint8)\n",
    "        \n",
    "        if self.type_camera.find(\"CSI\")>=0:\n",
    "            self.cap = cv2.VideoCapture(self._gstreamer_pipeline_CSI(),cv2.CAP_GSTREAMER)\n",
    "        else:\n",
    "            self.cap = cv2.VideoCapture(self._gstreamer_pipeline_USB(),cv2.CAP_GSTREAMER)\n",
    "\n",
    "        if self.cap.isOpened():\n",
    "            print(\"Caméra initialisée\")\n",
    "        else:\n",
    "            print(\"Erreur d'ouverture du flux vidéo\")\n",
    "        atexit.register(self.cap.release)\n",
    "    \n",
    "    # Lecture d'une frame\n",
    "    def capture_image(self):\n",
    "        re, image = self.cap.read()\n",
    "        if re:\n",
    "            image_resized = cv2.resize(image,(int(self.display_width),int(self.display_height)))\n",
    "        return image_resized\n",
    "    \n",
    "    # ON/OFF de la capture vidéo\n",
    "    def capture_video(self,run=False):\n",
    "        if run is True:\n",
    "            self.video_on = True\n",
    "        else:\n",
    "            self.video_on = False\n",
    "    \n",
    "    # Lecture d'un flux vidéo\n",
    "    def _capture_video(self):\n",
    "        while True:\n",
    "            if not self._running:\n",
    "                break\n",
    "            self.image = self.capture_image()\n",
    "\n",
    "            \n",
    "    # Détachement de la caméra\n",
    "    def release(self):\n",
    "        self.cap.release()\n",
    "\n",
    "    # Définition du pipeline pour la caméra CSI\n",
    "    def _gstreamer_pipeline_CSI(self):\n",
    "        return(\"nvarguscamerasrc sensor-id=%d ! \"\n",
    "                \"video/x-raw(memory:NVMM),\"\n",
    "                \"width=(int)%d,height=(int)%d,\"\n",
    "                \"format=(string)NV12, framerate=(fraction)%d/1 ! \"\n",
    "                \"nvvidconv flip-method=%d ! \"\n",
    "                \"video/x-raw,\"\n",
    "                \"width=(int)%d,height=(int)%d,\"\n",
    "                \"format=(string)BGRx ! videoconvert ! \"\n",
    "                \"video/x-raw, format=(string)BGR ! \"\n",
    "                \"appsink drop=true\"\n",
    "        %(self.capture_device,self.capture_width,self.capture_height,self.fps,self.flip, self.display_width,self.display_height))\n",
    "\n",
    "    # Définition du pipeline pour la USB\n",
    "    def _gstreamer_pipeline_USB(self):\n",
    "        return(\"v4l2src device=/dev/video%d ! \"\n",
    "               \"video/x-raw, width=(int)%d, height=(int)%d, framerate=(fraction)%d/1 ! \"\n",
    "               \"videoflip method=%d ! \"\n",
    "               \"videoconvert ! \"\n",
    "               \"video/x-raw, format=(string)BGR ! appsink drop=true\"\n",
    "        %(self.capture_device,self.capture_width,self.capture_height,self.fps,self.flip))\n",
    "    \n",
    "    # Surveillance de la variable \"video_on\"\n",
    "    @traitlets.observe('video_on')\n",
    "    def _on_running(self, change):\n",
    "        if change['new'] and not change['old']:\n",
    "            # not running -> running\n",
    "            self._running = True\n",
    "            self.thread = threading.Thread(target=self._capture_video)\n",
    "            self.thread.start()\n",
    "        elif change['old'] and not change['new']:\n",
    "            # running -> not running\n",
    "            self._running = False\n",
    "            self.thread.join()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def InitCamera():\n",
    "    camera = Camera(type_camera=\"USB\",capture_device=0,\n",
    "                capture_width=640,capture_height=480,\n",
    "                display_width=224,display_height=224,\n",
    "                fps=30,flip=0)\n",
    "    return camera"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bgr8_to_jpeg(value, quality=75):\n",
    "    return bytes(cv2.imencode('.jpg', value)[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Définition des catégories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TACHE = \"visage\"\n",
    "CATEGORIES = ['nez']\n",
    "\n",
    "datasets = {}\n",
    "for name in CATEGORIES:\n",
    "    datasets[name] = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fonction d'extraction des coordonnées"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "\n",
    "def GetCoordFromHeatmap(heatmap,sampling=4):\n",
    "    hmax = tf.keras.layers.MaxPooling2D(3, 1, padding=\"same\")(heatmap)\n",
    "    keep = tf.cast(tf.equal(heatmap, hmax), tf.float32)\n",
    "    prod = hmax*keep\n",
    "\n",
    "    # Applatissement\n",
    "    prod_applati = tf.reshape(prod, (1, -1))\n",
    "    \n",
    "    # Récupère les index du maximum\n",
    "    scores, index = tf.nn.top_k(prod_applati, k=1)\n",
    "    \n",
    "    # Calcul des coordonnées\n",
    "    xs = tf.cast(index % heatmap[0,:,:,0].shape[1], tf.float32)\n",
    "    ys = tf.cast(index // heatmap[0,:,:,0].shape[0], tf.float32)\n",
    "    return xs*sampling,ys*sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interface d'acquisition en temps réel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config IPCompleter.greedy=True\n",
    "\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import shutil\n",
    "import tensorrt as trt\n",
    "import pycuda.autoinit\n",
    "import pycuda.driver as cuda\n",
    "\n",
    "class ThreadLive(threading.Thread):\n",
    "    def __init__(self, state_widget, category_widget, preview_widget, camera, fichier_engine, PRECISION):\n",
    "        threading.Thread.__init__(self)\n",
    "\n",
    "        self.cfx = cuda.Device(0).make_context()\n",
    "        \n",
    "        TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE)\n",
    "        runtime = trt.Runtime(TRT_LOGGER)\n",
    "\n",
    "        # Chargement du moteur\n",
    "        print(\"Chargement du moteur...\")\n",
    "        with open(fichier_engine, 'rb') as f:\n",
    "            buf = f.read()\n",
    "            engine = runtime.deserialize_cuda_engine(buf)\n",
    "\n",
    "        # Création du context\n",
    "        print(\"Création du context...\")\n",
    "        context = engine.create_execution_context()\n",
    "\n",
    "        # Allocation des buffers d'entrée / sortie dans la mémoire GPU\n",
    "        print(\"Allocation de la mémoire ...\")\n",
    "        size_input = trt.volume(engine.get_binding_shape(0))* engine.max_batch_size\n",
    "        size_output = trt.volume(engine.get_binding_shape(1))* engine.max_batch_size\n",
    "        input_host_mem = cuda.pagelocked_empty(size_input, trt.nptype(PRECISION))\n",
    "        output_host_mem = cuda.pagelocked_empty(size_output, trt.nptype(PRECISION))\n",
    "        input_device_mem = cuda.mem_alloc(input_host_mem.nbytes)\n",
    "        output_device_mem = cuda.mem_alloc(output_host_mem.nbytes)    \n",
    "        bindings = [int(input_device_mem), int(output_device_mem)]        \n",
    "        print(\"Modèle initialisé ...\")\n",
    "        \n",
    "        # Sauvegarde dand les variables internes de la classe\n",
    "        self.context = context\n",
    "        self.engine  = engine\n",
    "        self.input_host_mem = input_host_mem\n",
    "        self.output_host_mem = output_host_mem\n",
    "        self.input_device_mem = input_device_mem\n",
    "        self.output_device_mem = output_device_mem\n",
    "        self.bindings = bindings\n",
    "        self.PRECISION = PRECISION\n",
    "        self.state_widget = state_widget\n",
    "        self.category_widget = category_widget\n",
    "        self.camera = camera\n",
    "        self.preview_widget = preview_widget\n",
    "        \n",
    "    def run(self):\n",
    "        context = self.context\n",
    "        input_host_mem = self.input_host_mem\n",
    "        output_host_mem = self.output_host_mem\n",
    "        input_device_mem = self.input_device_mem\n",
    "        output_device_mem = self.output_device_mem\n",
    "        bindings = self.bindings\n",
    "        PRECISION = self.PRECISION\n",
    "        state_widget = self.state_widget\n",
    "        camera = self.camera\n",
    "        category_widget = self.category_widget\n",
    "        preview_widget = self.preview_widget\n",
    "        PRECISION = self.PRECISION\n",
    "        \n",
    "        while state_widget.value == 'live':\n",
    "            # Capture de l'image\n",
    "            image_camera = camera.image\n",
    "\n",
    "            # Récupération de la catégorie\n",
    "            categorie = CATEGORIES.index(category_widget.value)\n",
    "\n",
    "            # Prédiciton de la heatmap\n",
    "            image = np.asarray(image_camera).astype(trt.nptype(PRECISION))\n",
    "            image = np.expand_dims(image,axis=0)\n",
    "            np.copyto(input_host_mem,image.ravel())\n",
    "\n",
    "            self.cfx.push()\n",
    "            cuda.memcpy_htod(input_device_mem, input_host_mem)\n",
    "            context.execute(batch_size=1,bindings=bindings) \n",
    "            cuda.memcpy_dtoh(output_host_mem, output_device_mem)    \n",
    "            self.cfx.pop()\n",
    "            \n",
    "            # Extraction des coordonnées\n",
    "            heatmap = np.reshape(output_host_mem,(1,56,56,1))\n",
    "            x,y = GetCoordFromHeatmap(heatmap)\n",
    "\n",
    "            # Affichage de l'image avec les coordonnées\n",
    "            prediction = image_camera.copy()\n",
    "            prediction = cv2.circle(prediction, (x, y), 8, (255, 0, 0), 3)\n",
    "            preview_widget.value = bgr8_to_jpeg(prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "import traitlets\n",
    "from IPython.display import display\n",
    "\n",
    "fichier_engine = \"Regression_Resnet18_saved_model/Regression_Resnet18_FP16.engine\"\n",
    "PRECISION = trt.float32\n",
    "\n",
    "# Initialise la caméra\n",
    "try :\n",
    "    camera.capture_video(run=False)\n",
    "    camera.release()\n",
    "    del camera\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "camera = InitCamera()\n",
    "\n",
    "# Création du widget de la vidéo\n",
    "camera_widget = ipywidgets.Image()\n",
    "traitlets.dlink((camera, 'image'), (camera_widget, 'value'), transform=bgr8_to_jpeg)\n",
    "\n",
    "# Lancement de la vidéo\n",
    "camera.capture_video(run=True)\n",
    "camera_link = traitlets.dlink((camera, 'image'), (camera_widget, 'value'), transform=bgr8_to_jpeg)\n",
    "\n",
    "# Création des widgets\n",
    "state_widget = ipywidgets.ToggleButtons(options=['stop', 'live'], description='state', value='stop')\n",
    "category_widget = ipywidgets.Dropdown(options=CATEGORIES, description='Catégorie')\n",
    "preview_widget = ipywidgets.Image(format=\"jpeg\",width=camera.display_width, height=camera.display_height, value=bgr8_to_jpeg(camera.image))\n",
    "\n",
    "def start_live(change):\n",
    "    if change['new'] == 'live':\n",
    "        global execute_thread\n",
    "        execute_thread = ThreadLive(state_widget, category_widget, preview_widget, camera, fichier_engine, PRECISION)\n",
    "        execute_thread.start()\n",
    "    else:\n",
    "        execute_thread.join()\n",
    "       \n",
    "state_widget.observe(start_live, names='value')\n",
    "\n",
    "data_collection_widget = ipywidgets.VBox([\n",
    "    ipywidgets.HBox([camera_widget,preview_widget]),category_widget,state_widget])\n",
    "\n",
    "\n",
    "display(data_collection_widget)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}