tensorcore_test-mixed_float16.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU 0: GeForce RTX 2070 (UUID: GPU-3c1ac771-91d3-c3f2-bedf-547b29fa0086)\n"
     ]
    }
   ],
   "source": [
    "!\"C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe\" -L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[name: \"/device:CPU:0\"\n",
      "device_type: \"CPU\"\n",
      "memory_limit: 268435456\n",
      "locality {\n",
      "}\n",
      "incarnation: 17880370301554667691\n",
      ", name: \"/device:XLA_CPU:0\"\n",
      "device_type: \"XLA_CPU\"\n",
      "memory_limit: 17179869184\n",
      "locality {\n",
      "}\n",
      "incarnation: 16407914299645491143\n",
      "physical_device_desc: \"device: XLA_CPU device\"\n",
      ", name: \"/device:GPU:0\"\n",
      "device_type: \"GPU\"\n",
      "memory_limit: 6531337689\n",
      "locality {\n",
      "  bus_id: 1\n",
      "  links {\n",
      "  }\n",
      "}\n",
      "incarnation: 3696133514411392922\n",
      "physical_device_desc: \"device: 0, name: GeForce RTX 2070, pci bus id: 0000:01:00.0, compute capability: 7.5\"\n",
      ", name: \"/device:XLA_GPU:0\"\n",
      "device_type: \"XLA_GPU\"\n",
      "memory_limit: 17179869184\n",
      "locality {\n",
      "}\n",
      "incarnation: 4494453348671403920\n",
      "physical_device_desc: \"device: XLA_GPU device\"\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.3.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.mixed_precision import experimental as mixed_precision\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK\n",
      "Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: GeForce RTX 2070, compute capability 7.5\n",
      "Compute dtype: float16\n",
      "Variable dtype: float32\n"
     ]
    }
   ],
   "source": [
    "policy = mixed_precision.Policy('mixed_float16')\n",
    "mixed_precision.set_policy(policy)\n",
    "print('Compute dtype: {}'.format(policy.compute_dtype))\n",
    "print('Variable dtype: {}'.format(policy.variable_dtype))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model will run with 4096 units on a GPU\n"
     ]
    }
   ],
   "source": [
    "inputs = keras.Input(shape=(784,), name='digits')\n",
    "if tf.config.list_physical_devices('GPU'):\n",
    "  print('The model will run with 4096 units on a GPU')\n",
    "  num_units = 4096\n",
    "else:\n",
    "  # Use fewer units on CPUs so the model finishes in a reasonable amount of time\n",
    "  print('The model will run with 64 units on a CPU')\n",
    "  num_units = 64\n",
    "dense1 = layers.Dense(num_units, activation='relu', name='dense_1')\n",
    "x = dense1(inputs)\n",
    "dense2 = layers.Dense(num_units, activation='relu', name='dense_2')\n",
    "x = dense2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.dtype: float16\n",
      "dense1.kernel.dtype: float32\n"
     ]
    }
   ],
   "source": [
    "print('x.dtype: %s' % x.dtype.name)\n",
    "# 'kernel' is dense1's variable\n",
    "print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Outputs dtype: float16\n"
     ]
    }
   ],
   "source": [
    "# INCORRECT: softmax and model output will be float16, when it should be float32\n",
    "outputs = layers.Dense(10, activation='softmax', name='predictions')(x)\n",
    "print('Outputs dtype: %s' % outputs.dtype.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Outputs dtype: float32\n"
     ]
    }
   ],
   "source": [
    "# CORRECT: softmax and model output are float32\n",
    "x = layers.Dense(10, name='dense_logits')(x)\n",
    "outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)\n",
    "print('Outputs dtype: %s' % outputs.dtype.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The linear activation is an identity function. So this simply casts 'outputs'\n",
    "# to float32. In this particular case, 'outputs' is already float32 so this is a\n",
    "# no-op.\n",
    "outputs = layers.Activation('linear', dtype='float32')(outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = keras.Model(inputs=inputs, outputs=outputs)\n",
    "model.compile(loss='sparse_categorical_crossentropy',\n",
    "              optimizer=keras.optimizers.RMSprop(),\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
    "x_train = x_train.reshape(60000, 784).astype('float32') / 255\n",
    "x_test = x_test.reshape(10000, 784).astype('float32') / 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_weights = model.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      "6/6 [==============================] - 0s 77ms/step - loss: 3.8920 - accuracy: 0.4055 - val_loss: 0.7667 - val_accuracy: 0.7740\n",
      "Epoch 2/5\n",
      "6/6 [==============================] - 0s 53ms/step - loss: 0.7409 - accuracy: 0.7681 - val_loss: 0.4063 - val_accuracy: 0.8844\n",
      "Epoch 3/5\n",
      "6/6 [==============================] - 0s 53ms/step - loss: 0.3158 - accuracy: 0.9053 - val_loss: 0.4937 - val_accuracy: 0.8342\n",
      "Epoch 4/5\n",
      "6/6 [==============================] - 0s 50ms/step - loss: 0.3538 - accuracy: 0.8878 - val_loss: 0.2201 - val_accuracy: 0.9359\n",
      "Epoch 5/5\n",
      "6/6 [==============================] - 0s 49ms/step - loss: 0.1910 - accuracy: 0.9427 - val_loss: 0.2384 - val_accuracy: 0.9235\n",
      "313/313 - 1s - loss: 0.2396 - accuracy: 0.9224\n",
      "Test loss: 0.23963715136051178\n",
      "Test accuracy: 0.9223999977111816\n"
     ]
    }
   ],
   "source": [
    "history = model.fit(x_train, y_train,\n",
    "                    batch_size=8192,\n",
    "                    epochs=5,\n",
    "                    validation_split=0.2)\n",
    "test_scores = model.evaluate(x_test, y_test, verbose=2)\n",
    "print('Test loss:', test_scores[0])\n",
    "print('Test accuracy:', test_scores[1])"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}