jax-library-and-mlp-implementation.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Jax - Understanding the Library and Implementing a MLP","metadata":{}},{"cell_type":"code","source":"import time\nimport numpy as np\n\nimport matplotlib.pyplot as plt\n\nimport jax.numpy as jnp\nfrom jax import grad, jit, vmap\nfrom jax import random\nfrom jax.scipy.special import logsumexp\n\nimport tensorflow as tf\nimport tensorflow_datasets as tfds","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:17.207301Z","iopub.execute_input":"2022-02-15T10:21:17.207682Z","iopub.status.idle":"2022-02-15T10:21:25.10625Z","shell.execute_reply.started":"2022-02-15T10:21:17.207562Z","shell.execute_reply":"2022-02-15T10:21:25.105621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Importing the Dataset\n\nLet's use the MNIST dataset. For that, we will use the tensorflow built-in functions to download and process it.the","metadata":{}},{"cell_type":"code","source":"def one_hot(x, k, dtype=jnp.float32):\n    \"\"\"\n    Create a one-hot encoding of x of size k.\n    \n    x: array\n        The array to be one hot encoded\n    k: interger\n        The number of classes\n    dtype: jnp.dtype, optional(default=float32)\n        The dtype to be used on the encoding\n    \n    \"\"\"\n    return jnp.array(x[:, None] == jnp.arange(k), dtype)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:25.107751Z","iopub.execute_input":"2022-02-15T10:21:25.108077Z","iopub.status.idle":"2022-02-15T10:21:25.113023Z","shell.execute_reply.started":"2022-02-15T10:21:25.10805Z","shell.execute_reply":"2022-02-15T10:21:25.112174Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_dir = '/tmp/tfds'\n\nmnist_data, info = tfds.load(name=\"mnist\", batch_size=-1, data_dir=data_dir, with_info=True)\nmnist_data = tfds.as_numpy(mnist_data)\n\ntrain_data, test_data = mnist_data['train'], mnist_data['test']\n\nnum_labels = info.features['label'].num_classes\nh, w, c = info.features['image'].shape\nnum_pixels = h * w * c\n\n# Full train set\ntrain_images, train_labels = train_data['image'], train_data['label']\ntrain_images = jnp.reshape(train_images, (len(train_images), num_pixels))\ntrain_labels = one_hot(train_labels, num_labels)\n\n# Full test set\ntest_images, test_labels = test_data['image'], test_data['label']\ntest_images = jnp.reshape(test_images, (len(test_images), num_pixels))\ntest_labels = one_hot(test_labels, num_labels)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:25.114329Z","iopub.execute_input":"2022-02-15T10:21:25.114653Z","iopub.status.idle":"2022-02-15T10:21:32.467206Z","shell.execute_reply.started":"2022-02-15T10:21:25.114596Z","shell.execute_reply":"2022-02-15T10:21:32.466372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sets the random key for Jax, this is different from the numpy way\nrandom_state = 42\nkey = random.PRNGKey(random_state)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.468479Z","iopub.execute_input":"2022-02-15T10:21:32.468729Z","iopub.status.idle":"2022-02-15T10:21:32.548161Z","shell.execute_reply.started":"2022-02-15T10:21:32.4687Z","shell.execute_reply":"2022-02-15T10:21:32.547395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Ok, but why this random state handling is different from numpy? Well, the numpy method defines a global random state that is shared between all of the random functions. This is not a problem when we are dealing with sequential execution, however, it becomes a problem when we parallelize our functions.\n\nSee, the random state will make sure that the value is the same given an execution of that function. See for example the following code snippet:","metadata":{}},{"cell_type":"code","source":"np.random.seed(0)\n\ndef bar(): return np.random.uniform()\ndef baz(): return np.random.uniform()\n\ndef foo(): return bar() + 2 * baz()\n\nprint(foo())","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.551182Z","iopub.execute_input":"2022-02-15T10:21:32.5515Z","iopub.status.idle":"2022-02-15T10:21:32.557478Z","shell.execute_reply.started":"2022-02-15T10:21:32.551458Z","shell.execute_reply":"2022-02-15T10:21:32.55663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"np.random.seed(0)\n\ndef bar(): return np.random.uniform()\ndef baz(): return np.random.uniform()\n\ndef foo(): return 2 * baz() + bar()\n\nprint(foo())","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.558817Z","iopub.execute_input":"2022-02-15T10:21:32.559437Z","iopub.status.idle":"2022-02-15T10:21:32.569009Z","shell.execute_reply.started":"2022-02-15T10:21:32.559403Z","shell.execute_reply":"2022-02-15T10:21:32.568164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"If we change the order of the execution, the result is different. If we paralelize the calculation of these functions, we cannot guarantee that any of them will return first, so our code will be stochastic. To deal with that, every function in Jax must receive a key and that key must be unique, in a way that it does not matter in which way our calculation happens, the result will always be the same.\n\nThe way of making this work is to never use tha same key twice. But how the hell are we going to create a different key? Well, Jax has a built-in function for that, called split.\n\nThis function will receive a key and will generate two new keys that can be used on the following functions. This split is deterministic, so we do not have to worry about it changing our results.","metadata":{}},{"cell_type":"code","source":"key","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.570052Z","iopub.execute_input":"2022-02-15T10:21:32.570263Z","iopub.status.idle":"2022-02-15T10:21:32.582559Z","shell.execute_reply.started":"2022-02-15T10:21:32.570239Z","shell.execute_reply":"2022-02-15T10:21:32.58182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Here we split our original key into three subkeys\nrandom.split(key, num=3)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.583894Z","iopub.execute_input":"2022-02-15T10:21:32.584569Z","iopub.status.idle":"2022-02-15T10:21:32.825837Z","shell.execute_reply.started":"2022-02-15T10:21:32.584517Z","shell.execute_reply":"2022-02-15T10:21:32.824075Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## MLP Implementation\n\nNow, let's implement a Multilayer Perceptron (MLP) using Jax to better grab a grasp of how we can use it for Machine Learning.","metadata":{}},{"cell_type":"markdown","source":"### Parameters initialization\n\nLet's define some functions to initialize our neural network parameters with random normal values.","metadata":{}},{"cell_type":"code","source":"def random_layer_params(m, n, key, scale=1e-2):\n    \"\"\"\n    This function returns two matrices, a W matrix with shape (n, m) and a b matrix with shape (n,)\n    \n    m: integer\n        The first dimension of the W matrix\n    n: integer\n        The second dimension of the b matrix\n    key: PRNGKey\n        A Jax PRNGKey\n    scale: float, optional(default=1e-2)\n        The scale of the random numbers on the matrices\n    \"\"\"\n    # Split our key into two new keys, one for each matrix\n    w_key, b_key = random.split(key, num=2)\n    return scale * random.normal(w_key, (m,n)), scale * random.normal(b_key, (n,))","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.827011Z","iopub.execute_input":"2022-02-15T10:21:32.827221Z","iopub.status.idle":"2022-02-15T10:21:32.832811Z","shell.execute_reply.started":"2022-02-15T10:21:32.827194Z","shell.execute_reply":"2022-02-15T10:21:32.832002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_network_params(layers_sizes, key):\n    \"\"\"\n    Given a list of weights for a neural network, initializes the weights of the network\n    \n    layers_sizes: list of integers\n        The number of neurons on each layer of the network\n    key: PRNGKey\n        A Jax PRNGKey\n    \"\"\"\n    # Generate one subkey for layer in the network\n    keys = random.split(key, len(layers_sizes))\n    return [random_layer_params(m, n, k) for m, n, k in zip(layers_sizes[:-1], layers_sizes[1:], keys)]","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.834292Z","iopub.execute_input":"2022-02-15T10:21:32.834749Z","iopub.status.idle":"2022-02-15T10:21:32.847592Z","shell.execute_reply.started":"2022-02-15T10:21:32.834708Z","shell.execute_reply":"2022-02-15T10:21:32.846846Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Creating a prediction function\n\nNow we will create a prediction function for our MLP. For the activation function we are going to use the ReLU.","metadata":{}},{"cell_type":"code","source":"def relu(x):\n    return jnp.maximum(0, x)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.848815Z","iopub.execute_input":"2022-02-15T10:21:32.84914Z","iopub.status.idle":"2022-02-15T10:21:32.859252Z","shell.execute_reply.started":"2022-02-15T10:21:32.849113Z","shell.execute_reply":"2022-02-15T10:21:32.858338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict(params, x):\n    \"\"\"\n    Function to generate a prediction given weights and the activation\n    \n    params: list of matrices\n        The weights for every layer of the network, including the bias\n    x: matrix\n        The activation, or the features, to be predicted\n    \"\"\"\n    activations = x\n    \n    for w, b in params[:-1]:\n        output = jnp.dot(w.T, activations) + b\n        activations = relu(output)\n        \n    final_w, final_b = params[-1]\n    logits = jnp.dot(final_w.T, activations) + final_b\n    \n    return logits - logsumexp(logits)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.860455Z","iopub.execute_input":"2022-02-15T10:21:32.861375Z","iopub.status.idle":"2022-02-15T10:21:32.870487Z","shell.execute_reply.started":"2022-02-15T10:21:32.861329Z","shell.execute_reply":"2022-02-15T10:21:32.869861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Notice that we created a method that will output only the result for one image. As one can expect, if you pass a batch of, let's say, 30 images to this function, then it would crash because we would have a shape problem on our matrices operations. The beauty of Jax is it ability to auto batch methods for us, so we don't need to worry about handling the batches sizes.","metadata":{}},{"cell_type":"code","source":"batched_predict = vmap(predict, in_axes=(None, 0))","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.871699Z","iopub.execute_input":"2022-02-15T10:21:32.872335Z","iopub.status.idle":"2022-02-15T10:21:32.882059Z","shell.execute_reply.started":"2022-02-15T10:21:32.872268Z","shell.execute_reply":"2022-02-15T10:21:32.881523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This batched_predict method now will be able to deal with a batch of images for us. To do so, we must pass the in_axes parameters which will tell us how the parameters for our function must be mapped.\n\nNotice we have two parameters: weights and x, the weights of our neural network and the image on which they should be applied. The None for the first parameter tells vmap that we should not map this parameter anywhere, i.e, it is not batcheable. The 0 tells us that the x (or the images) should be mapped into the axis 0, i.e, the rows.","metadata":{}},{"cell_type":"markdown","source":"### Utility and loss","metadata":{}},{"cell_type":"code","source":"def accuracy(params, images, targets):\n    \"\"\"\n    Calculates the accuracy of the neural network on a set of images\n\n    params: list of matrices\n        The weights for every layer of the network, including the bias\n    images: list of matrices\n        The images to be used on the calculation\n    targets: list of labels\n        The true labels for each of the targets\n\n    \"\"\"\n    target_class = jnp.argmax(targets, axis=1)\n    \n    # Predicts the probabilities for each class and get the maximum\n    predicted_class = jnp.argmax(batched_predict(params, images), axis=1)\n\n    return jnp.mean(predicted_class == target_class)\n\ndef loss(params, images, targets):\n    preds = batched_predict(params, images)\n    return -jnp.mean(preds * targets)\n\ndef update(params, x, y):\n    grads = grad(loss)(params, x, y)\n    return [(w - step_size * dw, b - step_size * db)\n          for (w, b), (dw, db) in zip(params, grads)]","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.884535Z","iopub.execute_input":"2022-02-15T10:21:32.885311Z","iopub.status.idle":"2022-02-15T10:21:32.894986Z","shell.execute_reply.started":"2022-02-15T10:21:32.885276Z","shell.execute_reply":"2022-02-15T10:21:32.894101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Our MLP will have three hiddens layers, each one with 784, 512 and 512 neurons.\nlayer_sizes = [784, 512, 512, 10]\n\n# Training parameters\nstep_size = 0.01\nnum_epochs = 10\nbatch_size = 128\n\n# Number of labels\nn_targets = 10\n\n# Initializing the network parameters with random values\nparams = init_network_params(layer_sizes, random.PRNGKey(0))","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:32.896319Z","iopub.execute_input":"2022-02-15T10:21:32.896707Z","iopub.status.idle":"2022-02-15T10:21:35.039822Z","shell.execute_reply.started":"2022-02-15T10:21:32.89667Z","shell.execute_reply":"2022-02-15T10:21:35.038868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_batches(batch_size):\n    \"\"\"\n    This function loads the MNIST and returns a batch of images given the batch size\n    \n    batch_size: integer\n        The batch size, i.e, the number of images to be retrieved at each step\n    \n    \"\"\"\n    ds = tfds.load(name='mnist', split='train', as_supervised=True, data_dir=data_dir)\n    ds = ds.batch(batch_size).prefetch(1)\n    return tfds.as_numpy(ds)","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:35.042435Z","iopub.execute_input":"2022-02-15T10:21:35.042981Z","iopub.status.idle":"2022-02-15T10:21:35.049386Z","shell.execute_reply.started":"2022-02-15T10:21:35.042931Z","shell.execute_reply":"2022-02-15T10:21:35.048642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"non_jit_time = []\n\nfor epoch in range(num_epochs):\n    start_time = time.time()\n    for x, y in get_train_batches(batch_size):\n        x = jnp.reshape(x, (len(x), num_pixels))\n        y = one_hot(y, num_labels)\n        params = update(params, x, y)\n    epoch_time = time.time() - start_time\n    non_jit_time.append(epoch_time)\n\n    train_acc = accuracy(params, train_images, train_labels)\n    test_acc = accuracy(params, test_images, test_labels)\n    print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n    print(\"Training set accuracy {}\".format(train_acc))\n    print(\"Test set accuracy {}\\n\".format(test_acc))","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:21:35.050478Z","iopub.execute_input":"2022-02-15T10:21:35.050734Z","iopub.status.idle":"2022-02-15T10:25:11.549499Z","shell.execute_reply.started":"2022-02-15T10:21:35.050697Z","shell.execute_reply":"2022-02-15T10:25:11.548665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"jit_update = jit(update)\n\njit_time = []\n\nfor epoch in range(num_epochs):\n    start_time = time.time()\n    for x, y in get_train_batches(batch_size):\n        x = jnp.reshape(x, (len(x), num_pixels))\n        y = one_hot(y, num_labels)\n        params = jit_update(params, x, y)\n    epoch_time = time.time() - start_time\n    jit_time.append(epoch_time)\n\n    train_acc = accuracy(params, train_images, train_labels)\n    test_acc = accuracy(params, test_images, test_labels)\n    print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n    print(\"Training set accuracy {}\".format(train_acc))\n    print(\"Test set accuracy {}\\n\".format(test_acc))","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:25:11.551044Z","iopub.execute_input":"2022-02-15T10:25:11.551398Z","iopub.status.idle":"2022-02-15T10:26:32.694946Z","shell.execute_reply.started":"2022-02-15T10:25:11.551355Z","shell.execute_reply":"2022-02-15T10:26:32.694003Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(11,7))\nplt.plot(non_jit_time)\nplt.plot(jit_time)\nplt.xlabel('Epochs')\nplt.ylabel('Elapsed Time')\nplt.title('Time per epoch (s)')\nplt.legend(['Without Jit on Update', 'With Jit on Update'])\n\nprint('Non Jit Average epoch time (s): ', np.mean(non_jit_time))\nprint('Jit Average epoch time (s): ', np.mean(jit_time))","metadata":{"execution":{"iopub.status.busy":"2022-02-15T10:26:32.696301Z","iopub.execute_input":"2022-02-15T10:26:32.697176Z","iopub.status.idle":"2022-02-15T10:26:32.950191Z","shell.execute_reply.started":"2022-02-15T10:26:32.697135Z","shell.execute_reply":"2022-02-15T10:26:32.949596Z"},"trusted":true},"execution_count":null,"outputs":[]}]}