{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.18","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":3004,"databundleVersionId":861823,"sourceType":"competition"}],"dockerImageVersionId":31091,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:30.282573Z","iopub.execute_input":"2025-08-25T19:23:30.283347Z","iopub.status.idle":"2025-08-25T19:23:30.303135Z","shell.execute_reply.started":"2025-08-25T19:23:30.283313Z","shell.execute_reply":"2025-08-25T19:23:30.297099Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport optax","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:30.305558Z","iopub.execute_input":"2025-08-25T19:23:30.305788Z","iopub.status.idle":"2025-08-25T19:23:30.318549Z","shell.execute_reply.started":"2025-08-25T19:23:30.305766Z","shell.execute_reply":"2025-08-25T19:23:30.312710Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = pd.read_csv('/kaggle/input/digit-recognizer/train.csv')\ntest_dataset = pd.read_csv('/kaggle/input/digit-recognizer/test.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:30.320497Z","iopub.execute_input":"2025-08-25T19:23:30.320729Z","iopub.status.idle":"2025-08-25T19:23:32.933698Z","shell.execute_reply.started":"2025-08-25T19:23:30.320706Z","shell.execute_reply":"2025-08-25T19:23:32.929098Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = jnp.array(train_dataset)\ntest_dataset = jnp.array(test_dataset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:32.936142Z","iopub.execute_input":"2025-08-25T19:23:32.936402Z","iopub.status.idle":"2025-08-25T19:23:33.152902Z","shell.execute_reply.started":"2025-08-25T19:23:32.936380Z","shell.execute_reply":"2025-08-25T19:23:33.147933Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"labels = train_dataset[:, 0]\none_hot_encoded = jax.nn.one_hot(labels, num_classes=10)\ntrain_ds = train_dataset[:, 1:].reshape(-1, 28, 28)\ntrain_ds = train_ds[..., jnp.newaxis]/255.0\ntest_ds = test_dataset.reshape(-1, 28, 28)/255.0\ntest_ds = test_ds[..., jnp.newaxis]\n\nprint(one_hot_encoded.shape)\nprint(train_ds.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:33.154981Z","iopub.execute_input":"2025-08-25T19:23:33.155236Z","iopub.status.idle":"2025-08-25T19:23:33.171595Z","shell.execute_reply.started":"2025-08-25T19:23:33.155211Z","shell.execute_reply":"2025-08-25T19:23:33.167025Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.imshow(train_ds[3])\nprint(labels[3])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:33.172729Z","iopub.execute_input":"2025-08-25T19:23:33.172955Z","iopub.status.idle":"2025-08-25T19:23:33.308735Z","shell.execute_reply.started":"2025-08-25T19:23:33.172932Z","shell.execute_reply":"2025-08-25T19:23:33.304001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from typing import List\nimport flax.linen as nn\nfrom flax.linen.activation import relu, softmax\n\nclass DigitRecogConv(nn.Module):\n\n    ConvLayers: List[int]\n    DenseLayers: List[int]\n    \n\n    @nn.compact\n    def __call__(self, x, train: bool):\n        \n        for i, layer in enumerate(self.ConvLayers):\n            if i%2 == 0:\n                x = nn.Conv(features=layer, kernel_size=(3,3), padding=1)(x)\n            else:\n                x = nn.Conv(features=layer, kernel_size=(3,3), strides=(2,2))(x)\n\n            x = nn.BatchNorm(use_running_average=not train)(x)\n            \n            x = relu(x)\n           \n        x = x.reshape(x.shape[0], -1)\n        \n        for i, layer in enumerate(self.DenseLayers):\n            x = nn.Dense(features=layer)(x)\n            if i < len(self.DenseLayers)-1:\n                x = nn.BatchNorm(use_running_average=not train)(x)\n                x = relu(x)\n\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:33.309475Z","iopub.execute_input":"2025-08-25T19:23:33.309682Z","iopub.status.idle":"2025-08-25T19:23:33.323769Z","shell.execute_reply.started":"2025-08-25T19:23:33.309654Z","shell.execute_reply":"2025-08-25T19:23:33.319297Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"convLayers = [32, 64, 128]\ndenseLayers = [256, 128, 64, 10]\n\nkey = jax.random.PRNGKey(42)\n\nmodel = DigitRecogConv(ConvLayers=convLayers, DenseLayers=denseLayers)\nvariables = model.init(key, jnp.ones((1, 28, 28, 1)), train=True)\nparams = variables['params']\nbatch_stats = variables['batch_stats']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:33.326121Z","iopub.execute_input":"2025-08-25T19:23:33.326511Z","iopub.status.idle":"2025-08-25T19:23:33.414053Z","shell.execute_reply.started":"2025-08-25T19:23:33.326486Z","shell.execute_reply":"2025-08-25T19:23:33.408400Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(train_ds.shape)\n\n\npred, updates = model.apply({'params': params, 'batch_stats': batch_stats}, train_ds, train=True, mutable=['batch_stats'])\nbatch_stats = updates['batch_stats']\npred.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:33.416112Z","iopub.execute_input":"2025-08-25T19:23:33.416371Z","iopub.status.idle":"2025-08-25T19:23:34.094017Z","shell.execute_reply.started":"2025-08-25T19:23:33.416347Z","shell.execute_reply":"2025-08-25T19:23:34.089284Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"lr = 0.001\nnum_steps = 400\n\noptz = optax.adam(lr)\noptz_params = optz.init(params)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:34.096134Z","iopub.execute_input":"2025-08-25T19:23:34.096419Z","iopub.status.idle":"2025-08-25T19:23:34.127377Z","shell.execute_reply.started":"2025-08-25T19:23:34.096395Z","shell.execute_reply":"2025-08-25T19:23:34.121396Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def loss_fn(params, batch_stats, batch, one_hot_enc): \n    pred, updates = model.apply({'params': params, 'batch_stats': batch_stats}, batch, train=True, mutable=['batch_stats'])\n    \n    return optax.softmax_cross_entropy(pred, one_hot_enc).mean(), (updates['batch_stats'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:34.129118Z","iopub.execute_input":"2025-08-25T19:23:34.129366Z","iopub.status.idle":"2025-08-25T19:23:34.140457Z","shell.execute_reply.started":"2025-08-25T19:23:34.129342Z","shell.execute_reply":"2025-08-25T19:23:34.134847Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@jax.jit\ndef train_loop(batch, params, batch_stats, optz_params):\n    (loss, (new_batch_stats)), grad = jax.value_and_grad(loss_fn, has_aux=True)(params, batch_stats, batch, one_hot_encoded)\n    \n    update, new_optz_params = optz.update(grad, optz_params, params)\n    new_params = optax.apply_updates(params, update)\n\n    return new_params, new_batch_stats, new_optz_params, loss\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:34.142850Z","iopub.execute_input":"2025-08-25T19:23:34.143089Z","iopub.status.idle":"2025-08-25T19:23:34.179423Z","shell.execute_reply.started":"2025-08-25T19:23:34.143064Z","shell.execute_reply":"2025-08-25T19:23:34.174711Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for i in range(num_steps):\n    params, batch_stats, optz_params, loss = train_loop(train_ds, params, batch_stats, optz_params)\n\n    print(f\"Step {i} completed with loss: {loss}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:23:34.181714Z","iopub.execute_input":"2025-08-25T19:23:34.181956Z","iopub.status.idle":"2025-08-25T19:26:15.101211Z","shell.execute_reply.started":"2025-08-25T19:23:34.181932Z","shell.execute_reply":"2025-08-25T19:26:15.096221Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Waiting for all training steps to complete...\")\njax.tree_util.tree_map(lambda x: x.block_until_ready(), (params, batch_stats))\nprint(\"All steps complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:26:15.102377Z","iopub.execute_input":"2025-08-25T19:26:15.102640Z","iopub.status.idle":"2025-08-25T19:26:15.117438Z","shell.execute_reply.started":"2025-08-25T19:26:15.102616Z","shell.execute_reply":"2025-08-25T19:26:15.110987Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import accuracy_score, confusion_matrix\n\npred = model.apply({'params': params, 'batch_stats': batch_stats}, train_ds, train=False)\npred = jax.nn.softmax(pred, axis=1)\npred = jnp.argmax(pred, axis=1)\naccuracy = accuracy_score(pred, labels)\nprint(accuracy)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:26:15.119013Z","iopub.execute_input":"2025-08-25T19:26:15.119237Z","iopub.status.idle":"2025-08-25T19:26:15.619610Z","shell.execute_reply.started":"2025-08-25T19:26:15.119217Z","shell.execute_reply":"2025-08-25T19:26:15.614068Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with open('./submission.csv', 'w') as f:\n    f.write('ImageId,Label\\n')\n    test_pred = model.apply({'params': params, 'batch_stats': batch_stats}, test_ds, train=False)\n    test_pred = jax.nn.softmax(test_pred, axis=1)\n    test_pred = jnp.argmax(test_pred, axis=1)\n\n    for i in range(len(test_pred)):\n        f.write(f\"{i+1},{test_pred[i]}\\n\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:26:15.621480Z","iopub.execute_input":"2025-08-25T19:26:15.621698Z","iopub.status.idle":"2025-08-25T19:26:37.456225Z","shell.execute_reply.started":"2025-08-25T19:26:15.621677Z","shell.execute_reply":"2025-08-25T19:26:37.449470Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.imshow(test_ds[1])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T19:26:37.458730Z","iopub.execute_input":"2025-08-25T19:26:37.458979Z","iopub.status.idle":"2025-08-25T19:26:37.597134Z","shell.execute_reply.started":"2025-08-25T19:26:37.458954Z","shell.execute_reply":"2025-08-25T19:26:37.591933Z"}},"outputs":[],"execution_count":null}]}