{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"TPU = 1\nVAL = 0\nLOAD = 0","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:23.836132Z","iopub.execute_input":"2023-08-20T06:39:23.836615Z","iopub.status.idle":"2023-08-20T06:39:23.842436Z","shell.execute_reply.started":"2023-08-20T06:39:23.83658Z","shell.execute_reply":"2023-08-20T06:39:23.841479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import time\n!pip install Levenshtein\nimport Levenshtein as lev\n","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:23.848957Z","iopub.execute_input":"2023-08-20T06:39:23.849634Z","iopub.status.idle":"2023-08-20T06:39:37.759029Z","shell.execute_reply.started":"2023-08-20T06:39:23.849599Z","shell.execute_reply":"2023-08-20T06:39:37.757452Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if 'tpu_init' not in globals():\n    import pandas as pd\n    import numpy as np\n    import tensorflow as tf\n    from tensorflow import keras\n    from tensorflow.keras import layers\n    import json\n    import tensorflow_addons as tfa","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:37.762838Z","iopub.execute_input":"2023-08-20T06:39:37.763239Z","iopub.status.idle":"2023-08-20T06:39:37.77114Z","shell.execute_reply.started":"2023-08-20T06:39:37.763199Z","shell.execute_reply":"2023-08-20T06:39:37.769465Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if 'tpu_init' not in globals():\n    tpu_init = True\n    if TPU:\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()\n        strategy = tf.distribute.experimental.TPUStrategy(tpu)\n    else:\n        ngpu = len(tf.config.experimental.list_physical_devices('GPU'))\n        if ngpu>1:\n            print(\"Using multi GPU\")\n            strategy = tf.distribute.MirroredStrategy()\n        elif ngpu==1:\n            print(\"Using single GPU\")\n            strategy = tf.distribute.get_strategy()\n            BATCH_SIZE = 64\n        else:\n            print(\"Using CPU\")\n            strategy = tf.distribute.get_strategy()\n            BATCH_SIZE = 2 ","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:37.77296Z","iopub.execute_input":"2023-08-20T06:39:37.773348Z","iopub.status.idle":"2023-08-20T06:39:37.786535Z","shell.execute_reply.started":"2023-08-20T06:39:37.773319Z","shell.execute_reply":"2023-08-20T06:39:37.785482Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import tensorflow.keras.mixed_precision as mixed_precision\n# try:\n#     policy = mixed_precision.Policy('mixed_bfloat16')\n#     mixed_precision.set_global_policy(policy)\n#     print('mixed_bfloat16')\n# except:\n#     policy = mixed_precision.Policy('mixed_float16')\n#     mixed_precision.set_global_policy(policy)\n#     print('mixed_float16')","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:37.788271Z","iopub.execute_input":"2023-08-20T06:39:37.788945Z","iopub.status.idle":"2023-08-20T06:39:37.806169Z","shell.execute_reply.started":"2023-08-20T06:39:37.788912Z","shell.execute_reply":"2023-08-20T06:39:37.805105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from shutil import copyfile\n# copy our file into the working directory (make sure it has .py suffix)\ncopyfile(src = \"/kaggle/input/ctc-tpu/CTC_TPU.py\", dst = \"/kaggle/working//CTC_TPU.py\")\n\n# import all our functions\nfrom CTC_TPU import classic_ctc_loss","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:37.809505Z","iopub.execute_input":"2023-08-20T06:39:37.810088Z","iopub.status.idle":"2023-08-20T06:39:37.826306Z","shell.execute_reply.started":"2023-08-20T06:39:37.810057Z","shell.execute_reply":"2023-08-20T06:39:37.824711Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"COMPUTE_TYPE = 'float32'\nUSE_Z_AXIS = True\nBATCH_SIZE = 64\nFRAME_LEN = 220  # 128\nPHRASE_MAX_LEN = 32 # + 2\nseed = 42\ntf.random.set_seed(seed)\nnp.random.seed(seed)\nDROP_OUT = 0.1\n# NORM = tf.keras.layers.BatchNormalization()\n# NORM = tf.keras.layers.LayerNormalization()\nMASKING = False\nif TPU:\n    BATCH_SIZE = 25 * strategy.num_replicas_in_sync\n    print(BATCH_SIZE)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:37.828897Z","iopub.execute_input":"2023-08-20T06:39:37.829342Z","iopub.status.idle":"2023-08-20T06:39:37.884784Z","shell.execute_reply.started":"2023-08-20T06:39:37.829306Z","shell.execute_reply":"2023-08-20T06:39:37.883628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"drive_path = \"/kaggle/input/asl-fingerspelling\"\nwith open(f\"{drive_path}/character_to_prediction_index.json\", \"r\") as f:\n    char_to_num = json.load(f)\n\npad_token = \"P\"\npad_token_idx = 59\n\nchar_to_num[pad_token] = pad_token_idx\n\nCLASS_NUM = len(char_to_num)\nprint(CLASS_NUM)\n\nnum_to_char = {j: i for i, j in char_to_num.items()}\n\ndf = pd.read_csv(f\"/kaggle/input/asl-fingerspelling/train.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:37.886492Z","iopub.execute_input":"2023-08-20T06:39:37.8871Z","iopub.status.idle":"2023-08-20T06:39:38.029318Z","shell.execute_reply.started":"2023-08-20T06:39:37.887065Z","shell.execute_reply":"2023-08-20T06:39:38.028296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NOSE=[\n    1,2,98,327\n]\nLNOSE = [98]\nRNOSE = [327]\nLIP = [ 0, \n    61, 185, 40, 39, 37, 267, 269, 270, 409,\n    291, 146, 91, 181, 84, 17, 314, 405, 321, 375,\n    78, 191, 80, 81, 82, 13, 312, 311, 310, 415,\n    95, 88, 178, 87, 14, 317, 402, 318, 324, 308,\n]\nLLIP = [84,181,91,146,61,185,40,39,37,87,178,88,95,78,191,80,81,82]\nRLIP = [314,405,321,375,291,409,270,269,267,317,402,318,324,308,415,310,311,312]\n\nPOSE = [490,491,492,493,494,495,496,497,498,499,500,501,502,503,504,505,506,507,508,509,510,511,512,513,514,515,516,517,518,519,520,521,522]\nPOSE = [i-1 for i in POSE]\n\nLPOSE = [494, 495, 497, 499, 501, 503, 505, 507, 509, 511, 513, 515, 517, 519]\nLPOSE = [i-1 for i in LPOSE]\n\nRPOSE = [491, 492, 496, 498, 500, 502, 504, 506, 508, 510, 512, 514, 516, 518]\nRPOSE = [i-1 for i in RPOSE]\n\n\n\nREYE = [\n    33, 7, 163, 144, 145, 153, 154, 155, 133,\n    246, 161, 160, 159, 158, 157, 173,\n]\nLEYE = [\n    263, 249, 390, 373, 374, 380, 381, 382, 362,\n    466, 388, 387, 386, 385, 384, 398,\n]\n\nLHAND = np.arange(468, 489).tolist()\nRHAND = np.arange(522, 543).tolist()\n\nPOINT_LANDMARKS = LIP + LHAND + RHAND + NOSE + REYE + LEYE  + POSE # use POSE\n\nwith open('/kaggle/input/training-args/training_args.json','r') as f:\n    get_cols = json.load(f)['columns']\n\ndef idx_to_cols(idx_array):\n    sample_cols = [get_cols[i+1] for i in idx_array]\n    # sample_cols = data.iloc[:,[i+1 for i in idx_array]].columns\n    \n    patterns = [s[1:] for s in sample_cols]\n    if USE_Z_AXIS:\n        matched_cols = [s for s in get_cols if s[1:] in patterns]\n    else:\n        matched_cols = [s for s in get_cols if s[1:] in patterns and s[0] != 'z']\n    return matched_cols\n\nSEL_COLS = idx_to_cols(POINT_LANDMARKS)\nprint(len(SEL_COLS))\nprint(len(set(SEL_COLS)))\ndef get_index(arr):\n    cols = idx_to_cols(arr)\n    idx = [SEL_COLS.index(c) for c in cols]\n    return idx","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:38.03069Z","iopub.execute_input":"2023-08-20T06:39:38.031772Z","iopub.status.idle":"2023-08-20T06:39:38.065509Z","shell.execute_reply.started":"2023-08-20T06:39:38.031724Z","shell.execute_reply":"2023-08-20T06:39:38.063852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_DIM = 3 if USE_Z_AXIS else 2\nHAND_NUMS = len(LHAND)+ len(RHAND)\nFACE_NUMS = len(LIP) + len(REYE) + len(LEYE) + len(NOSE)\nPOSE_NUMS = len(POSE)\nprint(HAND_NUMS, FACE_NUMS, POSE_NUMS)\nprint(HAND_NUMS+FACE_NUMS+POSE_NUMS)\n\nLIP_IDX = get_index(LIP)\nLHAND_IDX = get_index(LHAND)\nRHAND_IDX = get_index(RHAND)\nNOSE_IDX = get_index(NOSE)\nREYE_IDX = get_index(REYE)\nLEYE_IDX = get_index(LEYE)\n\nLLIP_IDX = get_index(LLIP)\nRLIP_IDX = get_index(RLIP)\nLNOSE_IDX = get_index(LNOSE)\nRNOSE_IDX = get_index(RNOSE)\n\nPOSE_IDX = get_index(POSE)\nLPOSE_IDX = get_index(LPOSE)\nRPOSE_IDX = get_index(RPOSE)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:38.067571Z","iopub.execute_input":"2023-08-20T06:39:38.068007Z","iopub.status.idle":"2023-08-20T06:39:38.093928Z","shell.execute_reply.started":"2023-08-20T06:39:38.067971Z","shell.execute_reply":"2023-08-20T06:39:38.09249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NAN_FILL_VALUE = tf.constant(0, dtype=tf.float32)\nNAN_VALUE = tf.constant(np.nan, dtype=tf.float32)\nPADDING_MASKING_VALUE = tf.constant(-100, dtype=tf.float32)\n\ndef resize_pad(x):\n    if tf.shape(x)[0] < FRAME_LEN:\n        # if MASKING:\n        #     if tf.shape(x)[0] < 32:\n        #         x = tf.pad(\n        #             x,\n        #             ([[0, (32 - tf.shape(x)[0])], [0, 0], [0, 0]]),\n        #             constant_values = PADDING_MASKING_VALUE,\n        #         )\n        x = tf.pad(\n            x,\n            ([[0, (FRAME_LEN - tf.shape(x)[0])], [0, 0], [0, 0]]),\n            constant_values = NAN_FILL_VALUE,\n        )\n    else:\n        x = tf.image.resize(x, (FRAME_LEN, tf.shape(x)[1]), \"nearest\")\n    return x\n\ndef tf_nan_mean(x, axis=0, keepdims=False):\n    return tf.reduce_sum(\n        tf.where(tf.math.is_nan(x), tf.zeros_like(x), x), axis=axis, keepdims=keepdims\n    ) / tf.reduce_sum(\n        tf.where(tf.math.is_nan(x), tf.zeros_like(x), tf.ones_like(x)),\n        axis=axis,\n        keepdims=keepdims,\n    )\n\n\ndef tf_nan_std(x, center=None, axis=0, keepdims=False):\n    if center is None:\n        center = tf_nan_mean(x, axis=axis, keepdims=True)\n    d = x - center\n    return tf.math.sqrt(tf_nan_mean(d * d, axis=axis, keepdims=keepdims))\n\n\ndef self_norm(x):\n    # input batch, 21 + 39 + 33, 2\n    mean_no_nan = tf_nan_mean(x, axis=[0,1],keepdims=True)\n    std_no_nan = tf_nan_std(x, center=mean_no_nan, axis=[0,1],keepdims=True)\n    x = (x - mean_no_nan) / (std_no_nan)\n    return x\n\ndef global_norm(x):\n    # face = x[:,:len(LIP_IDX),:]\n    pose = x[:,-len(POSE_IDX):,:]\n    mean_no_nan = tf_nan_mean(pose, axis=[0,1],keepdims=True)\n    std_no_nan = tf_nan_std(x, center=mean_no_nan, axis=[0,1],keepdims=True)\n    x = (x - mean_no_nan) / (std_no_nan) / 3\n    return x\n\n\ndef split_data(x):\n    # POINT_LANDMARKS = LIP + LHAND + RHAND + NOSE + REYE + LEYE\n\n    lip = tf.gather(x, LIP_IDX, axis=1)\n    if USE_Z_AXIS:\n        lip_x = lip[:, 0 * (len(LIP_IDX) // 3) : 1 * (len(LIP_IDX) // 3)]\n        lip_y = lip[:, 1 * (len(LIP_IDX) // 3) : 2 * (len(LIP_IDX) // 3)]\n        lip_z = lip[:, 2 * (len(LIP_IDX) // 3) : 3 * (len(LIP_IDX) // 3)]\n        lip = tf.concat(\n            [lip_x[..., tf.newaxis], lip_y[..., tf.newaxis], lip_z[..., tf.newaxis]],\n            axis=-1,\n        )\n    else:\n        lip_x = lip[:, 0 * (len(LIP_IDX) // 2) : 1 * (len(LIP_IDX) // 2)]\n        lip_y = lip[:, 1 * (len(LIP_IDX) // 2) : 2 * (len(LIP_IDX) // 2)]\n\n    lhand = tf.gather(x, LHAND_IDX, axis=1)\n    rhand = tf.gather(x, RHAND_IDX, axis=1)\n    if USE_Z_AXIS:\n        lhand_x = lhand[:, 0 * (len(LHAND_IDX) // 3) : 1 * (len(LHAND_IDX) // 3)]\n        lhand_y = lhand[:, 1 * (len(LHAND_IDX) // 3) : 2 * (len(LHAND_IDX) // 3)]\n        lhand_z = lhand[:, 2 * (len(LHAND_IDX) // 3) : 3 * (len(LHAND_IDX) // 3)]\n        lhand = tf.concat(\n            [lhand_x[..., tf.newaxis], lhand_y[..., tf.newaxis], lhand_z[..., tf.newaxis]],\n            axis=-1,\n        )\n        rhand_x = rhand[:, 0 * (len(RHAND_IDX) // 3) : 1 * (len(RHAND_IDX) // 3)]\n        rhand_y = rhand[:, 1 * (len(RHAND_IDX) // 3) : 2 * (len(RHAND_IDX) // 3)]\n        rhand_z = rhand[:, 2 * (len(RHAND_IDX) // 3) : 3 * (len(RHAND_IDX) // 3)]\n        rhand = tf.concat(\n            [rhand_x[..., tf.newaxis], rhand_y[..., tf.newaxis], rhand_z[..., tf.newaxis]],\n            axis=-1,\n        )\n    else:\n        lhand_x = lhand[:, 0 * (len(LHAND_IDX) // 2) : 1 * (len(LHAND_IDX) // 2)]\n        lhand_y = lhand[:, 1 * (len(LHAND_IDX) // 2) : 2 * (len(LHAND_IDX) // 2)]\n        rhand_x = rhand[:, 0 * (len(RHAND_IDX) // 2) : 1 * (len(RHAND_IDX) // 2)]\n        rhand_y = rhand[:, 1 * (len(RHAND_IDX) // 2) : 2 * (len(RHAND_IDX) // 2)]\n        lhand = tf.concat([lhand_x[..., tf.newaxis], lhand_y[..., tf.newaxis]], axis=-1)\n        rhand = tf.concat([rhand_x[..., tf.newaxis], rhand_y[..., tf.newaxis]], axis=-1)\n\n    nose = tf.gather(x, NOSE_IDX, axis=1)\n    if USE_Z_AXIS:\n        nose_x = nose[:, 0 * (len(NOSE_IDX) // 3) : 1 * (len(NOSE_IDX) // 3)]\n        nose_y = nose[:, 1 * (len(NOSE_IDX) // 3) : 2 * (len(NOSE_IDX) // 3)]\n        nose_z = nose[:, 2 * (len(NOSE_IDX) // 3) : 3 * (len(NOSE_IDX) // 3)]\n        nose = tf.concat(\n            [nose_x[..., tf.newaxis], nose_y[..., tf.newaxis], nose_z[..., tf.newaxis]],axis=-1)\n    else:\n        nose_x = nose[:, 0 * (len(NOSE_IDX) // 2) : 1 * (len(NOSE_IDX) // 2)]\n        nose_y = nose[:, 1 * (len(NOSE_IDX) // 2) : 2 * (len(NOSE_IDX) // 2)]\n        nose = tf.concat([nose_x[..., tf.newaxis], nose_y[..., tf.newaxis]], axis=-1)\n\n    reye = tf.gather(x, REYE_IDX, axis=1)\n    leye = tf.gather(x, LEYE_IDX, axis=1)\n    if USE_Z_AXIS:\n        reye_x = reye[:, 0 * (len(REYE_IDX) // 3) : 1 * (len(REYE_IDX) // 3)]\n        reye_y = reye[:, 1 * (len(REYE_IDX) // 3) : 2 * (len(REYE_IDX) // 3)]\n        reye_z = reye[:, 2 * (len(REYE_IDX) // 3) : 3 * (len(REYE_IDX) // 3)]\n        reye = tf.concat(\n            [reye_x[..., tf.newaxis], reye_y[..., tf.newaxis], reye_z[..., tf.newaxis]],\n            axis=-1,\n        )\n        leye_x = leye[:, 0 * (len(LEYE_IDX) // 3) : 1 * (len(LEYE_IDX) // 3)]\n        leye_y = leye[:, 1 * (len(LEYE_IDX) // 3) : 2 * (len(LEYE_IDX) // 3)]\n        leye_z = leye[:, 2 * (len(LEYE_IDX) // 3) : 3 * (len(LEYE_IDX) // 3)]\n        leye = tf.concat(\n            [leye_x[..., tf.newaxis], leye_y[..., tf.newaxis], leye_z[..., tf.newaxis]],\n            axis=-1,\n        )\n    else:\n        reye_x = reye[:, 0 * (len(REYE_IDX) // 2) : 1 * (len(REYE_IDX) // 2)]\n        reye_y = reye[:, 1 * (len(REYE_IDX) // 2) : 2 * (len(REYE_IDX) // 2)]\n        leye_x = leye[:, 0 * (len(LEYE_IDX) // 2) : 1 * (len(LEYE_IDX) // 2)]\n        leye_y = leye[:, 1 * (len(LEYE_IDX) // 2) : 2 * (len(LEYE_IDX) // 2)]\n        reye = tf.concat([reye_x[..., tf.newaxis], reye_y[..., tf.newaxis]], axis=-1)\n        leye = tf.concat([leye_x[..., tf.newaxis], leye_y[..., tf.newaxis]], axis=-1)\n\n    face = tf.concat([lip,nose,reye,leye], axis=1)\n\n\n    lpose = tf.gather(x, LPOSE_IDX, axis=1)\n    rpose = tf.gather(x, RPOSE_IDX, axis=1)\n    if USE_Z_AXIS:\n        lpose_x = lpose[:, 0 * (len(LPOSE_IDX) // 3) : 1 * (len(LPOSE_IDX) // 3)]\n        lpose_y = lpose[:, 1 * (len(LPOSE_IDX) // 3) : 2 * (len(LPOSE_IDX) // 3)]\n        lpose_z = lpose[:, 2 * (len(LPOSE_IDX) // 3) : 3 * (len(LPOSE_IDX) // 3)]\n        lpose = tf.concat([lpose_x[..., tf.newaxis], lpose_y[..., tf.newaxis], lpose_z[..., tf.newaxis]], axis=-1)\n        rpose_x = rpose[:, 0 * (len(RPOSE_IDX) // 3) : 1 * (len(RPOSE_IDX) // 3)]\n        rpose_y = rpose[:, 1 * (len(RPOSE_IDX) // 3) : 2 * (len(RPOSE_IDX) // 3)]\n        rpose_z = rpose[:, 2 * (len(RPOSE_IDX) // 3) : 3 * (len(RPOSE_IDX) // 3)]\n        rpose = tf.concat([rpose_x[..., tf.newaxis], rpose_y[..., tf.newaxis], rpose_z[..., tf.newaxis]], axis=-1)\n\n    pose = tf.concat([lpose,rpose], axis=1)\n\n    x = tf.concat([face[:,:len(LIP)],lhand,rhand,face[:,len(LIP):],pose], axis=1)\n    return x\n\n\n\ndef spatial_random_rotation_for_finger(\n    xyz,\n    degree=(-10, 10),\n):\n    # use first position for rotation center\n    if USE_Z_AXIS:\n        xy = xyz[:, :, 0:2]\n    else:\n        xy = xyz\n    center = xy[:,0:1,:]\n    # center = tf.reduce_mean(xy, axis=[0,1])\n    if degree is not None:\n        xy -= center\n        degree = tf.random.uniform(shape=[], minval=degree[0], maxval=degree[1], dtype=tf.float32)\n        radian = degree / 180 * np.pi\n        c = tf.math.cos(radian)\n        s = tf.math.sin(radian)\n        rotate_mat = tf.identity(\n            [\n                [c, s],\n                [-s, c],\n            ]\n        )\n        xy = xy @ rotate_mat\n        xy = xy + center\n    if USE_Z_AXIS:\n        return tf.concat([xy, xyz[:, :, 2:3]], axis=-1)\n    else:\n        return xy\n\ndef spatial_random_scale_for_finger(\n    xyz,\n    scale=(0.9, 1.1),\n):\n    # use first position for rotation center\n    if USE_Z_AXIS:\n        xy = xyz[:, :, 0:2]\n    else:\n        xy = xyz\n    center = xy[:,0:1,:]\n    # center = tf.reduce_mean(xy, axis=[0,1])\n    if scale is not None:\n        xy -= center\n        scale = tf.random.uniform(shape=[], minval=scale[0], maxval=scale[1], dtype=tf.float32)\n        xy = xy * scale\n        xy = xy + center\n    if USE_Z_AXIS:\n        return tf.concat([xy, xyz[:, :, 2:3]], axis=-1)\n    else:\n        return xy\n\ndef only_rotate(xyz,degree=(-15,15),shear = (-0.10,0.10),scale  = (0.75,1.5),\n                ):\n    scale = tf.random.uniform((),*scale)\n    xyz *= scale\n    if USE_Z_AXIS:\n        xy = xyz[:, :, 0:2]\n    else:\n        xy = xyz\n    center = tf_nan_mean(xy, axis=[0,1])\n    degree = tf.random.uniform((),*degree)\n    xy -= center\n\n    radian = degree/180*np.pi\n    c = tf.math.cos(radian)\n    s = tf.math.sin(radian)\n    rotate_mat = tf.identity([\n        [c,s],\n        [-s, c],\n    ])\n    xy = xy @ rotate_mat\n\n\n    shear_x = shear_y = tf.random.uniform((),*shear)\n    if tf.random.uniform(()) < 0.5:\n        shear_x = 0.\n    else:\n        shear_y = 0.\n    shear_mat = tf.identity([\n        [1.,shear_x],\n        [shear_y,1.]\n    ])\n    xy = xy @ shear_mat\n\n\n\n    xy = xy + center\n    if USE_Z_AXIS:\n        return tf.concat([xy, xyz[:, :, 2:3]], axis=-1)\n    else:\n        return xy\n\ndef inner_flip(x):\n    x,y,z = tf.unstack(x, axis=-1)\n    x = 2*tf_nan_mean(x, axis=[0,1], keepdims=True) -x\n    new_x = tf.stack([x,y,z], -1)\n    return new_x\n\ndef flip_lr(x):\n    x,y,z = tf.unstack(x, axis=-1)\n    face = tf.concat([x[:,:len(LIP)], x[:,len(LIP)+len(LHAND)+len(RHAND):-POSE_NUMS]], axis=1)\n    face_mean = tf_nan_mean(face, axis=[0,1], keepdims=True)\n    x = 2*face_mean -x\n    # x = 1 - x\n    new_x = tf.stack([x,y,z], -1)\n\n    face = tf.concat([new_x[:,:len(LIP),:], new_x[:,len(LIP)+len(LHAND)+len(RHAND):-POSE_NUMS,:]], axis=1)\n    leye = face[:,-len(LEYE):,:]\n    reye = face[:,-len(REYE)-len(LEYE):-len(LEYE),:]\n    face = tf.concat([face[:,:-len(LEYE)-len(REYE),:],leye,reye], axis=1)\n    rhand = new_x[:,len(LIP)+len(LHAND):len(LIP)+len(LHAND)+len(RHAND),:]\n    lhand = new_x[:,len(LIP):len(LIP)+len(LHAND),:]\n    pose = new_x[:,-POSE_NUMS:,:]\n    new_x = tf.concat([face[:,:len(LIP)],rhand, lhand,face[:,len(LIP):],pose], axis=1)\n\n    lip = new_x[:,:len(LIP),:]\n    lip = inner_flip(lip)\n    nose = new_x[:,len(LIP)+len(LHAND)+len(RHAND):len(LIP)+len(LHAND)+len(RHAND)+len(NOSE),:]\n    nose = inner_flip(nose)\n    pose = new_x[:,-POSE_NUMS:,:]\n    rpose = pose[:,POSE_NUMS//2:,:]\n    lpose = pose[:,:POSE_NUMS//2,:]\n    pose = tf.concat([rpose,lpose], axis=1)\n    new_x = tf.concat([lip, new_x[:,len(LIP):len(LIP)+len(LHAND),:], new_x[:,len(LIP)+len(LHAND):len(LIP)+len(LHAND)+len(RHAND),:], nose,new_x[:,len(LIP)+len(LHAND)+len(RHAND)+len(NOSE):-POSE_NUMS,:],pose], axis=1)\n    return new_x\n\n\n\ndef interp1d_(x, target_len, method=\"random\"):\n    target_len = tf.maximum(1, target_len)\n    if method == \"random\":\n        if tf.random.uniform(()) < 0.33:\n            x = tf.image.resize(x, (target_len, tf.shape(x)[1]), \"bilinear\")\n        else:\n            if tf.random.uniform(()) < 0.5:\n                x = tf.image.resize(x, (target_len, tf.shape(x)[1]), \"bicubic\")\n            else:\n                x = tf.image.resize(x, (target_len, tf.shape(x)[1]), \"nearest\")\n    else:\n        x = tf.image.resize(x, (target_len, tf.shape(x)[1]), method)\n    return x\n\n\ndef resample(x, rate):\n    # re-resample\n    rate = tf.random.uniform((),rate[0], rate[1])\n\n    length = tf.shape(x)[0]\n\n    new_size = tf.cast(rate * tf.cast(length, tf.float32), tf.int32)\n    return interp1d_(x, new_size)\n\ndef uniform_resample(x):\n    l = 10/tf.cast(tf.shape(x)[0], tf.float32)\n    r = FRAME_LEN/tf.cast(tf.shape(x)[0], tf.float32)\n    return resample(x, (l,r))\n\ndef resample_sub(x, rate=(0.5, 1.5)):\n\n    if tf.shape(x)[0] < 2:\n        x = resample(x, (1.0,5.0))\n        return x\n    if tf.random.uniform(()) < 0.5:\n        start = tf.random.uniform(shape=[], minval=0, maxval=tf.shape(x)[0]-1, dtype=tf.int32)\n        end = tf.random.uniform(shape=[], minval=start+1, maxval=tf.shape(x)[0], dtype=tf.int32)\n        if start > end:\n            start, end = end, start\n        x = tf.concat([x[:start], resample(x[start:end], rate), x[end:]], axis=0)\n    else:\n        x = resample(x, rate)\n    return x\n\n\ndef mask_along_axis(tensor, param_min, param_max, axis, mask_value=float('nan')):\n    tensor_shape = tf.shape(tensor)\n    dim_size = tensor_shape[axis]\n    min_mask_size = tf.cast(param_min * tf.cast(dim_size, tf.float32),tf.int32)\n    max_mask_size = tf.cast(param_max * tf.cast(dim_size, tf.float32), tf.int32)\n    mask_size = tf.cond(\n        tf.equal(min_mask_size, max_mask_size),\n        lambda: min_mask_size,\n        lambda: tf.random.uniform((), min_mask_size, max_mask_size+1, dtype=tf.int32)\n    )\n    mask_start = tf.random.uniform([], 0, dim_size - mask_size, tf.int32)\n    indices = tf.cast(tf.range(start=0, limit=dim_size, delta=1),tf.int32)\n    mask = tf.logical_or(indices < mask_start , indices >= (mask_start + mask_size))\n    if axis==1:\n        mask = tf.reshape(tf.cast(mask, tf.float32),(1,dim_size,1))\n    else:\n        mask = tf.reshape(tf.cast(mask, tf.float32),(dim_size,1,1))\n    masked_tensor = tf.where(mask==0,mask_value,tensor)\n    # masked_tensor = mask * tensor\n    return masked_tensor\n\ndef discrete_mask(tensor, param_min, param_max, axis,mask_value=float('nan')):\n    tensor_shape = tf.shape(tensor)\n    dim_size = tensor_shape[axis]\n\n    min_mask_size = tf.cast(param_min * tf.cast(dim_size, tf.float32), tf.int32)\n    max_mask_size = tf.cast(param_max * tf.cast(dim_size, tf.float32), tf.int32)\n\n    mask_size = tf.cond(\n        tf.equal(min_mask_size, max_mask_size),\n        lambda: min_mask_size,\n        lambda: tf.random.uniform((), min_mask_size, max_mask_size + 1, dtype=tf.int32)\n    )\n\n    mask_indices = tf.random.shuffle(tf.range(dim_size))[:mask_size]\n\n    mask = tf.scatter_nd(\n        tf.expand_dims(mask_indices, 1),\n        tf.ones(mask_size, dtype=tf.float32),\n        [dim_size]\n    )\n\n    if axis == 1:\n        mask = tf.reshape(mask, (1, dim_size, 1))\n    else:\n        mask = tf.reshape(mask, (dim_size, 1, 1))\n\n    masked_tensor = tf.where(mask==0,mask_value,tensor)\n    return masked_tensor\n\n\n\nfingers = [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16],[17,18,19,20]]\n\n\ndef random_rotate_fingers(x):\n    # x shape: (seq_len, 21, 2)\n    finger1 = tf.gather(x, fingers[0], axis=1)\n    finger2 = tf.gather(x, fingers[1], axis=1)\n    finger3 = tf.gather(x, fingers[2], axis=1)\n    finger4 = tf.gather(x, fingers[3], axis=1)\n    finger5 = tf.gather(x, fingers[4], axis=1)\n    finger1 = spatial_random_rotation_for_finger(finger1)\n    finger2 = spatial_random_rotation_for_finger(finger2)\n    finger3 = spatial_random_rotation_for_finger(finger3)\n    finger4 = spatial_random_rotation_for_finger(finger4)\n    finger5 = spatial_random_rotation_for_finger(finger5)\n    finger1 = spatial_random_scale_for_finger(finger1)\n    finger2 = spatial_random_scale_for_finger(finger2)\n    finger3 = spatial_random_scale_for_finger(finger3)\n    finger4 = spatial_random_scale_for_finger(finger4)\n    finger5 = spatial_random_scale_for_finger(finger5)\n    hand_root = tf.expand_dims(x[:,0,:], axis=1)\n    x = tf.concat([hand_root, finger1, finger2, finger3, finger4, finger5], axis=1)\n    return x\n\n\ndef spatial_mask(x, size=(0.1,0.4), mask_value=float('nan')):\n    x_min = tf.math.reduce_min(x[~tf.math.is_nan(x[...,0])])\n    x_max = tf.math.reduce_max(x[~tf.math.is_nan(x[...,0])])\n    y_min = tf.math.reduce_min(x[~tf.math.is_nan(x[...,1])])\n    y_max = tf.math.reduce_max(x[~tf.math.is_nan(x[...,1])])\n    mask_offset_x = tf.random.uniform((1,), x_min, x_max)\n    mask_offset_y = tf.random.uniform((1,), y_min, y_max)\n    mask_size_x = tf.random.uniform((1,), *size) * (x_max - x_min)\n    mask_size_y = tf.random.uniform((1,), *size) * (y_max - y_min)\n    mask_x = (mask_offset_x<x[...,0]) & (x[...,0] < mask_offset_x + mask_size_x)\n    mask_y = (mask_offset_y<x[...,1]) & (x[...,1] < mask_offset_y + mask_size_y)\n    mask = mask_x & mask_y\n    x = tf.where(mask[...,None], mask_value, x)\n    return x\n\n\ndef temporal_mask(x, size=(0.1,0.3), mask_value=float('nan')):\n    l = tf.shape(x)[0]\n    mask_size = tf.random.uniform((), *size)\n    mask_size = tf.cast(tf.cast(l, tf.float32) * mask_size, tf.int32)\n    mask_offset = tf.random.uniform((), 0, tf.clip_by_value(l-mask_size,1,l), dtype=tf.int32)\n    x = tf.tensor_scatter_nd_update(x,tf.range(mask_offset, mask_offset+mask_size)[...,None],tf.fill([mask_size,KPT_NUM//DATA_DIM,DATA_DIM],mask_value))\n    return x\n\n\ndef random_shift(x, shift_range=0.1):\n    # x shape TxKxdim\n    shift_var = tf.random.uniform((1,1,3),-shift_range,shift_range)\n    shift_var = tf.tile(shift_var, (tf.shape(x)[0], tf.shape(x)[1], 1))\n    x = x + shift_var\n    return x\n\ndef rotate_partial(x,partial=False):\n    def apply(x):\n        face = tf.concat([x[:,:len(LIP),:], x[:,len(LIP)+len(LHAND)+len(RHAND):-POSE_NUMS,:]], axis=1)\n        face = only_rotate(face,degree=(-30,30))\n        x = tf.concat([face[:,:len(LIP),:],x[:,len(LIP):len(LIP)+len(LHAND)+len(RHAND),:], face[:,len(LIP):,:],x[:,-POSE_NUMS:,:]], axis=1)\n        lhand = x[:,len(LIP):len(LIP)+len(LHAND),:]\n        lhand = only_rotate(lhand,degree=(-30,30))\n        x = tf.concat([x[:,:len(LIP),:], lhand, x[:,len(LIP)+len(LHAND):,:]], axis=1)\n        rhand = x[:,len(LIP)+len(LHAND):len(LIP)+len(LHAND)+len(RHAND),:]\n        rhand = only_rotate(rhand,degree=(-30,30))\n        x = tf.concat([x[:,:len(LIP)+len(LHAND),:], rhand, x[:,len(LIP)+len(LHAND)+len(RHAND):,:]], axis=1)\n        pose = x[:,-POSE_NUMS:,:]\n        pose = only_rotate(pose,degree=(-30,30))\n        x = tf.concat([x[:,:-POSE_NUMS,:], pose], axis=1)\n        return x\n    if tf.shape(x)[0] > 10 and partial:\n        start = tf.random.uniform(shape=[], minval=0, maxval=tf.shape(x)[0]-1, dtype=tf.int32)\n        end = tf.random.uniform(shape=[], minval=start+1, maxval=tf.shape(x)[0], dtype=tf.int32)\n        x = tf.concat([x[:start], apply(x[start:end]), x[end:]], axis=0)\n    else:\n        return apply(x)\n    return x\n\ndef rotate_finger_partial(x,partial=False):\n    def apply(x):\n        lhand = x[:,len(LIP):len(LIP)+len(LHAND),:]\n        lhand = random_rotate_fingers(lhand)\n        rhand = x[:,len(LIP)+len(LHAND):len(LIP)+len(LHAND)+len(RHAND),:]\n        rhand = random_rotate_fingers(rhand)\n        x = tf.concat([x[:,:len(LIP),:], lhand, rhand, x[:,len(LIP)+len(LHAND)+len(RHAND):,:]], axis=1)\n        return x\n    if tf.shape(x)[0] > 10 and partial:\n        start = tf.random.uniform(shape=[], minval=0, maxval=tf.shape(x)[0]-1, dtype=tf.int32)\n        end = tf.random.uniform(shape=[], minval=start+1, maxval=tf.shape(x)[0], dtype=tf.int32)\n        x = tf.concat([x[:start], apply(x[start:end]), x[end:]], axis=0)\n    else:\n        return apply(x)\n    return x\n\ndef shift_partial(x,partial=False):\n    def apply(x):\n        face = tf.concat([x[:,:len(LIP),:], x[:,len(LIP)+len(LHAND)+len(RHAND):-POSE_NUMS,:]], axis=1)\n        face = random_shift(face)\n        lhand = x[:,len(LIP):len(LIP)+len(LHAND),:]\n        lhand = random_shift(lhand)\n        rhand = x[:,len(LIP)+len(LHAND):len(LIP)+len(LHAND)+len(RHAND),:]\n        rhand = random_shift(rhand)\n        pose = x[:,-POSE_NUMS:,:]\n        pose = random_shift(pose)\n        x = tf.concat([face[:,:len(LIP),:], lhand, rhand, face[:,len(LIP):,:],pose], axis=1)\n\n        return x\n    if tf.shape(x)[0] > 10 and partial:\n        start = tf.random.uniform(shape=[], minval=0, maxval=tf.shape(x)[0]-1, dtype=tf.int32)\n        end = tf.random.uniform(shape=[], minval=start+1, maxval=tf.shape(x)[0], dtype=tf.int32)\n        x = tf.concat([x[:start], apply(x[start:end]), x[end:]], axis=0)\n    else:\n        return apply(x)\n    return x\n\n\ndef combined_mask_along_axis(tensor, s1, s2, mask_value=float('nan')):\n    def get_mask(tensor, param_min, param_max, axis, mask_value=float('nan')):\n        tensor_shape = tf.shape(tensor)\n        dim_size = tensor_shape[axis]\n        min_mask_size = tf.cast(param_min * tf.cast(dim_size, tf.float32),tf.int32)\n        max_mask_size = tf.cast(param_max * tf.cast(dim_size, tf.float32), tf.int32)\n        mask_size = tf.cond(\n            tf.equal(min_mask_size, max_mask_size),\n            lambda: min_mask_size,\n            lambda: tf.random.uniform((), min_mask_size, max_mask_size+1, dtype=tf.int32)\n        )\n        mask_start = tf.random.uniform([], 0, dim_size - mask_size, tf.int32)\n        indices = tf.cast(tf.range(start=0, limit=dim_size, delta=1),tf.int32)\n        mask = tf.logical_or(indices < mask_start , indices >= (mask_start + mask_size))\n        if axis==1:\n            mask = tf.reshape(mask,(1,dim_size,1))\n        else:\n            mask = tf.reshape(mask,(dim_size,1,1))\n        return mask\n    t_mask = get_mask(tensor, s1[0], s1[1], 1, mask_value)\n    f_mask = get_mask(tensor, s2[0], s2[1], 0, mask_value)\n    mask = tf.logical_or(t_mask, f_mask)\n    mask = tf.cast(mask, tf.float32)\n    masked_tensor = tf.where(mask==0,mask_value,tensor)\n    return masked_tensor\n\ndef augment_fn(x):\n    if tf.random.uniform(()) > 0.33:\n        x = resample_sub(x)\n        # if tf.random.uniform(()) > 0.5:\n        #     x = resample_sub(x) # temporal\n        # else:\n        #     x = uniform_resample(x) # temporal\n    if tf.random.uniform(()) > 0.5:\n        # filp face\n        x = flip_lr(x)\n\n    if tf.random.uniform(()) < 0.5:\n        if tf.random.uniform(()) < 0.5:\n            x = rotate_partial(x,partial=True)\n        else:\n            x = rotate_partial(x,partial=False)\n\n    if tf.random.uniform(()) < 0.5:\n        if tf.random.uniform(()) < 0.5:\n            x = rotate_finger_partial(x,partial=True)\n        else:\n            x = rotate_finger_partial(x,partial=False)\n\n    # # ramdom shift\n    if tf.random.uniform(()) > 0.5:\n        if tf.random.uniform(()) < 0.5:\n            x = shift_partial(x,partial=True)\n        else:\n            x = shift_partial(x,partial=False)\n\n    if tf.random.uniform(()) < 0.5:\n        # T = tf.minimum(tf.cast(tf.shape(x)[0],tf.float32),200.0)\n        # factor = T/200 * 0.3\n        if tf.random.uniform(()) > 0.33:\n            x = mask_along_axis(x,0.0,0.4,0) # originally 0.2~0.4, can change to 0.1~0.3 for faster convergence\n        elif tf.random.uniform(()) > 0.5:\n            x = temporal_mask(x,size=(0.0,0.4))\n        else:\n            x = discrete_mask(x,0.0,0.4,0)\n\n    if tf.random.uniform(()) < 0.5:\n        # spatial masking\n        if tf.random.uniform(()) > 0.33:\n            x = mask_along_axis(x,0.0,0.4,1) # can be 0.2,0.4\n        elif tf.random.uniform(()) > 0.5:\n            x = spatial_mask(x,size=(0.0,0.4))\n        else:\n            x = discrete_mask(x,0.0,0.4,1)\n\n\n    # x = tf.where(tf.math.is_nan(x),NAN_FILL_VALUE,x)\n    return x\n\n\nKPT_NUM = (len(LIP) + len(LHAND) + len(RHAND) + len(NOSE) + len(LEYE) + len(REYE) + len(LPOSE) + len(RPOSE)) * DATA_DIM\n\n\ndef preprocess1(x):\n    x = split_data(x)\n    # x = tf.where(tf.math.is_nan(x), NAN_FILL_VALUE, x)\n    return x\n\n\ndef preprocess2(x):\n    # x = global_norm(x)\n    face = tf.concat([x[:,:len(LIP),:], x[:,len(LIP)+len(LHAND)+len(RHAND):-POSE_NUMS,:]], axis=1)\n    pose = x[:, -POSE_NUMS:, :]\n    lhand = x[:,len(LIP):len(LIP)+len(LHAND),:]\n    rhand = x[:,len(LIP)+len(LHAND):len(LIP)+len(LHAND)+len(RHAND),:]\n    face = self_norm(face)\n    lhand = self_norm(lhand)\n    rhand = self_norm(rhand)\n    pose = self_norm(pose)\n    x = tf.concat([face, lhand, rhand, pose], axis=1)\n    # x = tf.concat([face, lhand, rhand], axis=1)\n    # x = tf.concat([face[:,:len(LIP),:], x[:,len(LIP):len(LIP)+len(LHAND)+len(RHAND),:],face[:,len(LIP):,:],x[:,-POSE_NUMS:,:]], axis=1)\n    # x = x[...,:2]\n\n\n    x = resize_pad(x)\n\n\n    dx = x[1:,:,:] - x[:-1,:,:]\n    dx = tf.concat([tf.zeros((1, KPT_NUM//DATA_DIM, tf.shape(x)[-1])), dx], axis=0) # Tx21x2\n    if tf.shape(x)[0] > 1:\n        dx2 = x[2:,:,:] - x[:-2,:,:]\n        dx2 = tf.concat([tf.zeros((2, KPT_NUM//DATA_DIM, tf.shape(x)[-1])), dx2], axis=0) # Tx21x2\n    else:\n        dx2 = tf.zeros_like(dx)\n    x = tf.concat([x, dx,dx2], axis=-1)\n\n    x = tf.reshape(x, (FRAME_LEN, INPUT_DIM))\n    x = tf.where(tf.math.is_nan(x), NAN_FILL_VALUE, x)\n\n    return x\n\ndef total_process(x, phrase, augment=False):\n\n    x = preprocess1(x)\n    if augment:\n        x = augment_fn(x)\n    x = preprocess2(x)\n    return x,phrase\n\nINPUT_DIM = KPT_NUM * 3 # coordinate + velocity\nTEMPORAL_DIM = FRAME_LEN\n\ndef preprocess_(x):\n    x = preprocess1(x)\n    x = preprocess2(x)\n    return x\n\ndef preprocess_fn(x, phrase, augment=False):\n    batch = total_process(x, phrase,augment)\n    return batch\n\n\nprint(TEMPORAL_DIM,INPUT_DIM)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:38.096669Z","iopub.execute_input":"2023-08-20T06:39:38.097087Z","iopub.status.idle":"2023-08-20T06:39:38.293576Z","shell.execute_reply.started":"2023-08-20T06:39:38.097054Z","shell.execute_reply":"2023-08-20T06:39:38.292416Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"table = tf.lookup.StaticHashTable(\n    initializer=tf.lookup.KeyValueTensorInitializer(\n        keys=list(char_to_num.keys()),\n        values=list(char_to_num.values()),\n    ),\n    default_value=tf.constant(-1),\n    name=\"class_weight\",\n)\n\n\ndef decode_fn(record_bytes):\n    schema = {COL: tf.io.VarLenFeature(dtype=tf.float32) for COL in SEL_COLS}\n    schema[\"phrase\"] = tf.io.FixedLenFeature([], dtype=tf.string)\n    features = tf.io.parse_single_example(record_bytes, schema)\n\n    phrase = features[\"phrase\"]\n    landmarks = [tf.sparse.to_dense(features[COL]) for COL in SEL_COLS]\n    landmarks = tf.transpose(landmarks)\n    phrase = tf.strings.bytes_split(phrase)\n\n    phrase = table.lookup(phrase)\n    phrase = tf.pad(\n        phrase,\n        paddings=[[0, PHRASE_MAX_LEN - tf.shape(phrase)[0]]],\n        constant_values=pad_token_idx,\n    )\n    return landmarks, phrase\n\n\nval_file_ids = [234418913]#, 1967755728, 425182931]\nif VAL:\n    val_tffiles = df[df.file_id.isin(val_file_ids)].file_id.map(lambda x: f\"/kaggle/input/aslfr-preprocess-dataset/tfds-v2/{x}.tfrecord\").unique()\n    train_tffiles = df[~df.file_id.isin(val_file_ids)].file_id.map(lambda x: f\"/kaggle/input/aslfr-preprocess-dataset/tfds-v2/{x}.tfrecord\").unique()\nelse:\n    train_tffiles = df.file_id.map(lambda x: f\"/kaggle/input/aslfr-preprocess-dataset/tfds-v2/{x}.tfrecord\").unique()","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:38.294779Z","iopub.execute_input":"2023-08-20T06:39:38.295118Z","iopub.status.idle":"2023-08-20T06:39:38.374675Z","shell.execute_reply.started":"2023-08-20T06:39:38.295089Z","shell.execute_reply":"2023-08-20T06:39:38.373214Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# val_len = int(len(tffiles) * 0.05)\nval_size = BATCH_SIZE\nif VAL:\n    print(' Train Val Split')\n    # np.random.shuffle(tffiles)\n#     train_dataset = tf.data.TFRecordDataset(tffiles[val_len:],num_parallel_reads=tf.data.AUTOTUNE\n#                                             ).map(decode_fn,num_parallel_calls=tf.data.AUTOTUNE).map(lambda x, phrase: preprocess_fn(x, phrase,augment=True),num_parallel_calls=tf.data.AUTOTUNE).batch(BATCH_SIZE, drop_remainder=True).prefetch(buffer_size=tf.data.AUTOTUNE)\n\n#     val_dataset = tf.data.TFRecordDataset(tffiles[:val_len],num_parallel_reads=tf.data.AUTOTUNE\n#                                           ).map(decode_fn,num_parallel_calls=tf.data.AUTOTUNE).map(preprocess_fn,tf.data.AUTOTUNE).batch(BATCH_SIZE, drop_remainder=True).prefetch(buffer_size=tf.data.AUTOTUNE)\n    train_dataset = tf.data.TFRecordDataset(train_tffiles,num_parallel_reads=tf.data.AUTOTUNE\n                                            ).map(decode_fn,num_parallel_calls=tf.data.AUTOTUNE).map(lambda x, phrase: preprocess_fn(x, phrase,augment=True),num_parallel_calls=tf.data.AUTOTUNE).batch(BATCH_SIZE).repeat().prefetch(buffer_size=tf.data.AUTOTUNE)\n\n    val_dataset = tf.data.TFRecordDataset(val_tffiles,num_parallel_reads=tf.data.AUTOTUNE\n                                          ).map(decode_fn,num_parallel_calls=tf.data.AUTOTUNE).map(preprocess_fn,tf.data.AUTOTUNE).batch(BATCH_SIZE).repeat().prefetch(buffer_size=tf.data.AUTOTUNE)\n    aux_dataset = tf.data.TFRecordDataset(train_tffiles[:1],num_parallel_reads=tf.data.AUTOTUNE\n                                          ).map(decode_fn,num_parallel_calls=tf.data.AUTOTUNE).map(preprocess_fn,tf.data.AUTOTUNE).batch(BATCH_SIZE).repeat().prefetch(buffer_size=tf.data.AUTOTUNE)\n    print('train size', 64211+1997)\n    print('val size', 1000)\n    print('train datafiles',train_tffiles)\n    print('val datafiles',val_tffiles)\nelse:\n    print('Full Train')\n#     train_dataset = tf.data.TFRecordDataset(tffiles, num_parallel_reads=tf.data.AUTOTUNE\n#                                             ).map(decode_fn,tf.data.AUTOTUNE).map(lambda x, phrase: preprocess_fn(x, phrase,augment=True),num_parallel_calls=tf.data.AUTOTUNE).batch(BATCH_SIZE, drop_remainder=True).prefetch(buffer_size=tf.data.AUTOTUNE)\n    train_dataset = tf.data.TFRecordDataset(train_tffiles, num_parallel_reads=tf.data.AUTOTUNE\n                                             ).map(decode_fn,tf.data.AUTOTUNE).map(lambda x, phrase: preprocess_fn(x, phrase,augment=True),num_parallel_calls=tf.data.AUTOTUNE).batch(BATCH_SIZE).repeat().prefetch(buffer_size=tf.data.AUTOTUNE)\n\n    print('train size', 64211+2997)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:38.376236Z","iopub.execute_input":"2023-08-20T06:39:38.376595Z","iopub.status.idle":"2023-08-20T06:39:55.932096Z","shell.execute_reply.started":"2023-08-20T06:39:38.376566Z","shell.execute_reply":"2023-08-20T06:39:55.930852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LearningRateLogger(tf.keras.callbacks.Callback):\n    def on_epoch_end(self, epoch, logs=None):\n        lr = self.model.optimizer.lr\n        # If the learning rate is a decayed value, compute its value\n        if isinstance(lr, tf.keras.optimizers.schedules.LearningRateSchedule):\n            lr = lr(self.model.optimizer.iterations)\n        print(f\"\\nLearning rate at end of epoch {epoch}: {lr.numpy()}\\n\")\n\n# Instantiate the callback\nlr_logger = LearningRateLogger()","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:55.941927Z","iopub.execute_input":"2023-08-20T06:39:55.942364Z","iopub.status.idle":"2023-08-20T06:39:55.951151Z","shell.execute_reply.started":"2023-08-20T06:39:55.942329Z","shell.execute_reply":"2023-08-20T06:39:55.949317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class OneCycleLR(tf.keras.optimizers.schedules.LearningRateSchedule):\n    '''\n    Unified single-cycle learning rate scheduler for tensorflow.\n    2022 Hoyeol Sohn <hoeyol0730@gmail.com>\n    '''\n    def __init__(self,\n                lr=1e-3,\n                epochs=10,\n                steps_per_epoch=100,\n                steps_per_update=1,\n                resume_epoch=0,\n                decay_epochs=10,\n                sustain_epochs=0,\n                warmup_epochs=10,\n                lr_start=0,\n                lr_min=0,\n                warmup_type='linear',\n                decay_type='cosine',\n                finetune_steps=0,\n                finetune_lr=1e-5,\n                **kwargs):\n\n        super().__init__(**kwargs)\n        self.lr = float(lr)\n        self.epochs = float(epochs)\n        self.steps_per_update = float(steps_per_update)\n        self.resume_epoch = float(resume_epoch)\n        self.steps_per_epoch = float(steps_per_epoch)\n        self.decay_epochs = float(decay_epochs)\n        self.sustain_epochs = float(sustain_epochs)\n        self.warmup_epochs = float(warmup_epochs)\n        self.lr_start = float(lr_start)\n        self.lr_min = float(lr_min)\n        self.decay_type = decay_type\n        self.warmup_type = warmup_type\n        self.finetune_steps = finetune_steps\n        self.finetune_lr = float(finetune_lr)\n\n    def __call__(self, step):\n        step = tf.cast(step, tf.float32)\n        total_steps = self.epochs * self.steps_per_epoch\n        warmup_steps = self.warmup_epochs * self.steps_per_epoch\n        sustain_steps = self.sustain_epochs * self.steps_per_epoch\n        decay_steps = self.decay_epochs * self.steps_per_epoch\n\n        if self.resume_epoch > 0:\n            step = step + self.resume_epoch * self.steps_per_epoch\n\n        step = tf.cond(step > decay_steps, lambda :decay_steps, lambda :step)\n        step = tf.math.truediv(step, self.steps_per_update) * self.steps_per_update\n\n        warmup_cond = step < warmup_steps\n        decay_cond = step >= (warmup_steps + sustain_steps)\n        finetune_cond = step >= (total_steps - self.finetune_steps)\n\n        lr = tf.cond(warmup_cond, lambda: tf.math.divide_no_nan(self.lr-self.lr_start , warmup_steps) * step + self.lr_start, lambda: self.lr)\n\n        lr = tf.cond(decay_cond, lambda: 0.5 * (self.lr - self.lr_min) * (1 + tf.cos(3.14159265359 * (step - warmup_steps - sustain_steps) / (decay_steps - warmup_steps - sustain_steps))) + self.lr_min, lambda:lr)\n\n        lr = tf.cond(finetune_cond, lambda: tf.constant(self.finetune_lr,tf.float32), lambda:lr)\n        return lr\n","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:55.957545Z","iopub.execute_input":"2023-08-20T06:39:55.958034Z","iopub.status.idle":"2023-08-20T06:39:55.978902Z","shell.execute_reply.started":"2023-08-20T06:39:55.957996Z","shell.execute_reply":"2023-08-20T06:39:55.977752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LateDropout(tf.keras.layers.Layer):\n    def __init__(self, rate, noise_shape=None, start_step=0, **kwargs):\n        super().__init__(**kwargs)\n        self.supports_masking = True\n        self.rate = rate\n        self.start_step = start_step\n        self.dropout = tf.keras.layers.Dropout(rate, noise_shape=noise_shape)\n\n    def build(self, input_shape):\n        super().build(input_shape)\n        agg = tf.VariableAggregation.ONLY_FIRST_REPLICA\n        self._train_counter = tf.Variable(0, dtype=\"int64\", aggregation=agg, trainable=False)\n\n    def call(self, inputs, training=False):\n        x = tf.cond(self._train_counter < self.start_step, lambda:inputs, lambda:self.dropout(inputs, training=training))\n        if training:\n            self._train_counter.assign_add(1)\n        return x\n\nclass CausalDWConv1D(tf.keras.layers.Layer):\n    def __init__(self,\n        kernel_size=17,\n        dilation_rate=1,\n        use_bias=True,\n        depthwise_initializer='glorot_uniform',\n        name='', **kwargs):\n        super().__init__(name=name,**kwargs)\n        self.causal_pad = tf.keras.layers.ZeroPadding1D((dilation_rate*(kernel_size-1),0))\n        self.dw_conv = tf.keras.layers.DepthwiseConv1D(\n                            kernel_size,\n                            strides=1,\n                            dilation_rate=dilation_rate,\n                            padding='valid',\n                            use_bias=use_bias,\n                            depthwise_initializer=depthwise_initializer,)\n        self.supports_masking = True\n\n    def call(self, inputs):\n        x = self.causal_pad(inputs)\n        x = self.dw_conv(x)\n        return x\n\nclass GLU(tf.keras.layers.Layer):\n    def __init__(self, dim, **kwargs):\n        super(GLU, self).__init__(**kwargs)\n        self.dim = dim\n        self.supports_masking = True\n\n    def call(self, inputs):\n        out, gate = tf.split(inputs, 2, axis=self.dim)\n        return out * tf.sigmoid(gate)\n\n    def compute_mask(self, inputs, mask=None):\n        return mask\n\n\n\nclass CausalConv1D(tf.keras.layers.Layer):\n    def __init__(self,\n        hid_dim,\n        kernel_size=17,\n        dilation_rate=1,\n        groups = 1,\n        name='', **kwargs):\n        super().__init__()\n        self.causal_pad = tf.keras.layers.ZeroPadding1D((dilation_rate*(kernel_size-1),0))\n        self.dw_conv = tf.keras.layers.Conv1D(\n                            hid_dim,\n                            kernel_size,\n                            strides=1,\n                            dilation_rate=dilation_rate,\n                            padding='valid',groups=groups)\n        self.supports_masking = True\n\n    def call(self, inputs):\n        x = self.causal_pad(inputs)\n        x = self.dw_conv(x)\n        return x\n\n\nclass ConformerConvModule(tf.keras.layers.Layer):\n    def __init__(self, dim, kernel_size, dropout=DROP_OUT,expansion_factor=2):\n        super().__init__()\n        hid_dim = dim * expansion_factor\n        # self.postNorm = tf.keras.layers.BatchNormalization\n        self.preNorm = tf.keras.layers.BatchNormalization() #tf.keras.layers.BatchNormalization\n\n        if not MASKING:\n            self.pointwise_conv1 = tf.keras.layers.Conv1D(hid_dim, 1, padding='same')\n        else:\n            self.pointwise_conv1 = CausalConv1D(hid_dim, 1 )\n\n        self.glu = GLU(-1)\n\n        if not MASKING:\n            self.depthwise_conv = tf.keras.layers.Conv1D(dim, kernel_size, padding='same', groups=dim)\n        else:\n            self.depthwise_conv = CausalConv1D(dim, kernel_size, groups=dim)\n\n        self.batchnorm = tf.keras.layers.BatchNormalization()\n        self.activation = tf.keras.layers.Activation('swish')\n\n        if not MASKING:\n            self.pointwise_conv2 = tf.keras.layers.Conv1D(dim, 1, padding='same')\n        else:\n            self.pointwise_conv2 = CausalConv1D(dim, 1)\n\n        self.add = tf.keras.layers.Add()\n        self.dropout = tf.keras.layers.Dropout(dropout)\n        self.supports_masking = True\n        self.dropout2 = tf.keras.layers.Dropout(dropout) # DW dropout\n        self.dropout1 = tf.keras.layers.Dropout(dropout)\n\n    def call(self, x):\n        identity = x\n        x = self.preNorm(x)\n        x = self.pointwise_conv1(x)\n        x = self.glu(x)\n        x = self.dropout1(x)\n\n        x = self.depthwise_conv(x)\n        x = self.batchnorm(x)\n        x = self.activation(x)\n        x = self.dropout2(x)\n\n        x = self.pointwise_conv2(x)\n        x = self.dropout(x)\n        x = self.add([x,identity])\n        return x\n\n\nclass MultiHeadSelfAttention(tf.keras.layers.Layer):\n    def __init__(self, dim=256, num_heads=4, dropout=DROP_OUT, **kwargs):\n        super().__init__(**kwargs)\n        self.dim = dim\n        self.scale = self.dim ** -0.5\n        self.num_heads = num_heads\n        self.qkv = tf.keras.layers.Dense(3 * dim)\n        self.drop1 = tf.keras.layers.Dropout(dropout)\n        self.proj = tf.keras.layers.Dense(dim)\n        self.rel_embedding = self.add_weight(\"rel_embedding\", shape=[FRAME_LEN * 2 - 1, dim // num_heads])\n        self.supports_masking = True\n\n    def call(self, inputs, mask=None):\n        qkv = self.qkv(inputs)\n        qkv = tf.keras.layers.Permute((2, 1, 3))(\n            tf.keras.layers.Reshape(\n                (-1, self.num_heads, self.dim * 3 // self.num_heads)\n            )(qkv)\n        )\n        q, k, v = tf.split(qkv, [self.dim // self.num_heads] * 3, axis=-1)\n\n        # seq_len = tf.shape(q)[-2]\n        seq_len = FRAME_LEN\n        rel_indices = tf.range(seq_len)[:, None] - tf.range(seq_len)[None, :] + seq_len - 1\n        rel_indices = rel_indices[:tf.shape(q)[-2], :tf.shape(q)[-2]]\n        rel_k = tf.nn.embedding_lookup(self.rel_embedding, rel_indices)\n        rel_logits = tf.einsum('bhid,ijd->bhij', q, rel_k)\n\n        attn = (tf.matmul(q, k, transpose_b=True) + rel_logits) * self.scale\n\n        if mask is not None:\n            mask1 = tf.cast(mask[:, None, None,:],tf.float32)\n            mask2 = tf.cast(mask[:, None, : , None],tf.float32)\n            mask = mask2 @ mask1\n            attn = attn + ((1-mask) * -1e9)\n\n        attn = tf.keras.layers.Softmax(axis=-1)(attn)\n        attn = self.drop1(attn)\n\n        x = attn @ v\n        x = tf.keras.layers.Reshape((-1, self.dim))(\n            tf.keras.layers.Permute((2, 1, 3))(x)\n        )\n        x = self.proj(x)\n        return x\n\n\ndef ConformerBlock(dim=256, num_heads=8, expand=2, attn_dropout=DROP_OUT, drop_rate=DROP_OUT, activation='swish',ksize=11):\n    def apply(inputs):\n        x = inputs\n        # mlp1\n        x = tf.keras.layers.BatchNormalization()(x) #tf.keras.layers.BatchNormalization(x)\n        x = tf.keras.layers.Dense(dim*expand, activation=activation)(x)\n        x = tf.keras.layers.Dropout(drop_rate)(x)\n        x = tf.keras.layers.Dense(dim, kernel_initializer=\"he_normal\")(x)\n        x = tf.keras.layers.Dropout(drop_rate)(x)\n        mlp1_x = tf.keras.layers.Add()([inputs, 0.5*x])\n        # attn\n        x = tf.keras.layers.BatchNormalization()(mlp1_x) #tf.keras.layers.BatchNormalization(mlp1_x)\n        x = MultiHeadSelfAttention(dim=dim,num_heads=num_heads,dropout=attn_dropout)(x)\n        x = tf.keras.layers.Dropout(drop_rate)(x)\n        attn_out = tf.keras.layers.Add()([mlp1_x, x])\n        # attn_out = tf.keras.layers.BatchNormalization(x)\n        # conv\n        conv_out = ConformerConvModule(dim,ksize)(attn_out)\n        # mlp2\n        x = tf.keras.layers.BatchNormalization()(conv_out) #tf.keras.layers.BatchNormalization(conv_out)\n        x = tf.keras.layers.Dense(dim*expand, activation=activation)(x)\n        x = tf.keras.layers.Dropout(drop_rate)(x)\n        x = tf.keras.layers.Dense(dim, kernel_initializer=\"he_normal\")(x)\n        x = tf.keras.layers.Dropout(drop_rate)(x)\n        x = tf.keras.layers.Add()([conv_out, 0.5*x])\n        # x = tf.keras.layers.BatchNormalization()(x)\n        return x\n    return apply\n\n# copy from https://www.kaggle.com/code/markwijkhuizen/aslfr-transformer-training-inference#Landmark-Embedding\nclass LandmarkEmbedding(tf.keras.layers.Layer):\n    def __init__(self, units, name='emb'):\n        super(LandmarkEmbedding, self).__init__(name=f'{name}_embedding')\n        self.supports_masking = True\n\n        # self.empty_embedding = self.add_weight(\n        #     name=f'{name}_empty_embedding',\n        #     shape=[units],\n        #     initializer=tf.keras.initializers.constant(0.0),\n        #     trainable=True,\n        # )\n        # Embedding\n        self.dense = tf.keras.Sequential([\n            tf.keras.layers.Dense((units+INPUT_DIM)//2,activation='swish'),\n            tf.keras.layers.Dense(units),\n            tf.keras.layers.BatchNormalization(),\n        ], name=f'{name}_dense')\n\n    def call(self, x):\n        return self.dense(x)\n        # return tf.where(\n        #         # Checks whether landmark is missing in frame\n        #         tf.reduce_all(x==PADDING_MASKING_VALUE, axis=-1, keepdims=True),\n        #         # If so, the empty embedding is used\n        #         self.empty_embedding,\n        #         # Otherwise the landmark data is embedded\n        #         self.dense(x),\n        #     )","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:55.981033Z","iopub.execute_input":"2023-08-20T06:39:55.981405Z","iopub.status.idle":"2023-08-20T06:39:56.048476Z","shell.execute_reply.started":"2023-08-20T06:39:55.981371Z","shell.execute_reply":"2023-08-20T06:39:56.046405Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def num_to_char_fn(y):\n    return [num_to_char.get(x, \"\") for x in y]\n\n# @tf.function()\n# def decode_phrase(pred):\n#     x = tf.argmax(pred, axis=1)\n#     diff = tf.not_equal(x[:-1], x[1:])\n#     adjacent_indices = tf.where(diff)[:, 0]\n    \n#     # Adding the index of the last token\n#     adjacent_indices = tf.concat([adjacent_indices, [tf.size(x) - 1]], 0)\n    \n#     x = tf.gather(x, adjacent_indices)\n#     mask = x != pad_token_idx\n#     x = tf.boolean_mask(x, mask, axis=0)\n#     return x\n\nspace_token_idx = char_to_num[' ']\n@tf.function()\ndef decode_phrase(pred):\n    x = tf.argmax(pred, axis=1)\n    diff = tf.not_equal(x[:-1], x[1:])\n    adjacent_indices = tf.where(diff)[:, 0]\n    \n    # Adding the index of the last token\n    adjacent_indices = tf.concat([adjacent_indices, [tf.size(x) - 1]], 0)\n    \n    x = tf.gather(x, adjacent_indices)\n    mask = x != pad_token_idx\n    x = tf.boolean_mask(x, mask, axis=0)\n#     no_space_x = tf.boolean_mask(x, x != space_token_idx, axis=0)\n#     if tf.shape(no_space_x)[0] < 4:\n#         x = tf.concat([no_space_x,tf.convert_to_tensor([char_to_num[c] for c in ' -aero'],dtype=tf.int64)],axis=0)\n    if tf.shape(x)[0] < 4:\n        x = tf.concat([x,tf.convert_to_tensor([char_to_num[c] for c in ' -aero'],dtype=tf.int64)],axis=0)\n    return x\n\n# A utility function to decode the output of the network\ndef decode_batch_predictions(pred):\n    output_text = []\n    for result in pred:\n        result = \"\".join(num_to_char_fn(decode_phrase(result).numpy()))\n        output_text.append(result)\n    return output_text\n\n# A callback class to output a few transcriptions during training\nclass CallbackEval(tf.keras.callbacks.Callback):\n    \"\"\"Displays a batch of outputs after every epoch.\"\"\"\n\n    def __init__(self, dataset,aux_dataset,model):\n        super().__init__()\n        self.dataset = dataset\n        self.aux_dataset = aux_dataset\n        self.model = model\n        self.train_score = []\n        self.val_score = []\n    def on_epoch_end(self, epoch: int, logs=None):\n        if epoch % 5 !=0:\n            return\n        start_time = time.time()\n        d = 0\n        n = 0\n        for batch in self.dataset.take(1000//BATCH_SIZE):\n            X, y = batch\n            batch_predictions = self.model([X])\n            batch_predictions = decode_batch_predictions(batch_predictions)\n            # predictions.extend(batch_predictions)\n            for i,label in enumerate(y):\n                label = \"\".join(num_to_char_fn(label.numpy())).replace(pad_token,'')\n                n += len(label)\n                d += lev.distance(label, batch_predictions[i])\n        print('val metric: ', (n-d)/n)\n        self.val_score.append((n-d)/n)\n        d = 0\n        n = 0\n        for batch in self.aux_dataset.take(1000//BATCH_SIZE):\n            X, y = batch\n            batch_predictions = self.model([X])\n            batch_predictions = decode_batch_predictions(batch_predictions)\n            # predictions.extend(batch_predictions)\n            for i,label in enumerate(y):\n                label = \"\".join(num_to_char_fn(label.numpy())).replace(pad_token,'')\n                n += len(label)\n                d += lev.distance(label, batch_predictions[i])\n        print('train metric: ', (n-d)/n)\n        self.train_score.append((n-d)/n)\n        print('eval time: ', time.time() - start_time)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:56.0508Z","iopub.execute_input":"2023-08-20T06:39:56.051334Z","iopub.status.idle":"2023-08-20T06:39:56.08338Z","shell.execute_reply.started":"2023-08-20T06:39:56.051282Z","shell.execute_reply":"2023-08-20T06:39:56.078926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class RestoreBestWeightsIfIncrease(tf.keras.callbacks.Callback):\n    def __init__(self,optimizer,patience=1, restore_threshold=0.1,min_epoch=0):\n        super(RestoreBestWeightsIfIncrease, self).__init__()\n        self.optimizer = optimizer\n        self.best_optimizer_weights = None\n        self.patience = patience\n        self.restore_threshold = restore_threshold\n        # best_weights to store the weights at which the minimum loss occurs.\n        self.best_weights = None\n        # best_loss Will keep track of the lowest loss so far.\n        self.best_loss = np.Inf\n        # wait Will keep track of the number of epochs the training has waited when loss is no longer minimum.\n        self.wait = 0\n        self.min_epoch = min_epoch\n    def on_epoch_end(self, epoch, logs=None):\n        current_val_loss = logs.get(\"loss\")\n        if np.less(current_val_loss, self.best_loss):\n            self.best_loss = current_val_loss\n            self.wait = 0\n            # Record the best weights if current results is better (less).\n            self.best_weights = self.model.get_weights()\n            self.best_optimizer_weights = self.optimizer.get_weights()\n\n            print('record best loss = ',self.best_loss)\n        else:\n            self.wait += 1\n            if self.wait >= self.patience and epoch > self.min_epoch:\n                if current_val_loss - self.best_loss > self.restore_threshold:\n                    self.model.set_weights(self.best_weights)\n                    self.optimizer.set_weights(self.best_optimizer_weights)\n                    print(\"\\nRestoring model weights from the end of the best epoch.\")\n                    self.wait = 0\n","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:56.085854Z","iopub.execute_input":"2023-08-20T06:39:56.086975Z","iopub.status.idle":"2023-08-20T06:39:56.106798Z","shell.execute_reply.started":"2023-08-20T06:39:56.086927Z","shell.execute_reply":"2023-08-20T06:39:56.105076Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SAM(tf.keras.Model):\n    def __init__(self, *args, rho=.05, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.rho = rho\n    def _grad_norm(self, gradients):\n        norm = tf.norm(\n            tf.stack([\n                tf.norm(grad) for grad in gradients if grad is not None\n            ])\n        )\n        return norm\n    def train_step_sam(self, data):\n        # Unpack the data. Its structure depends on your model and\n        # on what you pass to `fit()`.\n        x, y = data\n        e_ws = []\n\n        with tf.GradientTape() as tape:\n            y_pred = self(x, training=True)\n            loss = self.compiled_loss(y, y_pred)\n\n        trainable_params = self.trainable_variables\n        gradients = tape.gradient(loss, trainable_params)\n        grad_norm = self._grad_norm(gradients)\n        scale = self.rho / (grad_norm + 1e-12)\n\n        for (grad, param) in zip(gradients, trainable_params):\n            e_w = grad * scale\n            param.assign_add(e_w)\n            e_ws.append(e_w)\n\n\n        with tf.GradientTape() as tape2:\n            y_pred = self(x, training=True)\n            new_loss = self.compiled_loss(y, y_pred)\n\n        sam_gradients = tape2.gradient(new_loss, trainable_params)\n\n        for (param, e_w) in zip(trainable_params, e_ws):\n            param.assign_sub(e_w)\n\n        self.optimizer.apply_gradients(\n            zip(sam_gradients, trainable_params))\n\n        self.compiled_metrics.update_state(y, y_pred)\n        return {m.name: m.result() for m in self.metrics}\n\n\n    def train_step(self, data):\n        return self.train_step_sam(data)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:56.108961Z","iopub.execute_input":"2023-08-20T06:39:56.10944Z","iopub.status.idle":"2023-08-20T06:39:56.132656Z","shell.execute_reply.started":"2023-08-20T06:39:56.109399Z","shell.execute_reply":"2023-08-20T06:39:56.130758Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomModelCheckpoint(tf.keras.callbacks.Callback):\n    def __init__(self,model):\n        super().__init__()\n        self.model = model\n    def on_epoch_end(self, epoch, logs=None):\n        if epoch >= 350 and (epoch + 1) % 10 == 0:  # Check if this is the 5th epoch (or multiple thereof)\n            self.model.save_weights(f'model_weights_{epoch + 1:02d}.h5')","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:56.134745Z","iopub.execute_input":"2023-08-20T06:39:56.136232Z","iopub.status.idle":"2023-08-20T06:39:56.157658Z","shell.execute_reply.started":"2023-08-20T06:39:56.136177Z","shell.execute_reply":"2023-08-20T06:39:56.156589Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def positional_encoding(maxlen, num_hid):\n    depth = num_hid/2\n    positions = tf.range(maxlen, dtype = tf.float32)[..., tf.newaxis]\n    depths = tf.range(depth, dtype = tf.float32)[np.newaxis, :]/depth\n    angle_rates = tf.math.divide(1, tf.math.pow(tf.cast(10000, tf.float32), depths))\n    angle_rads = tf.linalg.matmul(positions, angle_rates)\n    pos_encoding = tf.concat(\n      [tf.math.sin(angle_rads), tf.math.cos(angle_rads)],\n      axis=-1)\n    return pos_encoding","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:56.159403Z","iopub.execute_input":"2023-08-20T06:39:56.160342Z","iopub.status.idle":"2023-08-20T06:39:56.172871Z","shell.execute_reply.started":"2023-08-20T06:39:56.160295Z","shell.execute_reply":"2023-08-20T06:39:56.171474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"steps_per_epoch = int((64211 if VAL else 67208)//BATCH_SIZE) + 1\n\n\ndata_cfg = {}\ndata_cfg['FRAME_LEN'] = FRAME_LEN\ndata_cfg['INPUT_DIM'] = INPUT_DIM\ndata_cfg['batch_size'] = BATCH_SIZE\n\n\n\n\nmodel_cfg = {}\nmodel_cfg[\"num_layers_enc\"] = 6\nmodel_cfg[\"encoder_dim\"] = 384\nmodel_cfg[\"num_head\"] =  8\nmodel_cfg['kernel_size'] = 27 # should back to 13\nmodel_cfg[\"decoder_dim\"] = 384 # 256\n\n\nmodel_cfg[\"source_maxlen\"] = TEMPORAL_DIM\nmodel_cfg[\"target_maxlen\"] = PHRASE_MAX_LEN\nmodel_cfg[\"num_classes\"] = CLASS_NUM\n\n\ndef get_model():\n    inp = tf.keras.layers.Input(shape=(TEMPORAL_DIM,INPUT_DIM))\n    if not MASKING:\n        x = inp\n    else:\n        x = tf.keras.layers.Masking(mask_value=PADDING_MASKING_VALUE)(inp)\n    # x = tf.keras.layers.Dense(model_cfg[\"encoder_dim\"]*2, activation='swish')(x)\n    # x = tf.keras.layers.Dense(model_cfg[\"encoder_dim\"])(x)\n\n    x = LandmarkEmbedding(model_cfg[\"encoder_dim\"])(x)\n    # x = tf.keras.layers.BatchNormalization()(x)\n    # pe = positional_encoding(TEMPORAL_DIM, model_cfg[\"encoder_dim\"])\n    # x = x + pe\n    for i in range(model_cfg['num_layers_enc']):\n        x = ConformerBlock(dim=model_cfg[\"encoder_dim\"], num_heads=model_cfg[\"num_head\"],ksize=model_cfg['kernel_size'])(x)\n\n    x = tf.keras.layers.GRU(model_cfg['decoder_dim'],return_sequences=True)(x)\n    x = tf.keras.layers.Dense(model_cfg['num_classes'])(x)\n    model = tf.keras.Model(inputs=inp,outputs=x)\n    return model\n\n\ndef aux_ctc_loss(y_true, y_pred):\n        y_true = tf.ensure_shape(y_true, [BATCH_SIZE//strategy.num_replicas_in_sync,PHRASE_MAX_LEN])\n        y_pred = tf.ensure_shape(y_pred, [BATCH_SIZE//strategy.num_replicas_in_sync,TEMPORAL_DIM,CLASS_NUM])\n        label_length = tf.reduce_sum(tf.cast(y_true != pad_token_idx, tf.int32), axis=-1)\n        logit_length = tf.ones(tf.shape(y_pred)[0], dtype=tf.int32) * tf.shape(y_pred)[1]\n        loss = classic_ctc_loss(\n                labels=y_true,\n                logits=y_pred,\n                label_length=label_length,\n                logit_length=logit_length,\n                blank_index=pad_token_idx,\n        )\n        loss = tf.reduce_mean(loss)\n        return loss\n\nlr_cfg = {}\nlr_cfg['lr'] =  1e-3\nlr_cfg['weight_decay'] = 1e-3 #1e-6\nlr_cfg['epochs'] = 500\n\nlr_cfg['optimizer'] = tfa.optimizers.RectifiedAdam\nlr_cfg['alpha'] = 0.05 # final lr = lr * alpha\nlr_cfg['finetune_epochs'] = 0\nlr_cfg['warmup_epochs'] = 0 # 0.1*lr_cfg['epochs']\ninitial_learning_rate = lr_cfg['lr']\n\n\nlr_cfg['scheduler'] = OneCycleLR\n\nlr_schedule = lr_cfg['scheduler'](\n    lr = lr_cfg['lr'],\n    epochs = lr_cfg['epochs'] + lr_cfg['finetune_epochs'],\n    steps_per_epoch = steps_per_epoch,\n    steps_per_update = 1,\n    resume_epoch = 0,\n    decay_epochs = lr_cfg['epochs'] - lr_cfg['warmup_epochs'],\n    sustain_epochs = 0,\n    warmup_epochs = lr_cfg['warmup_epochs'],\n    finetune_steps = 0,\n    finetune_lr = 0.05 * lr_cfg['lr'],\n    lr_start = lr_cfg['lr']*0.01,\n    lr_min = lr_cfg['lr']*lr_cfg['alpha'],\n)\n\n\nwith strategy.scope():\n    if VAL:\n        batch = next(iter(val_dataset))\n    else:\n        batch = next(iter(train_dataset))\n    print(\"raw input shape\", batch[0].shape)\n    print(\"raw target shape\", batch[1].shape)\n\n\n    model_cb = tf.keras.callbacks.ModelCheckpoint(\n        'best.h5',\n        monitor = 'val_loss' if VAL else 'loss',\n        verbose = 1,\n        save_best_only = True,\n        save_weights_only= True,\n        mode = 'auto',\n        save_freq='epoch',\n        options=None,\n        initial_value_threshold=None,\n    )\n\n    optimizer = lr_cfg['optimizer'](lr_schedule, weight_decay = lr_cfg['weight_decay'],) # clipnorm = 5.0)\n    optimizer = tfa.optimizers.Lookahead(optimizer, sync_period=5)\n    model = get_model()\n#     model = SAM(model.input,model.output)\n    model.compile(optimizer=optimizer, loss=aux_ctc_loss,steps_per_execution=steps_per_epoch)\n    model.summary()\n    model.build((None,TEMPORAL_DIM,INPUT_DIM))\n\n    restore_cb = RestoreBestWeightsIfIncrease(optimizer,patience=1, restore_threshold=1.0,min_epoch=10)\n    epochwise_checkpoint = CustomModelCheckpoint(model)\n    callbacks = [lr_logger,model_cb, restore_cb,epochwise_checkpoint]\n\n    if not LOAD:\n        history = model.fit(\n            train_dataset,\n            validation_data = val_dataset if VAL else None,\n            callbacks=callbacks,\n            epochs=lr_cfg['epochs'],\n            steps_per_epoch=steps_per_epoch,\n            validation_steps=0 if not VAL else -(1000//-BATCH_SIZE),\n        )\n        model.save_weights('model.h5')\n\n# SAM -> 0.767 (100 -> 0.754)\n# No SAM 50/0.71","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:39:56.175476Z","iopub.execute_input":"2023-08-20T06:39:56.176306Z","iopub.status.idle":"2023-08-20T06:40:07.467406Z","shell.execute_reply.started":"2023-08-20T06:39:56.176261Z","shell.execute_reply":"2023-08-20T06:40:07.465928Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TFLiteModel(tf.Module):\n    def __init__(\n        self,\n        model,\n    ):\n        super(TFLiteModel, self).__init__()\n        self.model = model\n\n    @tf.function(\n        input_signature=[\n            tf.TensorSpec(shape=[None, len(SEL_COLS)], dtype=tf.float32, name=\"inputs\")\n        ]\n    )\n    def __call__(self, inputs, training=False):\n        # Preprocess Data\n        x = tf.cast(inputs, tf.float32)\n        x = x[None]\n        x = tf.cond(\n            tf.shape(x)[1] == 0,\n            lambda: tf.zeros((1, 1, len(SEL_COLS))),\n            lambda: tf.identity(x),\n        )\n        x = x[0]\n        x = preprocess1(x)\n        x = preprocess2(x)\n        x = x[None]\n        x = self.model(x)\n        x = x[0]\n        x = decode_phrase(x)\n        x = tf.cond(tf.shape(x)[0] == 0, lambda: tf.zeros(1, tf.int64), lambda: tf.identity(x))\n        x = tf.one_hot(x, 59)\n        return {\"outputs\": x}\n\ntflitemodel_base = TFLiteModel(model)","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:40:07.497036Z","iopub.execute_input":"2023-08-20T06:40:07.497516Z","iopub.status.idle":"2023-08-20T06:40:07.517473Z","shell.execute_reply.started":"2023-08-20T06:40:07.497472Z","shell.execute_reply":"2023-08-20T06:40:07.515778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"keras_model_converter = tf.lite.TFLiteConverter.from_keras_model(tflitemodel_base)\n# before\nkeras_model_converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]\n# keras_model_converter.optimizations = [tf.lite.Optimize.DEFAULT] # comment for speed\n\n# now\nkeras_model_converter.optimizations = [tf.lite.Optimize.DEFAULT]\nkeras_model_converter.target_spec.supported_types = [tf.float16]\n\ntflite_model = keras_model_converter.convert()\nwith open(\"model.tflite\", \"wb\") as f:\n    f.write(tflite_model)\n\ninfargs = {\"selected_columns\": SEL_COLS}\n\nwith open(\"inference_args.json\", \"w\") as json_file:\n    json.dump(infargs, json_file)\n","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:40:07.519452Z","iopub.execute_input":"2023-08-20T06:40:07.519918Z","iopub.status.idle":"2023-08-20T06:42:26.804984Z","shell.execute_reply.started":"2023-08-20T06:40:07.519878Z","shell.execute_reply":"2023-08-20T06:42:26.803468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import zipfile\n\n# list of file names to be zipped\nfile_names = [\"model.tflite\", \"inference_args.json\"]\n\n# name of the zip file\nzip_name = \"submission.zip\"\n\n# create a ZipFile object\nwith zipfile.ZipFile(zip_name, 'w') as zipf:\n    # write each file into the zip file\n    for file in file_names:\n        zipf.write(file)\n\nprint(f'{zip_name} file is created.')","metadata":{"execution":{"iopub.status.busy":"2023-08-20T06:42:26.807384Z","iopub.execute_input":"2023-08-20T06:42:26.807933Z","iopub.status.idle":"2023-08-20T06:42:30.322552Z","shell.execute_reply.started":"2023-08-20T06:42:26.807883Z","shell.execute_reply":"2023-08-20T06:42:30.321226Z"},"trusted":true},"execution_count":null,"outputs":[]}]}