{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[],"dockerImageVersionId":30733,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\n# os.environ[\"KERAS_BACKEND\"] = \"jax\"\n\nimport gc\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\n# import jax\n# import keras\n\nfrom sklearn import metrics\n\nfrom tqdm.notebook import tqdm\n\nprint(tf.__version__)\n# print(jax.__version__)\n\ndef _parse_function(example_proto):\n    feature_description = {\n        'x': tf.io.FixedLenFeature([556], tf.float32),\n        'targets': tf.io.FixedLenFeature([368], tf.float32)\n    }\n    e = tf.io.parse_single_example(example_proto, feature_description)\n    return e['x'], e['targets']\n\nDATA = \"/kaggle/input/leap-atmospheric-physics-ai-climsim\"\nDATA_TFREC = \"/kaggle/input/leap-train-tfrecords\"\n\ntrain_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(90)]\nvalid_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(90, 101)]\n\nBATCH_SIZE = 2048\n\ntrain_options = tf.data.Options()\ntrain_options.deterministic = True\n\nds_train = (\n    tf.data.Dataset.from_tensor_slices(train_files)\n    .with_options(train_options)\n    .shuffle(100)\n    .interleave(\n        lambda file: tf.data.TFRecordDataset(file).map(_parse_function, num_parallel_calls=tf.data.AUTOTUNE),\n        num_parallel_calls=tf.data.AUTOTUNE,\n        cycle_length=10,\n        block_length=1000,\n        deterministic=True\n    )\n    .shuffle(4 * BATCH_SIZE)\n    .batch(BATCH_SIZE)\n    #.repeat()\n    .prefetch(tf.data.AUTOTUNE)\n)\n\nds_valid = (\n    tf.data.TFRecordDataset(valid_files)\n    .map(_parse_function)\n    .batch(BATCH_SIZE)\n    #.repeat()\n    .prefetch(tf.data.AUTOTUNE)\n)\n\nnorm_x = tf.keras.layers.Normalization()\nnorm_x.adapt(ds_train.map(lambda x, y: x).take(5_000_000//BATCH_SIZE))\n\nmean_x = norm_x.mean\nstdd_x = tf.keras.ops.maximum(1e-15, norm_x.variance ** 0.5)\n\nnorm_y = tf.keras.layers.Normalization()\nnorm_y.adapt(ds_train.map(lambda x, y: y).take(5_000_000//BATCH_SIZE))\n\nmean_y = norm_y.mean\nstdd_y = tf.keras.ops.maximum(1e-10, norm_y.variance ** 0.5)\nnp.save('mean_y.npy', mean_y)\nnp.save('stdd_y.npy', stdd_y)\nnp.save('mean_x.npy', mean_x)\nnp.save('stdd_x.npy', stdd_x)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T01:39:28.962607Z","iopub.execute_input":"2024-07-09T01:39:28.963206Z","iopub.status.idle":"2024-07-09T01:42:48.506674Z","shell.execute_reply.started":"2024-07-09T01:39:28.963172Z","shell.execute_reply":"2024-07-09T01:42:48.505843Z"},"trusted":true},"execution_count":null,"outputs":[]}]}