keras-core-jax-distributed.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.8.17","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q keras-core","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-08-24T06:04:10.368232Z","iopub.execute_input":"2023-08-24T06:04:10.368615Z","iopub.status.idle":"2023-08-24T06:04:14.934489Z","shell.execute_reply.started":"2023-08-24T06:04:10.368583Z","shell.execute_reply":"2023-08-24T06:04:14.933242Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stdout","text":"\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.2.1\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n","output_type":"stream"}]},{"cell_type":"code","source":"import numpy as np\nimport os, pprint, collections\n\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n\n# Note that keras_core should only be imported after the backend\n# has been configured. The backend cannot be changed once the\n# package is imported.\nimport keras_core as keras","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:04:45.931559Z","iopub.execute_input":"2023-08-24T06:04:45.932104Z","iopub.status.idle":"2023-08-24T06:04:45.937512Z","shell.execute_reply.started":"2023-08-24T06:04:45.932059Z","shell.execute_reply":"2023-08-24T06:04:45.936614Z"},"trusted":true},"execution_count":6,"outputs":[]},{"cell_type":"code","source":"pp = pprint.PrettyPrinter()\n","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:04:48.332094Z","iopub.execute_input":"2023-08-24T06:04:48.332495Z","iopub.status.idle":"2023-08-24T06:04:48.336858Z","shell.execute_reply.started":"2023-08-24T06:04:48.332460Z","shell.execute_reply":"2023-08-24T06:04:48.335992Z"},"trusted":true},"execution_count":7,"outputs":[]},{"cell_type":"code","source":"import jax\nimport jax.numpy as jnp\nimport tensorflow as tf  # just for tf.data\nimport keras_core as keras  # Keras multi-backend\n\nimport numpy as np\nfrom tqdm import tqdm\n\nfrom jax.experimental import mesh_utils\nfrom jax.sharding import Mesh\nfrom jax.sharding import NamedSharding\nfrom jax.sharding import PartitionSpec as P\n","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:04:57.273775Z","iopub.execute_input":"2023-08-24T06:04:57.274189Z","iopub.status.idle":"2023-08-24T06:04:57.285721Z","shell.execute_reply.started":"2023-08-24T06:04:57.274154Z","shell.execute_reply":"2023-08-24T06:04:57.284779Z"},"trusted":true},"execution_count":8,"outputs":[]},{"cell_type":"code","source":"\"\"\" Dataset\nClassic MNIST, loaded using tf.data\n\"\"\"\n\nBATCH_SIZE = 192\n\n(x_train, train_labels), (\n    x_eval,\n    eval_labels,\n) = keras.datasets.mnist.load_data()\nx_train = np.expand_dims(x_train, axis=-1).astype(\n    np.float32\n)  # from 28x28 to 28x28 x 1 color channel (B&W)\nx_eval = np.expand_dims(x_eval, axis=-1).astype(np.float32)\n\ntrain_data = tf.data.Dataset.from_tensor_slices((x_train, train_labels))\ntrain_data = train_data.shuffle(5000, reshuffle_each_iteration=True)\ntrain_data = train_data.batch(BATCH_SIZE, drop_remainder=True)\ntrain_data = train_data.repeat()\n\neval_data = tf.data.Dataset.from_tensor_slices((x_eval, eval_labels))\neval_data = eval_data.batch(10000)  # everything as one batch\n\nSTEPS_PER_EPOCH = len(train_labels) // BATCH_SIZE","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:05:08.500808Z","iopub.execute_input":"2023-08-24T06:05:08.501204Z","iopub.status.idle":"2023-08-24T06:05:12.969682Z","shell.execute_reply.started":"2023-08-24T06:05:08.501170Z","shell.execute_reply":"2023-08-24T06:05:12.968700Z"},"trusted":true},"execution_count":9,"outputs":[]},{"cell_type":"code","source":"# Keras \"sequential\" model building style\ndef make_backbone():\n    return keras.Sequential(\n        [\n            keras.layers.Rescaling(\n                1.0 / 255.0\n            ),  # input images are in the range [0, 255]\n            keras.layers.Conv2D(\n                filters=12, kernel_size=3, padding=\"same\", use_bias=False\n            ),\n            keras.layers.BatchNormalization(scale=False, center=True),\n            keras.layers.Activation(\"relu\"),\n            keras.layers.Conv2D(\n                filters=24,\n                kernel_size=6,\n                padding=\"same\",\n                use_bias=False,\n                strides=2,\n            ),\n            keras.layers.BatchNormalization(scale=False, center=True),\n            keras.layers.Activation(\"relu\"),\n            keras.layers.Conv2D(\n                filters=32,\n                kernel_size=6,\n                padding=\"same\",\n                use_bias=False,\n                strides=2,\n                name=\"large_k\",\n            ),\n            keras.layers.BatchNormalization(scale=False, center=True),\n            keras.layers.Activation(\"relu\"),\n        ],\n        name=\"backbone\",\n    )\n\n\ndef make_model():\n    input = keras.Input(shape=[28, 28, 1])\n    y = make_backbone()(input)\n    y = keras.layers.Flatten()(y)\n    y = keras.layers.Dense(200, activation=\"relu\")(y)\n    y = keras.layers.Dropout(0.4)(y)\n    y = keras.layers.Dense(10, activation=\"softmax\")(y)\n    model = keras.Model(inputs=input, outputs=y)\n    return model","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:05:47.144757Z","iopub.execute_input":"2023-08-24T06:05:47.145166Z","iopub.status.idle":"2023-08-24T06:05:47.159147Z","shell.execute_reply.started":"2023-08-24T06:05:47.145131Z","shell.execute_reply":"2023-08-24T06:05:47.158045Z"},"trusted":true},"execution_count":11,"outputs":[]},{"cell_type":"code","source":"\"\"\" JAX-native distribution with a Keras model\nFor now, you have to write a custom training loop for this\nNote: The features required by jax.sharding are not supported by the Colab TPU\nruntime at this time, but are available on Cloud TPU VMs and Kaggle TPU VMs.\n\"\"\"\n\nif len(jax.local_devices()) < 8:\n    raise Exception(\"This part requires 8 devices to run\")\nelse:\n    print(\"\\nIdentified local devices:\")\n    pp.pprint(jax.local_devices())\n\n# ----------------- Keras ---------------------\n\n# instantiate the model\nmodel = make_model()\n\n# learning rate\nlr = keras.optimizers.schedules.ExponentialDecay(0.01, STEPS_PER_EPOCH, 0.6)\n\n# optimizer\noptimizer = keras.optimizers.Adam(lr)\n\n# initialize all state with .build()\n(one_batch, one_batch_labels) = next(iter(train_data))\nmodel.build(one_batch)\noptimizer.build(model.trainable_variables)","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:05:59.178907Z","iopub.execute_input":"2023-08-24T06:05:59.179317Z","iopub.status.idle":"2023-08-24T06:06:01.149093Z","shell.execute_reply.started":"2023-08-24T06:05:59.179273Z","shell.execute_reply":"2023-08-24T06:06:01.148027Z"},"trusted":true},"execution_count":12,"outputs":[{"name":"stdout","text":"\nIdentified local devices:\n[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]\n","output_type":"stream"}]},{"cell_type":"code","source":"\"\"\" Distribution settings\n\n* Sharding the data on the batch axis\n* Replicating all model variables\n\nNote: this implements standard \"data parallel\" distributed training\n\n* Just for show, sharding the largest convolutional kernel along the\n  \"channels\" axis 4-ways and replicating 2-ways\n\nNote: this does not reflect a best practice but is intended to show\n      that you can split a very large kernel across multiple devices\n      if you have to\n\"\"\"\n\nprint(\n    \"\\nMostly data-parallel distribution. \"\n    \"Data is sharded across devices while the model is replicated. \"\n    \"For demo purposes, we split the largest kernel 4-ways \"\n    \"(and replicate 2-ways since we have 8 devices).\"\n)\n\n# ------------------ Jax ----------------------\n\ndevices = mesh_utils.create_device_mesh((8,))\n\n# data will be split along the batch axis\ndata_mesh = Mesh(devices, axis_names=(\"batch\",))  # naming axes of the mesh\n# naming axes of the sharded partition\ndata_sharding = NamedSharding(\n    data_mesh,\n    P(\n        \"batch\",\n    ),\n)\n# all variables will be replicated on all devices\nvar_mesh = Mesh(devices, axis_names=(\"_\"))\n# in NamedSharding, axes that are not mentioned are replicated (all axes here)\nvar_replication = NamedSharding(var_mesh, P())\n\n# for the demo, we will split the largest kernel 4-ways (and replicate 2-ways since we have 8 devices)\nlarge_kernel_mesh = Mesh(\n    devices.reshape((-1, 4)), axis_names=(None, \"out_chan\")\n)  # naming axes of the mesh\nlarge_kernel_sharding = NamedSharding(\n    large_kernel_mesh, P(None, None, None, \"out_chan\")\n)  # naming axes of the sharded partition\n","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:06:17.333176Z","iopub.execute_input":"2023-08-24T06:06:17.333610Z","iopub.status.idle":"2023-08-24T06:06:17.344880Z","shell.execute_reply.started":"2023-08-24T06:06:17.333575Z","shell.execute_reply":"2023-08-24T06:06:17.343602Z"},"trusted":true},"execution_count":13,"outputs":[{"name":"stdout","text":"\nMostly data-parallel distribution. Data is sharded across devices while the model is replicated. For demo purposes, we split the largest kernel 4-ways (and replicate 2-ways since we have 8 devices).\n","output_type":"stream"}]},{"cell_type":"code","source":"# ----------------- Keras ---------------------\n\n# Use Keras APIs to find the variable of a specific layer (we will be sharding this one in a special way)\n# In a Conv2D or Dense layer, the variables are 'kernel' and 'bias'\nspecial_layer_var = model.get_layer(\"backbone\").get_layer(\"large_k\").kernel\n\n# ------------------ Jax ----------------------\n# - accessing variables in Keras lists model.trainable_variables,\n# - model.non_trainable_variables and optimizer.variables\n\n# Apply the distribution settings to the model variables\nnon_trainable_variables = jax.device_put(\n    model.non_trainable_variables, var_replication\n)\noptimizer_variables = jax.device_put(optimizer.variables, var_replication)\n# this is what you would do replicate all trainable variables:\n# trainable_variables = jax.device_put(model.trainable_variables, var_replication)\n\n# For the demo, we split the largest kernel 4-ways instead of replicating it.\n# We still replicate all other trainable variables as in standard \"data-parallel\"\n# distributed training.\nprint_once = True\ntrainable_variables = model.trainable_variables\nfor i, v in enumerate(trainable_variables):\n    if v is special_layer_var:\n        # Apply distribution settings: sharding\n        sharded_v = jax.device_put(v, large_kernel_sharding)\n        trainable_variables[i] = sharded_v\n\n        print(\"Sharding of convolutional\", v.name, v.shape)\n        jax.debug.visualize_array_sharding(\n            jnp.reshape(sharded_v, [-1, v.shape[-1]])\n        )\n    else:\n        # Apply distribution settings: replication\n        replicated_v = jax.device_put(v, var_replication)\n        trainable_variables[i] = replicated_v\n\n        if print_once:\n            print_once = False\n            print(\n                \"\\nSharding of all other model variables (they are replicated)\"\n            )\n            jax.debug.visualize_array_sharding(\n                jnp.reshape(replicated_v, [-1, v.shape[-1]])\n            )\n\n# collect state in a handy named tuple\nTrainingState = collections.namedtuple(\n    \"TrainingState\",\n    [\"trainable_variables\", \"non_trainable_variables\", \"optimizer_variables\"],\n)\ndevice_train_state = TrainingState(\n    trainable_variables=trainable_variables,\n    non_trainable_variables=non_trainable_variables,\n    optimizer_variables=optimizer_variables,\n)\n# display data sharding\nx, y = next(iter(train_data))\nsharded_x = jax.device_put(x.numpy(), data_sharding)\nprint(\"Data sharding\")\njax.debug.visualize_array_sharding(jnp.reshape(sharded_x, [-1, 28 * 28]))\n","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:06:33.861296Z","iopub.execute_input":"2023-08-24T06:06:33.862069Z","iopub.status.idle":"2023-08-24T06:06:34.622234Z","shell.execute_reply.started":"2023-08-24T06:06:33.862032Z","shell.execute_reply":"2023-08-24T06:06:34.621028Z"},"trusted":true},"execution_count":14,"outputs":[{"name":"stdout","text":"\nSharding of all other model variables (they are replicated)\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m      \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3,4,5,6,7\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m       \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                \u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">      TPU 0,1,2,3,4,5,6,7       </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                </span>\n</pre>\n"},"metadata":{}},{"name":"stdout","text":"Sharding of convolutional kernel (6, 6, 24, 32)\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,6\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1,7\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2,4\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3,5\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> TPU 0,6 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\"> TPU 1,7 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\"> TPU 2,4 </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\"> TPU 3,5 </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n</pre>\n"},"metadata":{}},{"name":"stdout","text":"Data sharding\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;214;97;107m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;214;97;107m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;162;82m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;162;82m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n\u001b[38;2;0;0;0;48;2;231;203;148m                                     \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mTPU 6\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m                                      \u001b[0m\n\u001b[38;2;0;0;0;48;2;231;203;148m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;107;110;207m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mTPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;107;110;207m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;165;81;148m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mTPU 4\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;165;81;148m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;109;49m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mTPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;109;49m                                                                                \u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     TPU 0                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">                                     TPU 1                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">                                                                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">                                     TPU 2                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">                                                                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                     TPU 3                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                                                </span>\n<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">                                     TPU 6                                      </span>\n<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">                                                                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">                                     TPU 7                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">                                                                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">                                     TPU 4                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">                                                                                </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">                                     TPU 5                                      </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">                                                                                </span>\n</pre>\n"},"metadata":{}}]},{"cell_type":"code","source":"# ------------------ Jax ----------------------\n# - Using Keras-provided stateless APIs\n# - model.stateless_call\n# - optimizer.stateless_apply\n# These functions also work on other backends.\n\n# define loss\nloss = keras.losses.SparseCategoricalCrossentropy()\n\n\n# This is the loss function that will be differentiated.\n# Keras provides a pure functional forward pass: model.stateless_call\ndef compute_loss(trainable_variables, non_trainable_variables, x, y):\n    y_pred, updated_non_trainable_variables = model.stateless_call(\n        trainable_variables, non_trainable_variables, x\n    )\n    loss_value = loss(y, y_pred)\n    return loss_value, updated_non_trainable_variables\n\n\n# function to compute gradients\ncompute_gradients = jax.value_and_grad(compute_loss, has_aux=True)\n\n\n# Trainig step: Keras provides a pure functional optimizer.stateless_apply\n@jax.jit\ndef train_step(train_state, x, y):\n    (loss_value, non_trainable_variables), grads = compute_gradients(\n        train_state.trainable_variables,\n        train_state.non_trainable_variables,\n        x,\n        y,\n    )\n\n    trainable_variables, optimizer_variables = optimizer.stateless_apply(\n        train_state.optimizer_variables, grads, train_state.trainable_variables\n    )\n\n    return loss_value, TrainingState(\n        trainable_variables, non_trainable_variables, optimizer_variables\n    )\n\n\n# training loop\nEPOCHS = 200\nprint(\"\\nTrainig:\")\ndata_iter = iter(train_data)\nfor epoch in range(EPOCHS):\n    for i in tqdm(range(STEPS_PER_EPOCH)):\n        x, y = next(data_iter)\n        sharded_x = jax.device_put(x.numpy(), data_sharding)\n        loss_value, device_train_state = train_step(\n            device_train_state, sharded_x, y.numpy()\n        )\n    print(\"Epoch\", epoch, \"loss:\", loss_value)\n\n# The output of the model is still sharded. Sharding follows the data.\n\ndata, labels = next(iter(eval_data))\nsharded_data = jax.device_put(data.numpy(), data_sharding)\n\n\n@jax.jit\ndef predict(data):\n    predictions, updated_non_trainable_variables = model.stateless_call(\n        device_train_state.trainable_variables,\n        device_train_state.non_trainable_variables,\n        data,\n    )\n    return predictions\n\n\npredictions = predict(sharded_data)\nprint(\"\\nModel output sharding follows data sharding:\")\njax.debug.visualize_array_sharding(predictions)\n\n# Post-processing model state update to write them back into the model\nupdate = lambda variable, value: variable.assign(value)\n\njax.tree_map(\n    update, model.trainable_variables, device_train_state.trainable_variables\n)\njax.tree_map(\n    update,\n    model.non_trainable_variables,\n    device_train_state.non_trainable_variables,\n)\njax.tree_map(\n    update, optimizer.variables, device_train_state.optimizer_variables\n)\n\n# check that the model has the new state by running an eval\n# known issue: the optimizer should not be required here\nmodel.compile(\n    loss=keras.losses.SparseCategoricalCrossentropy(),\n    metrics=[keras.metrics.SparseCategoricalAccuracy()],\n)\nprint(\"\\nUpdating model and running an eval:\")\nloss, accuracy = model.evaluate(eval_data)\nprint(\"The model achieved an evaluation accuracy of:\", accuracy)","metadata":{"execution":{"iopub.status.busy":"2023-08-24T06:07:49.782317Z","iopub.execute_input":"2023-08-24T06:07:49.783378Z","iopub.status.idle":"2023-08-24T06:14:58.317125Z","shell.execute_reply.started":"2023-08-24T06:07:49.783338Z","shell.execute_reply":"2023-08-24T06:14:58.315646Z"},"trusted":true},"execution_count":16,"outputs":[{"name":"stdout","text":"\nTrainig:\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:04<00:00, 64.95it/s] \n","output_type":"stream"},{"name":"stdout","text":"Epoch 0 loss: 0.0051613655\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 154.30it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 1 loss: 0.023073694\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.40it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 2 loss: 0.01335797\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.26it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 3 loss: 0.004434864\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.67it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 4 loss: 0.002694505\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.03it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 5 loss: 0.011267702\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.24it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 6 loss: 0.005084862\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.90it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 7 loss: 0.0037868065\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.00it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 8 loss: 0.004354721\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.06it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 9 loss: 0.0033651057\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.88it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 10 loss: 0.007970343\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.56it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 11 loss: 0.0022554111\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.90it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 12 loss: 0.0025315904\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.35it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 13 loss: 0.0027730335\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.88it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 14 loss: 0.09231086\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.05it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 15 loss: 0.00050703855\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.97it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 16 loss: 0.005900763\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.30it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 17 loss: 0.0047382293\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.39it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 18 loss: 0.0057516964\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 19 loss: 0.000799091\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.69it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 20 loss: 0.0020407606\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.54it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 21 loss: 0.0017153182\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.87it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 22 loss: 0.003606855\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.99it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 23 loss: 0.005014763\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.06it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 24 loss: 0.0015248026\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.57it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 25 loss: 0.0032154913\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.67it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 26 loss: 0.002289772\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.41it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 27 loss: 0.0077875424\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 28 loss: 0.0018546786\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.81it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 29 loss: 0.003362091\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.59it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 30 loss: 0.0019750698\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.48it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 31 loss: 0.0077667153\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.61it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 32 loss: 0.00509238\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.86it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 33 loss: 0.09413421\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.38it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 34 loss: 0.010569198\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.15it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 35 loss: 0.004389434\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.04it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 36 loss: 0.0046832967\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 37 loss: 0.0033779717\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.25it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 38 loss: 0.0012988928\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.23it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 39 loss: 0.00148514\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 153.47it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 40 loss: 0.08799639\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.90it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 41 loss: 0.102935396\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 154.20it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 42 loss: 0.0064084753\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 154.79it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 43 loss: 0.0059722695\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.49it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 44 loss: 0.003496557\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.02it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 45 loss: 0.0066671125\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.05it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 46 loss: 0.0010258653\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.70it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 47 loss: 0.001148876\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.08it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 48 loss: 0.0022844695\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 153.12it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 49 loss: 0.0019500902\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.60it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 50 loss: 0.0025554872\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.82it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 51 loss: 0.006012059\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.54it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 52 loss: 0.006816498\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.91it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 53 loss: 0.012626072\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.17it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 54 loss: 0.0020623181\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 154.11it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 55 loss: 0.00055154663\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.58it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 56 loss: 0.0005175198\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.14it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 57 loss: 0.08928703\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.05it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 58 loss: 0.0010914998\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.39it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 59 loss: 0.0023103198\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.00it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 60 loss: 0.00221428\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.17it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 61 loss: 0.00079013285\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.70it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 62 loss: 0.009109486\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.87it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 63 loss: 0.006482146\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.95it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 64 loss: 0.009222313\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.95it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 65 loss: 0.0043873936\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.83it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 66 loss: 0.0023372069\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 67 loss: 0.01008275\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.99it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 68 loss: 0.00083925517\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.24it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 69 loss: 0.011575425\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.08it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 70 loss: 0.0026201992\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.46it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 71 loss: 0.0012501717\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.57it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 72 loss: 0.0050113937\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.58it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 73 loss: 0.001466934\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.11it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 74 loss: 0.0005016887\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 75 loss: 0.0052876724\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.68it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 76 loss: 0.005690371\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.04it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 77 loss: 0.003983121\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 78 loss: 0.0013184206\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 79 loss: 0.010736575\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.26it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 80 loss: 0.019568603\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.57it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 81 loss: 0.004313573\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.69it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 82 loss: 0.010991074\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.66it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 83 loss: 0.0012171332\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.07it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 84 loss: 0.00055077707\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 153.65it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 85 loss: 0.0073316023\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.77it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 86 loss: 0.0028434987\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 87 loss: 0.00817379\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 143.19it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 88 loss: 0.004961028\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 140.57it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 89 loss: 0.012947354\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 138.62it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 90 loss: 0.001499664\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 140.84it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 91 loss: 0.0037881478\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.77it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 92 loss: 0.009549415\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.53it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 93 loss: 0.0078542745\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.17it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 94 loss: 0.004703642\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 140.28it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 95 loss: 0.014925811\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 139.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 96 loss: 0.008275516\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 142.86it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 97 loss: 0.005108869\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.96it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 98 loss: 0.019720722\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.66it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 99 loss: 0.001974451\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.79it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 100 loss: 0.0014513393\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.89it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 101 loss: 0.012814565\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.26it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 102 loss: 0.005026892\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.33it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 103 loss: 0.001023326\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.51it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 104 loss: 0.0073186615\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.34it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 105 loss: 0.004433358\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.76it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 106 loss: 0.001427429\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.81it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 107 loss: 0.0007770895\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.20it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 108 loss: 0.0011203168\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.95it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 109 loss: 0.005042796\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 110 loss: 0.002338284\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 143.27it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 111 loss: 0.0012297763\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.27it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 112 loss: 0.008016388\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.19it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 113 loss: 0.0013119545\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.85it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 114 loss: 0.0035683354\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.06it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 115 loss: 0.0020889153\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.50it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 116 loss: 0.0018147466\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.14it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 117 loss: 0.003004505\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.44it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 118 loss: 0.0057763434\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.48it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 119 loss: 0.001639887\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.29it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 120 loss: 0.0042094085\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.79it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 121 loss: 0.0016605941\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.75it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 122 loss: 0.0046089813\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.33it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 123 loss: 0.00039930813\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.72it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 124 loss: 0.0012195\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.08it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 125 loss: 0.004652785\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.83it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 126 loss: 0.0021859822\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 140.12it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 127 loss: 0.010130579\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.52it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 128 loss: 0.0033728941\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.09it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 129 loss: 0.0073223854\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.13it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 130 loss: 0.0061236704\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.07it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 131 loss: 0.0016905258\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.58it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 132 loss: 0.002220683\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.24it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 133 loss: 0.0012796351\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.92it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 134 loss: 0.010886195\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.37it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 135 loss: 0.0010391874\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.79it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 136 loss: 0.0029520458\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.66it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 137 loss: 0.0020032595\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.14it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 138 loss: 0.08622976\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.25it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 139 loss: 0.109192364\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.30it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 140 loss: 0.0012985271\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.04it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 141 loss: 0.0012599411\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.09it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 142 loss: 0.002931129\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.44it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 143 loss: 0.026181644\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.53it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 144 loss: 0.008663185\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.27it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 145 loss: 0.0047908197\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.94it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 146 loss: 0.0014185531\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.06it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 147 loss: 0.0027114234\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.33it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 148 loss: 0.002854012\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.56it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 149 loss: 0.0036366757\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.35it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 150 loss: 0.008331113\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.92it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 151 loss: 0.0060296264\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.78it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 152 loss: 0.0013188244\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.75it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 153 loss: 0.0011017967\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.80it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 154 loss: 0.0027540554\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.64it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 155 loss: 0.016113192\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.69it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 156 loss: 0.0049481853\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.84it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 157 loss: 0.0024123907\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.31it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 158 loss: 0.011917135\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.83it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 159 loss: 0.0059766974\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.95it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 160 loss: 0.007746275\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.18it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 161 loss: 0.0021171786\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.42it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 162 loss: 0.0025407393\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.47it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 163 loss: 0.0012321952\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.32it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 164 loss: 0.0050710808\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.75it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 165 loss: 0.0018525075\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.31it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 166 loss: 0.0073213195\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.10it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 167 loss: 0.0028305366\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 153.94it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 168 loss: 0.001279901\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.78it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 169 loss: 0.0007328101\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.52it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 170 loss: 0.005517888\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.51it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 171 loss: 0.0039075078\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.07it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 172 loss: 0.00078372326\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.90it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 173 loss: 0.00035636983\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.01it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 174 loss: 0.002435878\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.56it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 175 loss: 0.0015519147\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.11it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 176 loss: 0.00381902\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.62it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 177 loss: 0.000856303\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.77it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 178 loss: 0.0051490823\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 146.69it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 179 loss: 0.008321327\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.64it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 180 loss: 0.002521487\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.83it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 181 loss: 0.0009861202\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.69it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 182 loss: 0.010210268\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 150.27it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 183 loss: 0.0035283607\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.55it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 184 loss: 0.008553078\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.20it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 185 loss: 0.00052486436\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 152.22it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 186 loss: 0.0019686937\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.86it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 187 loss: 0.0027586464\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.57it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 188 loss: 0.0058374805\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.67it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 189 loss: 0.0048602917\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.52it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 190 loss: 0.0057701813\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 149.10it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 191 loss: 0.003368033\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 151.71it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 192 loss: 0.003662585\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 148.22it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 193 loss: 0.09084996\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 147.28it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 194 loss: 0.005811976\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 142.47it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 195 loss: 0.004816776\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.73it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 196 loss: 0.005332428\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 144.56it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 197 loss: 0.0086028315\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 142.43it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 198 loss: 0.0024974593\n","output_type":"stream"},{"name":"stderr","text":"100%|██████████| 312/312 [00:02<00:00, 145.95it/s]\n","output_type":"stream"},{"name":"stdout","text":"Epoch 199 loss: 0.0036550718\n\nModel output sharding follows data sharding:\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mTPU 6\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\n\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mTPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mTPU 4\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mTPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  TPU 1  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  TPU 2  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 3  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  TPU 6  </span>\n<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  TPU 7  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  TPU 4  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  TPU 5  </span>\n<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n</pre>\n"},"metadata":{}},{"name":"stdout","text":"\nUpdating model and running an eval:\n\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m3s\u001b[0m 3s/step - loss: 0.0406 - sparse_categorical_accuracy: 0.9891\nThe model achieved an evaluation accuracy of: 0.9890999794006348\n","output_type":"stream"}]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}