{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":3649,"databundleVersionId":46718,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Imports & Paths","metadata":{}},{"cell_type":"code","source":"!pip install py7zr\nimport os, glob\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport py7zr\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\n\n# Path to the uploaded CIFAR10 7z dataset\nDATA_PATH = \"/kaggle/input/cifar-10/train.7z\"   # adjust if different filename\nTEST_PATH = \"/kaggle/input/cifar-10/test.7z\"    # adjust if different filename\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:10:34.605333Z","iopub.execute_input":"2025-10-09T10:10:34.605520Z","iopub.status.idle":"2025-10-09T10:10:53.991825Z","shell.execute_reply.started":"2025-10-09T10:10:34.605503Z","shell.execute_reply":"2025-10-09T10:10:53.991048Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Extract Data","metadata":{}},{"cell_type":"code","source":"# Extract train and test archives\nos.makedirs(\"/kaggle/working/cifar10/train\", exist_ok=True)\nos.makedirs(\"/kaggle/working/cifar10/test\", exist_ok=True)\n\nprint(\"Extracting train...\")\nwith py7zr.SevenZipFile(DATA_PATH, mode='r') as archive:\n    archive.extractall(path=\"/kaggle/working/cifar10/train\")\n\nprint(\"Extracting test...\")\nwith py7zr.SevenZipFile(TEST_PATH, mode='r') as archive:\n    archive.extractall(path=\"/kaggle/working/cifar10/test\")\n\nprint(\"Extraction complete!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:10:53.992837Z","iopub.execute_input":"2025-10-09T10:10:53.993334Z","iopub.status.idle":"2025-10-09T10:13:25.695546Z","shell.execute_reply.started":"2025-10-09T10:10:53.993301Z","shell.execute_reply":"2025-10-09T10:13:25.694880Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load Data","metadata":{}},{"cell_type":"code","source":"# Training data: usually in /train/{class_name}/{images}\ntrain_dir = \"/kaggle/working/cifar10/train\"\ntest_dir  = \"/kaggle/working/cifar10/test/test\"\n\n# Build class map\nclasses = sorted(os.listdir(train_dir))\nclass_to_idx = {c: i for i, c in enumerate(classes)}\nprint(\"Classes:\", classes)\n\nX_train, y_train = [], []\nfor c in classes:\n    folder = os.path.join(train_dir, c)\n    for img_file in glob.glob(os.path.join(folder, \"*.png\")):\n        img = Image.open(img_file).convert(\"RGB\").resize((32,32))\n        X_train.append(np.array(img))\n        y_train.append(class_to_idx[c])\n\nX_train = np.array(X_train, dtype=\"float32\") / 255.0\ny_train = np.array(y_train, dtype=\"int32\")\n\nprint(\"Train shape:\", X_train.shape, y_train.shape)\n\n# Test data\nX_test, test_ids = [], []\nfor img_file in sorted(glob.glob(os.path.join(test_dir, \"*.png\"))):\n    img_id = os.path.basename(img_file).split(\".\")[0]\n    img = Image.open(img_file).convert(\"RGB\").resize((32,32))\n    X_test.append(np.array(img))\n    test_ids.append(img_id)\n\nX_test = np.array(X_test, dtype=\"float32\") / 255.0\nprint(\"Test shape:\", X_test.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:13:25.697321Z","iopub.execute_input":"2025-10-09T10:13:25.697532Z","iopub.status.idle":"2025-10-09T10:14:25.638561Z","shell.execute_reply.started":"2025-10-09T10:13:25.697515Z","shell.execute_reply":"2025-10-09T10:14:25.637877Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train/Validation Split","metadata":{}},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n\nX_train, X_val, y_train, y_val = train_test_split(\n    X_train, y_train, test_size=0.1, stratify=y_train, random_state=42\n)\n\nprint(\"Train:\", X_train.shape, y_train.shape)\nprint(\"Val:\", X_val.shape, y_val.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:14:25.639226Z","iopub.execute_input":"2025-10-09T10:14:25.639442Z","iopub.status.idle":"2025-10-09T10:14:25.942198Z","shell.execute_reply.started":"2025-10-09T10:14:25.639424Z","shell.execute_reply":"2025-10-09T10:14:25.941536Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Building The Model","metadata":{}},{"cell_type":"code","source":"\nimport os, math, time, sys, json, random\nimport numpy as np\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\n\nprint(tf.__version__)\n\n# Reproducibility (still allows GPU nondeterminism in some ops)\nSEED = 1337\nrandom.seed(SEED); np.random.seed(SEED); tf.random.set_seed(SEED)\n\nNUM_CLASSES = 10\nINPUT_SHAPE = (32, 32, 3)  # keep as-is for CIFAR-10\nBATCH_SIZE = 256\nEPOCHS = 100\n\n# Safety: make sure memory growth is enabled (prevents OOM on some Kaggle GPUs)\ngpus = tf.config.list_physical_devices('GPU')\nfor g in gpus:\n    try:\n        tf.config.experimental.set_memory_growth(g, True)\n    except:\n        pass\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:14:25.943073Z","iopub.execute_input":"2025-10-09T10:14:25.943366Z","iopub.status.idle":"2025-10-09T10:14:26.545805Z","shell.execute_reply.started":"2025-10-09T10:14:25.943341Z","shell.execute_reply":"2025-10-09T10:14:26.544966Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Advanced data augmentation: CutMix + CutOut (probabilistic)","metadata":{}},{"cell_type":"code","source":"\nSEED = 1337\nAUTO = tf.data.AUTOTUNE\nNUM_CLASSES_TF = tf.constant(10, dtype=tf.int32)  # adjust if needed\n\nbasic_augment = keras.Sequential([\n    layers.RandomFlip(\"horizontal\", seed=SEED),\n    layers.RandomTranslation(0.1, 0.1, seed=SEED),\n    layers.RandomZoom(0.1, seed=SEED),\n], name=\"basic_augment\")\n\ndef ensure_one_hot_float(labels):\n    \"\"\"Accepts [B] ints or [B,C] one-hots (int/float). Returns float32 one-hot [B,C].\"\"\"\n    rank = tf.rank(labels)\n    def _from_indices():\n        idx = tf.cast(labels, tf.int32)              # [B]\n        return tf.one_hot(idx, NUM_CLASSES_TF)       # float32 [B,C]\n    def _from_onehot():\n        return tf.cast(labels, tf.float32)           # float32 [B,C]\n    return tf.cond(tf.equal(rank, 1), _from_indices, _from_onehot)\n\ndef _rand_bbox(height, width, lam):\n    height = tf.cast(height, tf.float32)\n    width  = tf.cast(width,  tf.float32)\n    cut_rat = tf.sqrt(1.0 - lam)                     # float32\n    cut_h = tf.cast(height * cut_rat, tf.int32)\n    cut_w = tf.cast(width  * cut_rat, tf.int32)\n    cy = tf.random.uniform([], 0, tf.cast(height, tf.int32), dtype=tf.int32)\n    cx = tf.random.uniform([], 0, tf.cast(width,  tf.int32), dtype=tf.int32)\n    y1 = tf.clip_by_value(cy - cut_h // 2, 0, tf.cast(height, tf.int32))\n    y2 = tf.clip_by_value(cy + cut_h // 2, 0, tf.cast(height, tf.int32))\n    x1 = tf.clip_by_value(cx - cut_w // 2, 0, tf.cast(width,  tf.int32))\n    x2 = tf.clip_by_value(cx + cut_w // 2, 0, tf.cast(width,  tf.int32))\n    return y1, x1, y2, x2\n\n@tf.function\ndef cutmix(images, labels, alpha=1.0):\n    \"\"\"Graph-safe CutMix (expects images float32 [B,H,W,C], labels one-hot float32 [B,C]).\"\"\"\n    images = tf.cast(images, tf.float32)\n    labels = ensure_one_hot_float(labels)\n\n    B = tf.shape(images)[0]\n    H = tf.shape(images)[1]\n    W = tf.shape(images)[2]\n\n    idx = tf.random.shuffle(tf.range(B))\n    images_shuf = tf.gather(images, idx)\n    labels_shuf = tf.gather(labels, idx)\n\n    g1 = tf.random.gamma([1], alpha=alpha, dtype=tf.float32)\n    g2 = tf.random.gamma([1], alpha=alpha, dtype=tf.float32)\n    lam = tf.squeeze(g1 / (g1 + g2))\n    lam = tf.clip_by_value(lam, 0.0, 1.0)\n\n    y1, x1, y2, x2 = _rand_bbox(H, W, lam)\n    box_h = tf.maximum(y2 - y1, 0)\n    box_w = tf.maximum(x2 - x1, 0)\n\n    def mix_with_box():\n        M_core = tf.ones([B, box_h, box_w, 1], dtype=images.dtype)\n        M = tf.pad(M_core, [[0,0],[y1, H - y2],[x1, W - x2],[0,0]])\n        mixed = images * (1.0 - M) + images_shuf * M\n        box_area = tf.cast(box_h * box_w, tf.float32)\n        lam_adj = 1.0 - box_area / tf.cast(H * W, tf.float32)\n        new_labels = lam_adj * labels + (1.0 - lam_adj) * labels_shuf\n        return mixed, new_labels\n\n    return tf.cond(tf.logical_and(box_h > 0, box_w > 0), mix_with_box, lambda: (images, labels))\n\n@tf.function\ndef cutout_single(image, mask_size=12):\n    image = tf.cast(image, tf.float32)\n    H = tf.shape(image)[0]; W = tf.shape(image)[1]\n    cy = tf.random.uniform([], 0, H, dtype=tf.int32)\n    cx = tf.random.uniform([], 0, W, dtype=tf.int32)\n    y1 = tf.clip_by_value(cy - mask_size // 2, 0, H)\n    y2 = tf.clip_by_value(cy + mask_size // 2, 0, H)\n    x1 = tf.clip_by_value(cx - mask_size // 2, 0, W)\n    x2 = tf.clip_by_value(cx + mask_size // 2, 0, W)\n    box_h = tf.maximum(y2 - y1, 0); box_w = tf.maximum(x2 - x1, 0)\n\n    def erase():\n        rect = tf.ones([box_h, box_w, tf.shape(image)[-1]], dtype=image.dtype)\n        rect = tf.pad(rect, [[y1, H - y2],[x1, W - x2],[0,0]])\n        return image * (1.0 - rect)\n\n    return tf.cond(tf.logical_and(box_h > 0, box_w > 0), erase, lambda: image)\n\n@tf.function\ndef apply_strong_aug(images, labels, p_cutmix=0.5, p_cutout=0.5):\n    images = tf.cast(images, tf.float32)\n    labels = ensure_one_hot_float(labels)\n\n    images = basic_augment(images, training=True)\n\n    rnd = tf.random.uniform([], 0, 1)\n    def _do_cutmix():\n        return cutmix(images, labels)\n\n    def _maybe_cutout():\n        B = tf.shape(images)[0]\n        flags = tf.random.uniform([B], 0, 1) < p_cutout\n        def _one(img, flag):\n            return tf.cond(flag, lambda: cutout_single(img), lambda: tf.cast(img, tf.float32))\n        imgs = tf.map_fn(lambda x: _one(x[0], x[1]),\n                         (images, flags),\n                         fn_output_signature=tf.float32)\n        return imgs, labels\n\n    return tf.cond(rnd < p_cutmix, _do_cutmix, _maybe_cutout)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:14:26.546595Z","iopub.execute_input":"2025-10-09T10:14:26.546856Z","iopub.status.idle":"2025-10-09T10:14:26.722405Z","shell.execute_reply.started":"2025-10-09T10:14:26.546829Z","shell.execute_reply":"2025-10-09T10:14:26.721626Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## tf.data pipeline (fast & shuffle-friendly)","metadata":{}},{"cell_type":"code","source":"def make_ds(X, y, batch_size, training=True):\n    # Convert labels to one-hot encoding first\n    y_one_hot = tf.one_hot(y, depth=NUM_CLASSES)\n    \n    ds = tf.data.Dataset.from_tensor_slices((X, y_one_hot))\n    if training:\n        ds = ds.shuffle(8192, seed=SEED, reshuffle_each_iteration=True)\n    ds = ds.batch(batch_size, drop_remainder=training)\n    if training:\n        # Apply augmentation and SET THE OUTPUT SIGNATURE explicitly\n        ds = ds.map(\n            lambda x, y: apply_strong_aug(x, y), \n            num_parallel_calls=AUTO\n        )\n        # Fix: Set the shape explicitly after augmentation\n        def set_shapes(images, labels):\n            images.set_shape([batch_size, 32, 32, 3])\n            labels.set_shape([batch_size, NUM_CLASSES])\n            return images, labels\n        ds = ds.map(set_shapes, num_parallel_calls=AUTO)\n    else:\n        # For validation, also set shapes\n        def set_shapes_val(images, labels):\n            images.set_shape([None, 32, 32, 3])\n            labels.set_shape([None, NUM_CLASSES])\n            return images, labels\n        ds = ds.map(set_shapes_val, num_parallel_calls=AUTO)\n    \n    ds = ds.prefetch(AUTO)\n    return ds\n\n# Rebuild datasets with the fixed pipeline\ntrain_ds = make_ds(X_train, y_train, BATCH_SIZE, training=True)\nval_ds   = make_ds(X_val,   y_val,   BATCH_SIZE, training=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T13:05:42.849266Z","iopub.execute_input":"2025-10-09T13:05:42.849579Z","iopub.status.idle":"2025-10-09T13:05:42.867415Z","shell.execute_reply.started":"2025-10-09T13:05:42.849556Z","shell.execute_reply":"2025-10-09T13:05:42.866452Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Custom NN","metadata":{}},{"cell_type":"code","source":"def se_block(x, se_ratio=0.25):\n    in_ch = x.shape[-1]\n    squeeze = layers.GlobalAveragePooling2D()(x)\n    squeeze = layers.Dense(int(in_ch * se_ratio), activation=\"relu\")(squeeze)\n    excite  = layers.Dense(in_ch, activation=\"sigmoid\")(squeeze)\n    excite  = layers.Reshape((1,1,in_ch))(excite)\n    return layers.Multiply()([x, excite])\n\ndef DSConv(x, filters, kernel_size=3, stride=1, drop=0.0):\n    # Depthwise-Separable Conv block with BN+GELU+SE\n    x = layers.DepthwiseConv2D(kernel_size, strides=stride, padding=\"same\", use_bias=False)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Activation(\"gelu\")(x)\n\n    x = layers.Conv2D(filters, 1, padding=\"same\", use_bias=False)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Activation(\"gelu\")(x)\n\n    x = se_block(x, se_ratio=0.25)\n    if drop > 0:\n        x = layers.Dropout(drop)(x)\n    return x\n\ndef Stem(x, filters=64):\n    x = layers.Conv2D(filters, 3, padding=\"same\", use_bias=False)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Activation(\"gelu\")(x)\n    return x\n\ndef downsample(x, filters):\n    x = layers.Conv2D(filters, 3, strides=2, padding=\"same\", use_bias=False)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Activation(\"gelu\")(x)\n    return x\n\ndef build_cnn(input_shape=INPUT_SHAPE, num_classes=NUM_CLASSES, drop=0.1):\n    inp = layers.Input(shape=input_shape)\n\n    # Optional per-pixel normalization layer (leave if your inputs are [0,1])\n    x = layers.Rescaling(1.0, offset=0.0)(inp)  # no change; acts as identity\n\n    x = Stem(x, 64)             # 32x32\n    x = DSConv(x,  96, 3, 1, drop=drop)\n    x = DSConv(x,  96, 3, 1, drop=drop)\n\n    x = downsample(x, 160)      # 16x16\n    x = DSConv(x, 160, 3, 1, drop=drop)\n    x = DSConv(x, 160, 3, 1, drop=drop)\n\n    x = downsample(x, 256)      # 8x8\n    x = DSConv(x, 256, 3, 1, drop=drop)\n    x = DSConv(x, 256, 3, 1, drop=drop)\n\n    x = layers.GlobalAveragePooling2D()(x)\n    x = layers.Dropout(0.2)(x)\n    out = layers.Dense(num_classes, activation=\"softmax\")(x)\n\n    model = keras.Model(inp, out, name=\"CustomCIFAR10_DSConv_SE\")\n    return model\n\nmodel = build_cnn()\nmodel.summary()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:14:31.210036Z","iopub.execute_input":"2025-10-09T10:14:31.210254Z","iopub.status.idle":"2025-10-09T10:14:32.269890Z","shell.execute_reply.started":"2025-10-09T10:14:31.210236Z","shell.execute_reply":"2025-10-09T10:14:32.269108Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Cosine annealing with warmup + AdamW","metadata":{}},{"cell_type":"code","source":"'''class WarmupCosineSchedule(keras.optimizers.schedules.LearningRateSchedule):\n    def __init__(self, base_lr, total_steps, warmup_steps=0, min_lr=1e-5):\n        super().__init__()\n        self.base_lr = base_lr\n        self.total_steps = tf.cast(total_steps, tf.float32)\n        self.warmup_steps = tf.cast(warmup_steps, tf.float32)\n        self.min_lr = min_lr\n\n    def __call__(self, step):\n        step = tf.cast(step, tf.float32)\n        # Warmup\n        if self.warmup_steps > 0:\n            warmup_lr = self.base_lr * (step / tf.maximum(1.0, self.warmup_steps))\n        else:\n            warmup_lr = self.base_lr\n\n        # Cosine decay (after warmup)\n        progress = tf.minimum(1.0, tf.maximum(0.0, (step - self.warmup_steps) / tf.maximum(1.0, self.total_steps - self.warmup_steps)))\n        cosine_lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1.0 + tf.cos(math.pi * progress))\n\n        return tf.cond(step < self.warmup_steps, lambda: warmup_lr, lambda: cosine_lr)\n\n# Steps/epoch based on training dataset\nsteps_per_epoch = int(np.ceil(len(X_train) / BATCH_SIZE))\ntotal_steps = steps_per_epoch * EPOCHS\nwarmup_ratio = 0.1  # 10% warmup\nwarmup_steps = int(total_steps * warmup_ratio)\n\nBASE_LR = 3e-3     # good starting point for AdamW on CIFAR-10 with batch 256\nMIN_LR  = 3e-5\n\nlr_schedule = WarmupCosineSchedule(base_lr=BASE_LR, total_steps=total_steps, warmup_steps=warmup_steps, min_lr=MIN_LR)\noptimizer = keras.optimizers.AdamW(learning_rate=lr_schedule, weight_decay=1e-4)'''\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T10:14:32.271521Z","iopub.execute_input":"2025-10-09T10:14:32.271767Z","iopub.status.idle":"2025-10-09T10:14:32.283077Z","shell.execute_reply.started":"2025-10-09T10:14:32.271747Z","shell.execute_reply":"2025-10-09T10:14:32.282481Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tensorflow as tf\nfrom keras import ops\nfrom keras.optimizers.schedules import LearningRateSchedule\nfrom keras.saving import register_keras_serializable\n\n@register_keras_serializable(package=\"Custom\")\nclass WarmupCosineSchedule(LearningRateSchedule):\n    \"\"\"\n    Linear warmup to base_lr over `warmup_steps`, then cosine decay to `min_lr`\n    by `total_steps`. All args are Python scalars (ints/floats).\n    \"\"\"\n\n    def __init__(self, base_lr, total_steps, warmup_steps=0, min_lr=0.0, name=None):\n        super().__init__()\n        if total_steps <= 0:\n            raise ValueError(\"total_steps must be > 0\")\n        if warmup_steps < 0 or warmup_steps >= total_steps:\n            raise ValueError(\"warmup_steps must be in [0, total_steps)\")\n        if min_lr < 0:\n            raise ValueError(\"min_lr must be >= 0\")\n\n        self.base_lr = float(base_lr)\n        self.total_steps = int(total_steps)\n        self.warmup_steps = int(warmup_steps)\n        self.min_lr = float(min_lr)\n        self.name = name or \"WarmupCosineSchedule\"\n\n        # Precompute the cosine portion length\n        self._decay_steps = self.total_steps - self.warmup_steps\n\n    def __call__(self, step):\n        step = ops.cast(step, \"float32\")\n        warmup_steps = ops.cast(self.warmup_steps, \"float32\")\n        decay_steps  = ops.cast(self._decay_steps, \"float32\")\n\n        # Linear warmup: 0 -> base_lr\n        def warmup():\n            # avoid div by zero if warmup_steps == 0\n            denom = ops.maximum(warmup_steps, 1.0)\n            return (step / denom) * self.base_lr\n\n        # Cosine decay: base_lr -> min_lr\n        def cosine():\n            # progress from 0..1 over decay steps\n            t = ops.minimum(step - warmup_steps, decay_steps) / ops.maximum(decay_steps, 1.0)\n            cosine_decay = 0.5 * (1.0 + ops.cos(ops.constant(3.141592653589793) * t))\n            return self.min_lr + (self.base_lr - self.min_lr) * cosine_decay\n\n        return ops.where(step < warmup_steps, warmup(), cosine())\n\n    def get_config(self):\n        # Must return JSON-serializable types\n        return {\n            \"base_lr\": self.base_lr,\n            \"total_steps\": self.total_steps,\n            \"warmup_steps\": self.warmup_steps,\n            \"min_lr\": self.min_lr,\n            \"name\": self.name,\n        }\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T13:03:02.702386Z","iopub.execute_input":"2025-10-09T13:03:02.702659Z","iopub.status.idle":"2025-10-09T13:03:15.749933Z","shell.execute_reply.started":"2025-10-09T13:03:02.702638Z","shell.execute_reply":"2025-10-09T13:03:15.749319Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Compile, callbacks, and train","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nimport keras\n\n# --- Training config ---\nEPOCHS = 100                         # adjust if needed\nBASE_LR = 3e-4                       # peak LR during/after warmup\nMIN_LR  = 1e-6                       # floor LR at end of cosine\nWARMUP_FRACTION = 0.10               # first 10% of steps are warmup\nWEIGHT_DECAY = 1e-4\n\n# --- Infer steps/total_steps from the dataset ---\ncard = tf.data.experimental.cardinality(train_ds).numpy()\nsteps_per_epoch = int(card) if card > 0 else None\nif steps_per_epoch is None:\n    raise ValueError(\n        \"train_ds must have a finite cardinality. \"\n        \"Make sure it's batched and not infinite/repeat().\"\n    )\ntotal_steps  = steps_per_epoch * EPOCHS\nwarmup_steps = max(1, int(WARMUP_FRACTION * total_steps))\n\n# --- Learning-rate schedule (your custom class must be defined above) ---\nlr_schedule = WarmupCosineSchedule(\n    base_lr=BASE_LR,\n    total_steps=total_steps,\n    warmup_steps=warmup_steps,\n    min_lr=MIN_LR,\n)\n\n# --- Optimizer & compile ---\noptimizer = keras.optimizers.AdamW(\n    learning_rate=lr_schedule,\n    weight_decay=WEIGHT_DECAY,\n)\n# Use 'categorical_crossentropy' if your labels are one-hot; switch to 'sparse_categorical_crossentropy' if labels are int-encoded.\nmodel.compile(\n    optimizer=optimizer,\n    loss=\"categorical_crossentropy\",\n    metrics=[\"accuracy\"],\n)\n\n# --- Callbacks ---\nckpt = keras.callbacks.ModelCheckpoint(\n    filepath=\"best_weights.h5\",\n    monitor=\"val_accuracy\",\n    mode=\"max\",\n    save_best_only=True,\n    save_weights_only=True,  # faster & avoids full optimizer serialization in experiments\n    verbose=1,\n)\nearly = keras.callbacks.EarlyStopping(\n    monitor=\"val_accuracy\",\n    mode=\"max\",\n    patience=10,\n    restore_best_weights=True,\n    verbose=1,\n)\n\n# (Optional) quick LR logger each epoch\ndef _lr_for_epoch(epoch, logs):\n    # current global step = (seen samples / batch), but we can query the optimizer variable directly:\n    opt_step = optimizer.iterations.numpy()\n    # Beware: querying inside the graph may differ; this is for a rough log.\n    # Evaluate the schedule at the current step:\n    current_lr = float(lr_schedule(opt_step).numpy())\n    print(f\"\\n[LR] step {opt_step} -> lr={current_lr:.6g}\")\n\nlr_logger = keras.callbacks.LambdaCallback(on_epoch_end=_lr_for_epoch)\n\n# --- Train ---\nhistory = model.fit(\n    train_ds,\n    validation_data=val_ds,\n    epochs=EPOCHS,\n    callbacks=[ckpt, early, lr_logger],\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-09T13:04:32.901139Z","iopub.execute_input":"2025-10-09T13:04:32.902133Z","iopub.status.idle":"2025-10-09T13:04:33.202633Z","shell.execute_reply.started":"2025-10-09T13:04:32.902108Z","shell.execute_reply":"2025-10-09T13:04:33.201620Z"}},"outputs":[],"execution_count":null}]}