{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.10","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":"# 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","execution":{"iopub.status.busy":"2023-01-03T22:12:22.105655Z","iopub.execute_input":"2023-01-03T22:12:22.106031Z","iopub.status.idle":"2023-01-03T22:12:22.233792Z","shell.execute_reply.started":"2023-01-03T22:12:22.105933Z","shell.execute_reply":"2023-01-03T22:12:22.233018Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stdout","text":"/kaggle/input/tpu-getting-started/sample_submission.csv\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/09-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/14-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/01-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/13-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/00-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/07-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/12-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/11-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/05-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/04-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/15-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/06-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/02-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/03-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/08-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/val/10-224x224-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/04-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/06-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/05-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/08-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/15-224x224-452.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/09-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/10-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/07-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/00-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/01-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/12-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/14-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/03-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/02-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/11-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/test/13-224x224-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/13-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/11-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/05-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/00-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/02-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/14-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/06-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/10-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/01-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/03-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/08-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/09-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/15-224x224-783.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/07-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/04-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-224x224/train/12-224x224-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/05-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/08-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/01-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/03-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/13-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/10-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/07-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/04-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/14-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/15-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/12-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/11-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/02-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/00-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/06-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/val/09-331x331-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/14-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/13-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/03-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/05-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/09-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/04-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/01-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/11-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/07-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/08-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/00-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/02-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/15-331x331-452.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/10-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/12-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/test/06-331x331-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/15-331x331-783.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/02-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/14-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/04-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/08-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/12-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/03-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/05-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/00-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/06-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/07-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/09-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/10-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/01-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/13-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-331x331/train/11-331x331-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/05-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/07-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/15-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/08-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/06-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/00-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/11-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/04-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/09-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/02-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/10-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/12-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/14-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/13-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/03-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/val/01-192x192-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/06-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/12-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/08-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/07-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/00-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/03-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/13-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/05-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/01-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/14-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/10-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/04-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/15-192x192-452.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/09-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/11-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/test/02-192x192-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/06-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/11-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/12-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/03-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/09-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/08-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/04-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/05-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/14-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/07-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/13-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/15-192x192-783.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/02-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/10-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/01-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-192x192/train/00-192x192-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/02-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/12-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/09-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/00-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/13-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/03-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/14-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/04-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/10-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/05-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/08-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/11-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/07-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/15-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/01-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/val/06-512x512-232.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/15-512x512-452.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/02-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/03-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/01-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/00-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/13-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/07-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/10-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/04-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/14-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/09-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/08-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/11-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/05-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/06-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/test/12-512x512-462.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/05-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/00-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/11-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/15-512x512-783.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/10-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/12-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/01-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/14-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/13-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/07-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/03-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/06-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/08-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/02-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/09-512x512-798.tfrec\n/kaggle/input/tpu-getting-started/tfrecords-jpeg-512x512/train/04-512x512-798.tfrec\n","output_type":"stream"}]},{"cell_type":"code","source":"import os\nif 'TPU_NAME' in os.environ:\n  import requests\n  if 'TPU_DRIVER_MODE' not in globals():\n    url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'\n    resp = requests.post(url)\n    TPU_DRIVER_MODE = 1\n\n\n  from jax.config import config\n  config.FLAGS.jax_xla_backend = \"tpu_driver\"\n  config.FLAGS.jax_backend_target = os.environ['TPU_NAME']\n  print('Registered TPU:', config.FLAGS.jax_backend_target)\nelse:\n  print('No TPU detected. Can be changed under \"Runtime/Change runtime type\".')\n\n!pip install jax jaxlib flax optax datasets transformers\n","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:12:22.234951Z","iopub.execute_input":"2023-01-03T22:12:22.235168Z","iopub.status.idle":"2023-01-03T22:14:45.171563Z","shell.execute_reply.started":"2023-01-03T22:12:22.235143Z","shell.execute_reply":"2023-01-03T22:14:45.170286Z"},"trusted":true},"execution_count":2,"outputs":[{"name":"stdout","text":"Registered TPU: grpc://10.0.0.2:8470\nRequirement already satisfied: jax in /opt/conda/lib/python3.7/site-packages (0.2.19)\nRequirement already satisfied: jaxlib in /opt/conda/lib/python3.7/site-packages (0.1.70)\nCollecting flax\n  Downloading flax-0.6.3-py3-none-any.whl (197 kB)\n\u001b[K     |████████████████████████████████| 197 kB 2.1 MB/s eta 0:00:01\n\u001b[?25hCollecting optax\n  Downloading optax-0.1.4-py3-none-any.whl (154 kB)\n\u001b[K     |████████████████████████████████| 154 kB 23.6 MB/s eta 0:00:01\n\u001b[?25hRequirement already satisfied: datasets in /opt/conda/lib/python3.7/site-packages (1.12.1)\nRequirement already satisfied: transformers in /opt/conda/lib/python3.7/site-packages (4.5.1)\nRequirement already satisfied: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.0)\nRequirement already satisfied: numpy>=1.18 in /opt/conda/lib/python3.7/site-packages (from jax) (1.19.5)\nRequirement already satisfied: opt-einsum in /opt/conda/lib/python3.7/site-packages (from jax) (3.3.0)\nRequirement already satisfied: scipy in /opt/conda/lib/python3.7/site-packages (from jaxlib) (1.7.1)\nRequirement already satisfied: flatbuffers<3.0,>=1.12 in /opt/conda/lib/python3.7/site-packages (from jaxlib) (1.12)\nCollecting orbax\n  Downloading orbax-0.0.23-py3-none-any.whl (66 kB)\n\u001b[K     |████████████████████████████████| 66 kB 3.1 MB/s  eta 0:00:01\n\u001b[?25hCollecting typing-extensions>=4.1.1\n  Downloading typing_extensions-4.4.0-py3-none-any.whl (26 kB)\nCollecting tensorstore\n  Downloading tensorstore-0.1.28-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (8.3 MB)\n\u001b[K     |████████████████████████████████| 8.3 MB 33.1 MB/s eta 0:00:01\n\u001b[?25hRequirement already satisfied: msgpack in /opt/conda/lib/python3.7/site-packages (from flax) (1.0.2)\nRequirement already satisfied: PyYAML>=5.4.1 in /opt/conda/lib/python3.7/site-packages (from flax) (5.4.1)\nCollecting jax\n  Downloading jax-0.3.25.tar.gz (1.1 MB)\n\u001b[K     |████████████████████████████████| 1.1 MB 47.8 MB/s eta 0:00:01\n\u001b[?25hRequirement already satisfied: matplotlib in /opt/conda/lib/python3.7/site-packages (from flax) (3.4.3)\nCollecting rich>=11.1\n  Downloading rich-13.0.0-py3-none-any.whl (238 kB)\n\u001b[K     |████████████████████████████████| 238 kB 32.0 MB/s eta 0:00:01\n\u001b[?25hCollecting numpy>=1.18\n  Downloading numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (15.7 MB)\n\u001b[K     |████████████████████████████████| 15.7 MB 38.6 MB/s eta 0:00:01\n\u001b[?25hCollecting chex>=0.1.5\n  Downloading chex-0.1.5-py3-none-any.whl (85 kB)\n\u001b[K     |████████████████████████████████| 85 kB 2.9 MB/s  eta 0:00:01\n\u001b[?25hRequirement already satisfied: aiohttp in /opt/conda/lib/python3.7/site-packages (from datasets) (3.7.4.post0)\nRequirement already satisfied: packaging in /opt/conda/lib/python3.7/site-packages (from datasets) (21.0)\nRequirement already satisfied: huggingface-hub<0.1.0,>=0.0.14 in /opt/conda/lib/python3.7/site-packages (from datasets) (0.0.17)\nRequirement already satisfied: fsspec[http]>=2021.05.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (2021.8.1)\nRequirement already satisfied: dill in /opt/conda/lib/python3.7/site-packages (from datasets) (0.3.4)\nRequirement already satisfied: xxhash in /opt/conda/lib/python3.7/site-packages (from datasets) (2.0.2)\nRequirement already satisfied: multiprocess in /opt/conda/lib/python3.7/site-packages (from datasets) (0.70.12.2)\nRequirement already satisfied: importlib-metadata in /opt/conda/lib/python3.7/site-packages (from datasets) (3.4.0)\nRequirement already satisfied: pandas in /opt/conda/lib/python3.7/site-packages (from datasets) (1.3.2)\nRequirement already satisfied: requests>=2.19.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (2.25.1)\nRequirement already satisfied: tqdm>=4.62.1 in /opt/conda/lib/python3.7/site-packages (from datasets) (4.62.1)\nRequirement already satisfied: pyarrow!=4.0.0,>=1.0.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (5.0.0)\nRequirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers) (3.0.12)\nRequirement already satisfied: sacremoses in /opt/conda/lib/python3.7/site-packages (from transformers) (0.0.45)\nRequirement already satisfied: tokenizers<0.11,>=0.10.1 in /opt/conda/lib/python3.7/site-packages (from transformers) (0.10.3)\nRequirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.7/site-packages (from transformers) (2021.8.3)\nRequirement already satisfied: six in /opt/conda/lib/python3.7/site-packages (from absl-py->jax) (1.15.0)\nRequirement already satisfied: dm-tree>=0.1.5 in /opt/conda/lib/python3.7/site-packages (from chex>=0.1.5->optax) (0.1.6)\nRequirement already satisfied: toolz>=0.9.0 in /opt/conda/lib/python3.7/site-packages (from chex>=0.1.5->optax) (0.11.1)\nRequirement already satisfied: pyparsing>=2.0.2 in /opt/conda/lib/python3.7/site-packages (from packaging->datasets) (2.4.7)\nRequirement already satisfied: chardet<5,>=3.0.2 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (4.0.0)\nRequirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (2021.5.30)\nRequirement already satisfied: idna<3,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (2.10)\nRequirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (1.26.6)\nRequirement already satisfied: pygments<3.0.0,>=2.6.0 in /opt/conda/lib/python3.7/site-packages (from rich>=11.1->flax) (2.10.0)\nCollecting commonmark<0.10.0,>=0.9.0\n  Downloading commonmark-0.9.1-py2.py3-none-any.whl (51 kB)\n\u001b[K     |████████████████████████████████| 51 kB 4.1 MB/s  eta 0:00:01\n\u001b[?25hRequirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (21.2.0)\nRequirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (5.1.0)\nRequirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (1.6.3)\nRequirement already satisfied: async-timeout<4.0,>=3.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (3.0.1)\nRequirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata->datasets) (3.5.0)\nRequirement already satisfied: pillow>=6.2.0 in /opt/conda/lib/python3.7/site-packages (from matplotlib->flax) (8.2.0)\nRequirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.7/site-packages (from matplotlib->flax) (1.3.1)\nRequirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.7/site-packages (from matplotlib->flax) (0.10.0)\nRequirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.7/site-packages (from matplotlib->flax) (2.8.0)\nCollecting etils\n  Downloading etils-0.9.0-py3-none-any.whl (140 kB)\n\u001b[K     |████████████████████████████████| 140 kB 65.7 MB/s eta 0:00:01\n\u001b[?25hCollecting cached_property\n  Downloading cached_property-1.5.2-py2.py3-none-any.whl (7.6 kB)\nRequirement already satisfied: pytest in /opt/conda/lib/python3.7/site-packages (from orbax->flax) (6.2.5)\nRequirement already satisfied: importlib_resources in /opt/conda/lib/python3.7/site-packages (from orbax->flax) (5.2.2)\nRequirement already satisfied: pytz>=2017.3 in /opt/conda/lib/python3.7/site-packages (from pandas->datasets) (2021.1)\nRequirement already satisfied: toml in /opt/conda/lib/python3.7/site-packages (from pytest->orbax->flax) (0.10.2)\nRequirement already satisfied: iniconfig in /opt/conda/lib/python3.7/site-packages (from pytest->orbax->flax) (1.1.1)\nRequirement already satisfied: py>=1.8.2 in /opt/conda/lib/python3.7/site-packages (from pytest->orbax->flax) (1.10.0)\nRequirement already satisfied: pluggy<2.0,>=0.12 in /opt/conda/lib/python3.7/site-packages (from pytest->orbax->flax) (1.0.0)\nRequirement already satisfied: click in /opt/conda/lib/python3.7/site-packages (from sacremoses->transformers) (8.0.1)\nRequirement already satisfied: joblib in /opt/conda/lib/python3.7/site-packages (from sacremoses->transformers) (1.0.1)\nBuilding wheels for collected packages: jax\n  Building wheel for jax (setup.py) ... \u001b[?25ldone\n\u001b[?25h  Created wheel for jax: filename=jax-0.3.25-py3-none-any.whl size=1308509 sha256=c410af118b84a05d1df582f5ba0f3cfdc01763cc7e50c5f01fb16c93a493ebf7\n  Stored in directory: /root/.cache/pip/wheels/a5/79/44/dbae745c48940291d6d4d157f7faa6f0b27f198781073d8c78\nSuccessfully built jax\nInstalling collected packages: typing-extensions, numpy, jax, tensorstore, etils, commonmark, chex, cached-property, rich, orbax, optax, flax\n  Attempting uninstall: typing-extensions\n    Found existing installation: typing-extensions 3.7.4.3\n    Uninstalling typing-extensions-3.7.4.3:\n      Successfully uninstalled typing-extensions-3.7.4.3\n  Attempting uninstall: numpy\n    Found existing installation: numpy 1.19.5\n    Uninstalling numpy-1.19.5:\n      Successfully uninstalled numpy-1.19.5\n  Attempting uninstall: jax\n    Found existing installation: jax 0.2.19\n    Uninstalling jax-0.2.19:\n      Successfully uninstalled jax-0.2.19\n\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\ncaip-notebooks-serverextension 1.0.0 requires google-cloud-bigquery-storage, which is not installed.\nyellowbrick 1.3.post1 requires numpy<1.20,>=1.16.0, but you have numpy 1.21.6 which is incompatible.\ntensorflow 2.4.1 requires numpy~=1.19.2, but you have numpy 1.21.6 which is incompatible.\ntensorflow 2.4.1 requires typing-extensions~=3.7.4, but you have typing-extensions 4.4.0 which is incompatible.\npdpbox 0.2.1 requires matplotlib==3.1.1, but you have matplotlib 3.4.3 which is incompatible.\nmatrixprofile 1.1.10 requires protobuf==3.11.2, but you have protobuf 3.18.0 which is incompatible.\nimbalanced-learn 0.8.0 requires scikit-learn>=0.24, but you have scikit-learn 0.23.2 which is incompatible.\narviz 0.11.2 requires typing-extensions<4,>=3.7.4.3, but you have typing-extensions 4.4.0 which is incompatible.\naiobotocore 1.4.1 requires botocore<1.20.107,>=1.20.106, but you have botocore 1.21.44 which is incompatible.\u001b[0m\nSuccessfully installed cached-property-1.5.2 chex-0.1.5 commonmark-0.9.1 etils-0.9.0 flax-0.6.3 jax-0.3.25 numpy-1.21.6 optax-0.1.4 orbax-0.0.23 rich-13.0.0 tensorstore-0.1.28 typing-extensions-4.4.0\n\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\n","output_type":"stream"}]},{"cell_type":"code","source":"from kaggle_datasets import KaggleDatasets\nimport tensorflow as tf \nimport tensorflow_datasets as tfds\nimport jax.random as random\n\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('tpu-getting-started')\nprint(GCS_DS_PATH) # what do gcs paths look like?\n\nIMAGE_SIZE = [192, 192]\nGCS_PATH = GCS_DS_PATH + '/tfrecords-jpeg-192x192'\nAUTO = tf.data.experimental.AUTOTUNE\n\nTRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec')\nVALIDATION_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/val/*.tfrec')\nTEST_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/test/*.tfrec') \n\n","metadata":{"execution":{"iopub.status.busy":"2023-01-04T00:43:49.910511Z","iopub.execute_input":"2023-01-04T00:43:49.910876Z","iopub.status.idle":"2023-01-04T00:43:59.780707Z","shell.execute_reply.started":"2023-01-04T00:43:49.910777Z","shell.execute_reply":"2023-01-04T00:43:59.779699Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stderr","text":"2023-01-04 00:43:50.865451: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/conda/lib\n2023-01-04 00:43:50.865594: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n","output_type":"stream"},{"name":"stdout","text":"gs://kds-e46cf9604ba9eb344e23f7d9394e7721c23e6e2d058fc48f3020226f\n","output_type":"stream"},{"name":"stderr","text":"2023-01-04 00:43:59.535950: I tensorflow/core/platform/cloud/google_auth_provider.cc:180] Attempting an empty bearer token since no token was retrieved from files, and GCE metadata check was skipped.\n2023-01-04 00:43:59.629567: I tensorflow/core/platform/cloud/google_auth_provider.cc:180] Attempting an empty bearer token since no token was retrieved from files, and GCE metadata check was skipped.\n2023-01-04 00:43:59.703436: I tensorflow/core/platform/cloud/google_auth_provider.cc:180] Attempting an empty bearer token since no token was retrieved from files, and GCE metadata check was skipped.\n","output_type":"stream"}]},{"cell_type":"code","source":"CLASSES = ['pink primrose',    'hard-leaved pocket orchid', 'canterbury bells', 'sweet pea',     'wild geranium',     'tiger lily',           'moon orchid',              'bird of paradise', 'monkshood',        'globe thistle',         # 00 - 09\n           'snapdragon',       \"colt's foot\",               'king protea',      'spear thistle', 'yellow iris',       'globe-flower',         'purple coneflower',        'peruvian lily',    'balloon flower',   'giant white arum lily', # 10 - 19\n           'fire lily',        'pincushion flower',         'fritillary',       'red ginger',    'grape hyacinth',    'corn poppy',           'prince of wales feathers', 'stemless gentian', 'artichoke',        'sweet william',         # 20 - 29\n           'carnation',        'garden phlox',              'love in the mist', 'cosmos',        'alpine sea holly',  'ruby-lipped cattleya', 'cape flower',              'great masterwort', 'siam tulip',       'lenten rose',           # 30 - 39\n           'barberton daisy',  'daffodil',                  'sword lily',       'poinsettia',    'bolero deep blue',  'wallflower',           'marigold',                 'buttercup',        'daisy',            'common dandelion',      # 40 - 49\n           'petunia',          'wild pansy',                'primula',          'sunflower',     'lilac hibiscus',    'bishop of llandaff',   'gaura',                    'geranium',         'orange dahlia',    'pink-yellow dahlia',    # 50 - 59\n           'cautleya spicata', 'japanese anemone',          'black-eyed susan', 'silverbush',    'californian poppy', 'osteospermum',         'spring crocus',            'iris',             'windflower',       'tree poppy',            # 60 - 69\n           'gazania',          'azalea',                    'water lily',       'rose',          'thorn apple',       'morning glory',        'passion flower',           'lotus',            'toad lily',        'anthurium',             # 70 - 79\n           'frangipani',       'clematis',                  'hibiscus',         'columbine',     'desert-rose',       'tree mallow',          'magnolia',                 'cyclamen ',        'watercress',       'canna lily',            # 80 - 89\n           'hippeastrum ',     'bee balm',                  'pink quill',       'foxglove',      'bougainvillea',     'camellia',             'mallow',                   'mexican petunia',  'bromelia',         'blanket flower',        # 90 - 99\n           'trumpet creeper',  'blackberry lily',           'common tulip',     'wild rose']                                                                                                                                               # 100 - 102\n\n\ndef decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n    image = tf.reshape(image, [*IMAGE_SIZE, 3]) # explicit size needed for TPU\n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label # returns a dataset of (image, label) pairs\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"id\": tf.io.FixedLenFeature([], tf.string),  # shape [] means single element\n        # class is missing, this competitions's challenge is to predict flower classes for the test dataset\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['id']\n    return image, idnum # returns a dataset of image(s)\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    # Read from TFRecords. For optimal performance, reading from multiple files at once and\n    # disregarding data order. Order does not matter since we will be shuffling the data anyway.\n\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls=AUTO)\n    # returns a dataset of (image, label) pairs if labeled=True or (image, id) pairs if labeled=False\n    return tfds.as_numpy(dataset.batch(64))\n","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:53.055818Z","iopub.execute_input":"2023-01-03T22:14:53.056126Z","iopub.status.idle":"2023-01-03T22:14:53.075511Z","shell.execute_reply.started":"2023-01-03T22:14:53.056092Z","shell.execute_reply":"2023-01-03T22:14:53.074426Z"},"trusted":true},"execution_count":4,"outputs":[]},{"cell_type":"code","source":"train_dataset = load_dataset(TRAINING_FILENAMES)\ntest_dataset = load_dataset(TEST_FILENAMES)\nval_dataset = load_dataset(VALIDATION_FILENAMES)","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:53.076998Z","iopub.execute_input":"2023-01-03T22:14:53.077253Z","iopub.status.idle":"2023-01-03T22:14:53.616131Z","shell.execute_reply.started":"2023-01-03T22:14:53.077226Z","shell.execute_reply":"2023-01-03T22:14:53.615222Z"},"trusted":true},"execution_count":5,"outputs":[{"name":"stderr","text":"2023-01-03 22:14:53.321325: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/conda/lib\n2023-01-03 22:14:53.321400: W tensorflow/stream_executor/cuda/cuda_driver.cc:326] failed call to cuInit: UNKNOWN ERROR (303)\n","output_type":"stream"}]},{"cell_type":"code","source":"import jax.numpy as jnp\nimport flax.linen as nn\nimport jax\nclass CNN(nn.Module):\n    @nn.compact\n    def __call__(self, x):\n        x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n        x = nn.relu(x)\n        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n        x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n        x = nn.relu(x)\n        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n        x = x.reshape((x.shape[0], -1))  \n        x = nn.Dense(features=256)(x)\n        x = nn.relu(x)\n        x = nn.Dense(features=len(CLASSES))(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:53.619983Z","iopub.execute_input":"2023-01-03T22:14:53.620269Z","iopub.status.idle":"2023-01-03T22:14:54.068612Z","shell.execute_reply.started":"2023-01-03T22:14:53.62024Z","shell.execute_reply":"2023-01-03T22:14:54.06707Z"},"trusted":true},"execution_count":6,"outputs":[{"traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_20/2621179609.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mflax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinen\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mCNN\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mModule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompact\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/opt/conda/lib/python3.7/site-packages/flax/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\"\"\"Flax API.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0;34m.\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     19\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;34m.\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mjax_utils\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;34m.\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mlinen\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/opt/conda/lib/python3.7/site-packages/flax/core/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     26\u001b[0m )\n\u001b[1;32m     27\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m from .scope import (\n\u001b[0m\u001b[1;32m     29\u001b[0m   \u001b[0mScope\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mScope\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m   \u001b[0mArray\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mArray\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/opt/conda/lib/python3.7/site-packages/flax/core/scope.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mflax\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mflax\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mflax\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mstruct\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     29\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mflax\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtraceback_util\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;34m.\u001b[0m\u001b[0mfrozen_dict\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mfreeze\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/opt/conda/lib/python3.7/site-packages/flax/struct.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mtyping_extensions\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mdataclass_transform\u001b[0m  \u001b[0;31m# pytype: disable=not-supported-yet\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mImportError\u001b[0m: cannot import name 'dataclass_transform' from 'typing_extensions' (/opt/conda/lib/python3.7/site-packages/typing_extensions.py)"],"ename":"ImportError","evalue":"cannot import name 'dataclass_transform' from 'typing_extensions' (/opt/conda/lib/python3.7/site-packages/typing_extensions.py)","output_type":"error"}]},{"cell_type":"code","source":"import optax\nfrom flax.training import train_state  # Useful dataclass to keep train state\n\ndef create_train_state(rng, learning_rate):\n  cnn = CNN()\n  params = cnn.init(rng, jnp.ones([64, 192, 192, 3]))['params'] # initialize parameters by passing a template image\n  tx = optax.adam(learning_rate)\n  return train_state.TrainState.create(\n      apply_fn=cnn.apply, params=params, tx=tx)\n\nrng = jax.random.PRNGKey(0)\nrng, init_rng = jax.random.split(rng)\n\nstate = create_train_state(init_rng, 0.05)\ndel init_rng  # Must not be used anymore.","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:54.069749Z","iopub.status.idle":"2023-01-03T22:14:54.070299Z","shell.execute_reply.started":"2023-01-03T22:14:54.070019Z","shell.execute_reply":"2023-01-03T22:14:54.070048Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def cross_entropy_loss(*, logits, labels):\n  labels_onehot = jax.nn.one_hot(labels, num_classes=len(CLASSES))\n  return optax.softmax_cross_entropy(logits=logits, labels=labels_onehot).mean()\n\ndef compute_metrics(*, logits, labels):\n  loss = cross_entropy_loss(logits=logits, labels=labels)\n  accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n  metrics = {\n      'loss': loss,\n      'accuracy': accuracy,\n  }\n  return metrics\n\nfrom tqdm import tqdm\n\n@jax.jit\ndef train_step(state, image_batch, label_batch):\n    def cel(params):\n        logits = CNN().apply({'params': params}, image_batch)\n        loss = cross_entropy_loss(logits=logits, labels=label_batch)\n        return loss, logits\n    grad_fn = jax.grad(cel, has_aux=True)\n    grads, logits = grad_fn(state.params)\n    state = state.apply_gradients(grads=grads)\n    metrics = compute_metrics(logits=logits, labels=label_batch)\n    return state, metrics\n        \ndef train_epoch(state, train_dataset, epoch):\n    print(\"Training epoch.\")\n    batch_metrics = []\n    for i in tqdm(train_dataset):\n        state, metrics = train_step(state, i[0], i[1])\n        batch_metrics.append(metrics)\n\n    print(batch_metrics)\n    # compute mean of metrics across each batch in epoch.\n    batch_metrics_np = jax.device_get(batch_metrics)\n    epoch_metrics_np = {\n      k: np.mean([metrics[k] for metrics in batch_metrics_np])\n      for k in batch_metrics_np[0]} # jnp.mean does not work on lists\n\n    print('train epoch: %d, loss: %.4f, accuracy: %.2f' % (\n      epoch, epoch_metrics_np['loss'], epoch_metrics_np['accuracy'] * 100))\n\n    return state\n","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:54.072053Z","iopub.status.idle":"2023-01-03T22:14:54.072705Z","shell.execute_reply.started":"2023-01-03T22:14:54.072291Z","shell.execute_reply":"2023-01-03T22:14:54.072331Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i in range(5):\n    state = train_epoch(state, train_dataset, i)\n","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:54.07439Z","iopub.status.idle":"2023-01-03T22:14:54.074923Z","shell.execute_reply.started":"2023-01-03T22:14:54.074646Z","shell.execute_reply":"2023-01-03T22:14:54.074671Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@jax.jit\ndef eval_step(params, batch):\n  logits = CNN().apply({'params': params}, batch['image'])\n  return compute_metrics(logits=logits, labels=batch['label'])\n","metadata":{"execution":{"iopub.status.busy":"2023-01-03T22:14:54.076417Z","iopub.status.idle":"2023-01-03T22:14:54.076868Z","shell.execute_reply.started":"2023-01-03T22:14:54.076666Z","shell.execute_reply":"2023-01-03T22:14:54.076693Z"},"trusted":true},"execution_count":null,"outputs":[]}]}