{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":43873,"databundleVersionId":5585780,"sourceType":"competition"},{"sourceId":11127565,"sourceType":"datasetVersion","datasetId":6939019}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:24:28.830674Z","iopub.execute_input":"2025-03-22T14:24:28.830960Z","iopub.status.idle":"2025-03-22T14:24:29.159453Z","shell.execute_reply.started":"2025-03-22T14:24:28.830931Z","shell.execute_reply":"2025-03-22T14:24:29.155618Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! pip install lightning mlflow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:26:20.692728Z","iopub.execute_input":"2025-03-22T14:26:20.693021Z","iopub.status.idle":"2025-03-22T14:26:30.363498Z","shell.execute_reply.started":"2025-03-22T14:26:20.693000Z","shell.execute_reply":"2025-03-22T14:26:30.362653Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nfrom torch.utils.data import DataLoader, Dataset, WeightedRandomSampler\nimport h5py\nimport lightning as L\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nimport sys\nimport torchvision.transforms as transforms\nfrom PIL import Image\nimport random\nimport matplotlib.pyplot as plt\nimport logging\nfrom sklearn.model_selection import train_test_split\nimport mlflow\nfrom torch.utils.data import Dataset\nimport torch\nimport torchvision.transforms as transforms\nfrom PIL import Image\nimport random\n\nfrom transformers import ViTImageProcessor, ViTModel, ViTConfig\nfrom transformers.utils import cached_file\n\nfrom lightning.pytorch import Trainer\nfrom lightning.pytorch.loggers import MLFlowLogger\nfrom datetime import datetime\n\nimport torch.nn.functional as F\nfrom lightning.pytorch.callbacks import Callback, ModelCheckpoint, EarlyStopping","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:39:30.556224Z","iopub.execute_input":"2025-03-22T14:39:30.556565Z","iopub.status.idle":"2025-03-22T14:39:30.562221Z","shell.execute_reply.started":"2025-03-22T14:39:30.556539Z","shell.execute_reply":"2025-03-22T14:39:30.561313Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ImageClsDataset(Dataset):\n    def __init__(\n        self,\n        hdf5_file: str,\n        augment: bool = False,\n        indices: np.ndarray = None,\n        ds_total_len: float = 2,\n        image_height: int = 224,\n        image_width: int = 224,\n        device: torch.device = torch.device(\"cpu\"),\n    ):\n        \"\"\"\n        hdf5_file: the path of the hdf5 dataset\n        augment: choose if doing augmentation on image or not\n        indices: used to split the train/val dataset\n        ds_total_len: Define how many times the amount of data (data augmentation + original data) the dataset will process is the original dataset\n            e.g.: ds_total_len=2 represents that this dataset would process 2 times the amount of images than the original dataset size.\n                  In other words, to process all original images and the same number of augmented images as original images.\n        image_height | image_width: the size of the image as the input into the model\n        device: the CPU/GPU device\n        \"\"\"\n        self.hdf5_file = hdf5_file\n        self.augment = augment\n        self.indices = indices\n        self.ds_total_len = ds_total_len\n        self.image_height = image_height\n        self.image_width = image_width\n        self.device = device\n\n        with h5py.File(self.hdf5_file, \"r\") as hf:\n            self.base_length = hf[\"images\"].shape[0]\n            # load all labels into memory at one time (high efficiency)\n            self.labels = hf[\"labels\"][:]\n\n        # if without given indices, do not do any operations on dataset\n        if self.indices is None:\n            self.indices = np.arange(len(self.labels))\n\n        # if with indices, update the labels and dataset length\n        self.labels = self.labels[self.indices]\n        self.base_length = len(self.indices)\n\n        # data augmentation\n        self.augmentation_transforms = {\n            \"Horizontal Flip\": transforms.RandomHorizontalFlip(p=1.0),\n            \"Rotation\": transforms.RandomRotation(degrees=10),\n            \"Color Jitter\": transforms.ColorJitter(\n                brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1\n            ),\n            \"Random Erasing\": transforms.Compose(\n                [\n                    transforms.ToTensor(),\n                    transforms.RandomErasing(p=1.0, scale=(0.02, 0.33), value=0),\n                    transforms.ToPILImage(),\n                ]\n            ),\n        }\n\n        # resizing original images\n        self.base_transform = transforms.Resize((self.image_height, self.image_width))\n\n        # normalization(fit for ImageNet pretraining)\n        self.normalize_transform = transforms.Compose(\n            [\n                transforms.ToTensor(),\n                transforms.Normalize(\n                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]\n                ),\n            ]\n        )\n\n        # process all original images and the same number of augmented images as original images\n        self.length = self.ds_total_len * self.base_length\n\n    def random_resized_crop(self, img):\n        \"\"\"\n        Crop the image randomly\n        \"\"\"\n        crop_transform = transforms.RandomResizedCrop(\n            self.image_width, scale=(0.3, 1.0)\n        )\n\n        return crop_transform(img)\n\n    def center_included_random_crop(self, img):\n        \"\"\"\n        Ensure random cropping within the center point\n        \"\"\"\n        W, H = img.size\n        center_x, center_y = (\n            W // 2,\n            H // 2,\n        )  # calculate the center point of the original image\n\n        scale_min, scale_max = 0.3, 0.6  # allowable crop ratio\n        crop_ratio = random.uniform(scale_min, scale_max)  # generate the crop ratio\n        crop_W, crop_H = int(W * crop_ratio), int(H * crop_ratio)\n\n        # make sure the original image's center point in the crop frame\n        left_min = max(0, center_x - crop_W)\n        top_min = max(0, center_y - crop_H)\n        left_max = min(W - crop_W, center_x)\n        top_max = min(H - crop_H, center_y)\n\n        # random select the crop frame's coordinates\n        left = random.randint(left_min, left_max)\n        top = random.randint(top_min, top_max)\n        right = left + crop_W\n        bottom = top + crop_H\n\n        # resizing\n        cropped_img = img.crop((left, top, right, bottom))\n        resize_transform = transforms.Resize((self.image_height, self.image_width))\n\n        return resize_transform(cropped_img)\n\n    def apply_augmentations(self, img):\n        \"\"\"\n        Apply random augmentation to the image\n        \"\"\"\n\n        # random select the cropping method\n        crop_methods = [self.random_resized_crop, self.center_included_random_crop]\n        img = random.choice(crop_methods)(img)\n\n        # random select other augmentation methods\n        num_augmentations = random.randint(1, len(self.augmentation_transforms))\n        chosen_augmentations = random.sample(\n            list(self.augmentation_transforms.keys()), num_augmentations\n        )\n\n        for augmentation in chosen_augmentations:\n            img = self.augmentation_transforms[augmentation](img)\n\n        return img\n\n    def __len__(self):\n        \"\"\"\n        return the number of data in the dataset\n        \"\"\"\n        return self.length\n\n    def __getitem__(self, idx):\n        # judge if need do data augmentation (bool)\n        is_augmented = idx >= self.base_length\n\n        # if need data augmentation, need to map the idx into original dataset (subtract base_length)\n        # and use the original image corresponding to idx for data augmentation to generate the augmented image\n        true_idx = idx if not is_augmented else int(idx - int(self.base_length))\n\n        with h5py.File(self.hdf5_file, \"r\") as hf:\n            image = hf[\"images\"][self.indices[true_idx]]\n            label = self.labels[true_idx]\n\n        # convert NumPy -> PIL（fit with torchvision.transforms）\n        image = Image.fromarray(image)\n\n        # resizing first\n        image = self.base_transform(image)\n\n        # data augmentation\n        if is_augmented and self.augment:\n            image = self.apply_augmentations(image)\n\n        # DO normalization, we use pretained weights, so follow this normalization\n        image = self.normalize_transform(image)\n\n        # return image.to(self.device), torch.tensor(\n        #     label, dtype=torch.long, device=self.device\n        # )\n\n        return image, torch.tensor(label, dtype=torch.long)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:27:13.914438Z","iopub.execute_input":"2025-03-22T14:27:13.914766Z","iopub.status.idle":"2025-03-22T14:27:13.928761Z","shell.execute_reply.started":"2025-03-22T14:27:13.914742Z","shell.execute_reply":"2025-03-22T14:27:13.927762Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_class_weights(labels, ds_total_len=2, augment_prob=0.5):\n    \"\"\"\n    Compute sampling weights for original and augmented data.\n    :param labels: np.array, original dataset labels\n    :param ds_total_len: Dataset size multiplier (default = 2x original)\n    :param augment_prob: Probability of selecting augmented samples\n    \"\"\"\n    unique_labels, counts = np.unique(labels, return_counts=True)\n\n    # Calculate the category sampling weight (the fewer category samples, the greater the weight)\n    class_weights = {label: 1.0 / count for label, count in zip(unique_labels, counts)}\n\n    # Original data weight\n    original_sample_weights = np.array([class_weights[label] for label in labels])\n\n    # Copy labels so that augmented_labels and original_labels keep the same category ratio\n    augmented_labels = np.tile(labels, max(int(ds_total_len - 1), 1))\n    augmented_sample_weights = np.array([class_weights[label] for label in augmented_labels])\n\n    # Calculate how much original` vs. augmented should be in batch\n    num_original = len(original_sample_weights)\n    num_augmented = len(augmented_sample_weights)\n    total_samples = num_original + num_augmented\n\n    # Control the proportion of augmented data in batch\n    if num_augmented == 0:\n        augmented_ratio = 0\n    else:\n        augmented_ratio = augment_prob * total_samples / num_augmented\n\n    original_ratio = (1 - augment_prob) * total_samples / num_original\n\n    final_sample_weights = np.concatenate([\n        original_sample_weights * original_ratio,  \n        augmented_sample_weights * augmented_ratio \n    ])\n\n    return final_sample_weights","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:27:26.153598Z","iopub.execute_input":"2025-03-22T14:27:26.153896Z","iopub.status.idle":"2025-03-22T14:27:26.159933Z","shell.execute_reply.started":"2025-03-22T14:27:26.153875Z","shell.execute_reply":"2025-03-22T14:27:26.158938Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ImageClsDataModule(L.LightningDataModule):\n    def __init__(\n        self,\n        hdf5_file: str,\n        batch_size:int=32,\n        train_ratio: float=0.8,\n        ds_total_len:float=2.0,\n        initial_augment_prob:float=0.8,\n        final_augment_prob:float=0.3,\n        num_epochs:int=50,\n        image_height:int=224,\n        image_width:int=224,\n        device: torch.device = torch.device(\"cpu\"),\n        num_workers: int = 0,\n    ):\n        \"\"\"\n        Params:\n            hdf5_file: the path of the hdf5 dataset\n            batch_size: size of batch\n            train_ratio: train/val dataset ratio\n            ds_total_len: Define how many times the amount of data (data augmentation + original data) the dataset will process is the original dataset\n                e.g.: ds_total_len=2 represents that this dataset would process 2 times the amount of images than the original dataset size.\n                    In other words, to process all original images and the same number of augmented images as original images.\n            initial_augment_prob: Initial data augmentation ratio\n            final_augment_prob: final data augmentation ratio\n            num_epochs: the number of epochs\n            device: USE GPU/CPU deivce\n            num_worker: the number of workers for dataloader\n        \"\"\"\n        super().__init__()\n        self.hdf5_file = hdf5_file\n        self.batch_size = batch_size\n        self.train_ratio = train_ratio\n        self.ds_total_len = ds_total_len\n        self.initial_augment_prob = initial_augment_prob\n        self.final_augment_prob = final_augment_prob\n        self.num_epochs = num_epochs\n        self.device = device\n        self.num_workers = num_workers\n        self.image_height=image_height\n        self.image_width=image_width\n\n        self.current_epoch = 0\n\n    def setup(self, stage=None):\n        \"\"\"\n        Split dataset into train and testsets\n        \"\"\"\n\n        with h5py.File(self.hdf5_file, \"r\") as hf:\n            labels = np.array(hf[\"labels\"])\n            num_samples = len(labels)\n\n        # ensure the label balancing\n        train_indices, val_indices = train_test_split(\n            np.arange(num_samples),\n            test_size=1 - self.train_ratio,\n            stratify=labels,\n            random_state=42,\n        )\n\n        # create train and val dataset\n        self.train_dataset = ImageClsDataset(\n            hdf5_file=self.hdf5_file,\n            indices=train_indices,\n            augment=True,\n            ds_total_len=self.ds_total_len,\n            device=self.device,\n            image_height=self.image_height,\n            image_width=self.image_width\n        )\n\n        self.val_dataset = ImageClsDataset(\n            hdf5_file=self.hdf5_file,\n            indices=val_indices,\n            augment=False,\n            # hard coded as 1 because onlu use original images in val set\n            ds_total_len=1,\n            device=self.device,\n            image_height=self.image_height,\n            image_width=self.image_width\n        )\n\n    def update_augment_prob(self):\n        \"\"\"\n        update the augmented ratio for each epoch\n\n        use more augmented images in early-stage epoches.\n        use fewer augmented images in late-stage epoches.\n        \"\"\"\n        # sync the current training epoch\n        if self.trainer is not None:\n            self.current_epoch = self.trainer.current_epoch\n\n        self.augment_prob = self.initial_augment_prob - (\n            self.current_epoch / self.num_epochs\n        ) * (self.initial_augment_prob - self.final_augment_prob)\n\n        # print(f\"📢 Epoch {self.current_epoch}: Augment Prob = {self.augment_prob:.2f}\")\n        # logging.info(\n        #     f\"📢 Epoch {self.current_epoch}: Augment Prob = {self.augment_prob:.2f}\"\n        # )\n\n        return None\n\n    def train_dataloader(self):\n        \"\"\"\n        define how the dataloader created for each epoch\n        requirements:\n            set `reload_dataloaders_every_n_epochs=1` in L.trainer\n        \"\"\"\n        # calculate the current epoch's augment ratio\n        self.update_augment_prob()\n        # self.train_dataset.update_augment_prob(self.augment_prob)\n\n        # re-build the sampler\n        sample_weights = compute_class_weights(\n            labels=self.train_dataset.labels,\n            ds_total_len=self.ds_total_len,\n            augment_prob=self.augment_prob,\n        )\n\n        train_sampler = WeightedRandomSampler(\n            sample_weights, num_samples=len(sample_weights), replacement=True\n        )\n        # print(\n        #     f\"TRAIN_DATALODER DATAMODULE FLAG (current epoch num): {self.current_epoch+1}\"\n        # )\n\n        # logging.info(\n        #     f\"TRAIN_DATALODER DATAMODULE FLAG (current epoch num): {self.current_epoch+1}\"\n        # )\n\n        return DataLoader(\n            self.train_dataset,\n            batch_size=self.batch_size,\n            sampler=train_sampler,\n            num_workers=self.num_workers,\n            pin_memory=True,\n        )\n\n    def val_dataloader(self):\n        \"\"\"\n        define how the validation dataloader created for each epoch\n        requirements:\n            set `reload_dataloaders_every_n_epochs=1` in L.trainer\n        \"\"\"\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,\n            num_workers=self.num_workers,\n            pin_memory=True,\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:27:46.634431Z","iopub.execute_input":"2025-03-22T14:27:46.634744Z","iopub.status.idle":"2025-03-22T14:27:46.645760Z","shell.execute_reply.started":"2025-03-22T14:27:46.634720Z","shell.execute_reply":"2025-03-22T14:27:46.644914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# enable GPU\nif torch.cuda.is_available():\n    print(\"✅ GPU is available!\")\n    print(\"Device name:\", torch.cuda.get_device_name(0))\nelse:\n    print(\"❌ GPU not available. Using CPU.\")\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Using device:\", device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:30:15.824604Z","iopub.execute_input":"2025-03-22T14:30:15.824927Z","iopub.status.idle":"2025-03-22T14:30:15.831182Z","shell.execute_reply.started":"2025-03-22T14:30:15.824900Z","shell.execute_reply":"2025-03-22T14:30:15.830372Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# add kaggle path\nkaggle_output=\"/kaggle/working/\"\n\nos.makedirs(\"/kaggle/working/vit-base-patch16-224-in21k\", exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:33:09.109734Z","iopub.execute_input":"2025-03-22T14:33:09.110050Z","iopub.status.idle":"2025-03-22T14:33:09.113715Z","shell.execute_reply.started":"2025-03-22T14:33:09.110025Z","shell.execute_reply":"2025-03-22T14:33:09.112893Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load pre defined parameters\nlocal_config_path = \"/kaggle/working/vit-base-patch16-224-in21k\"\n\n# label encoder\nlabel_encoder = {\n    \"dot\": 0,\n    \"scatter\": 1,\n    \"horizontal_bar\": 2,\n    \"line\": 3,\n    \"vertical_bar\": 4,\n}\n\n# label decoder\nlabel_decoder = {v: k for k, v in label_encoder.items()}\n\nid2label = {idx: label for idx, label in label_decoder.items()}\nlabel2id = {label: idx for idx, label in label_encoder.items()}\n\n# Dataset parameters\nimage_cls_height = 320\nimage_cls_width = 320\n\n# dataset size multiplier\nds_multiplier = 2\n\n# Datamodule Parameters\ninitial_augment_prob = 0.8\nfinal_augment_prob = 0.2\n\n# Dataloader parameters\nnum_workers = 4\nbatch_size = 32\nnum_epochs = 30\n\n# data path\nhdf5_file = \"/kaggle/input/benetech-ds/ImageTransformation.h5\"\n\ntimestamp = datetime.now().strftime(\"%Y%m%d-%H%M%S\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:35:09.475083Z","iopub.execute_input":"2025-03-22T14:35:09.475415Z","iopub.status.idle":"2025-03-22T14:35:09.480759Z","shell.execute_reply.started":"2025-03-22T14:35:09.475389Z","shell.execute_reply":"2025-03-22T14:35:09.479876Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"pretrained model path: {local_config_path}\")\n\nprint(f\"label encoder: {label_encoder}\")\nprint(f\"label decoder: {label_decoder}\")\nprint(f\"id2label: {id2label}\")\nprint(f\"label2id: {label2id}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:35:20.292480Z","iopub.execute_input":"2025-03-22T14:35:20.292841Z","iopub.status.idle":"2025-03-22T14:35:20.298608Z","shell.execute_reply.started":"2025-03-22T14:35:20.292813Z","shell.execute_reply":"2025-03-22T14:35:20.297799Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mlflow_logger = MLFlowLogger(\n    experiment_name=\"ViT-image-recognition\",\n    tracking_uri=\"/kaggle/working/mlruns\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:36:39.409674Z","iopub.execute_input":"2025-03-22T14:36:39.409968Z","iopub.status.idle":"2025-03-22T14:36:39.418060Z","shell.execute_reply.started":"2025-03-22T14:36:39.409948Z","shell.execute_reply":"2025-03-22T14:36:39.417153Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load VIT model's config and processor\n\nconfig = ViTConfig.from_pretrained(\"google/vit-base-patch16-224-in21k\",cache_dir=local_config_path)\nprocessor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k', cache_dir=local_config_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:36:51.327252Z","iopub.execute_input":"2025-03-22T14:36:51.327578Z","iopub.status.idle":"2025-03-22T14:36:51.872955Z","shell.execute_reply.started":"2025-03-22T14:36:51.327554Z","shell.execute_reply":"2025-03-22T14:36:51.872019Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"type(config)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:37:11.113810Z","iopub.execute_input":"2025-03-22T14:37:11.114130Z","iopub.status.idle":"2025-03-22T14:37:11.119212Z","shell.execute_reply.started":"2025-03-22T14:37:11.114104Z","shell.execute_reply":"2025-03-22T14:37:11.118293Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# add customized paramaters into config file\n\nconfig.num_labels = len(id2label)\nconfig.id2label = id2label\nconfig.label2id = label2id\nconfig.hidden_dropout_prob = 0.1 \nconfig.image_size = 320\nconfig.patch_size = 16\n# 320/16=20,\n# input token num: 320*320 / 16*16 = 20*20=400 \n# add CLS: 400+1=401\n# each token seq len: 16*16*3=768\n\nmodel_name = \"google/vit-base-patch16-224-in21k\"\nconfig_path = cached_file(model_name, \"config.json\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:37:20.218196Z","iopub.execute_input":"2025-03-22T14:37:20.218513Z","iopub.status.idle":"2025-03-22T14:37:20.707283Z","shell.execute_reply.started":"2025-03-22T14:37:20.218489Z","shell.execute_reply":"2025-03-22T14:37:20.706412Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:37:24.532635Z","iopub.execute_input":"2025-03-22T14:37:24.532956Z","iopub.status.idle":"2025-03-22T14:37:24.538970Z","shell.execute_reply.started":"2025-03-22T14:37:24.532928Z","shell.execute_reply":"2025-03-22T14:37:24.538074Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load pretrained VIT model\n\nmodel = ViTModel.from_pretrained(\n    \"google/vit-base-patch16-224-in21k\",\n    config=config,\n    # I want fine tuning the model with Higher Resolution, change the input size from 224 to 320\n    # so set ignore mismatched_size right here\n    ignore_mismatched_sizes=True,\n    cache_dir=local_config_path,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:37:37.418327Z","iopub.execute_input":"2025-03-22T14:37:37.418707Z","iopub.status.idle":"2025-03-22T14:37:40.403297Z","shell.execute_reply.started":"2025-03-22T14:37:37.418678Z","shell.execute_reply":"2025-03-22T14:37:40.402431Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:37:49.504597Z","iopub.execute_input":"2025-03-22T14:37:49.504911Z","iopub.status.idle":"2025-03-22T14:37:49.510455Z","shell.execute_reply.started":"2025-03-22T14:37:49.504886Z","shell.execute_reply":"2025-03-22T14:37:49.509605Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for name, param in model.named_parameters():\n    if param.requires_grad:\n        print(f\"  ✅ {name}: {param.shape}\")\n    else:\n        print(f\"  ❌ {name}: (frozen)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:37:58.461418Z","iopub.execute_input":"2025-03-22T14:37:58.461751Z","iopub.status.idle":"2025-03-22T14:37:58.493444Z","shell.execute_reply.started":"2025-03-22T14:37:58.461722Z","shell.execute_reply":"2025-03-22T14:37:58.492729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# freeze some layers\n\n# fine tuning embedding layer\nfor param in model.embeddings.parameters():\n    param.requires_grad = True\n\n# freeze the first 6 layers(total 12)\nfor name, param in model.encoder.named_parameters():\n    if any(f\"layer.{i}\" in name for i in range(6)):  \n        param.requires_grad = False\n\n# 🔥 fine tuning the last 6 layers\nfor name, param in model.encoder.named_parameters():\n    if any(f\"layer.{i}\" in name for i in range(6, 12)):  \n        param.requires_grad = True\n\n# 🔥 fine tune pooling layer\nfor param in model.pooler.parameters():\n    param.requires_grad = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:38:35.728624Z","iopub.execute_input":"2025-03-22T14:38:35.728929Z","iopub.status.idle":"2025-03-22T14:38:35.735689Z","shell.execute_reply.started":"2025-03-22T14:38:35.728905Z","shell.execute_reply":"2025-03-22T14:38:35.734881Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\ntotal_params = sum(p.numel() for p in model.parameters())\n\nprint(f\"Trainable Params: {trainable_params:,} / {total_params:,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:38:41.534262Z","iopub.execute_input":"2025-03-22T14:38:41.534596Z","iopub.status.idle":"2025-03-22T14:38:41.541607Z","shell.execute_reply.started":"2025-03-22T14:38:41.534570Z","shell.execute_reply":"2025-03-22T14:38:41.540812Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for name, param in model.named_parameters():\n    if param.requires_grad:\n        print(f\"  ✅ {name}: {param.shape}\")\n    else:\n        print(f\"  ❌ {name}: (frozen)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:38:47.529641Z","iopub.execute_input":"2025-03-22T14:38:47.529964Z","iopub.status.idle":"2025-03-22T14:38:47.546531Z","shell.execute_reply.started":"2025-03-22T14:38:47.529937Z","shell.execute_reply":"2025-03-22T14:38:47.544472Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load pretrained VIT model\n\nmodel = ViTModel.from_pretrained(\n    \"google/vit-base-patch16-224-in21k\",\n    config=config,\n    # I want fine tuning the model with Higher Resolution, change the input size from 224 to 320\n    # so set ignore mismatched_size right here\n    ignore_mismatched_sizes=True,\n    cache_dir=local_config_path,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:40:06.765723Z","iopub.execute_input":"2025-03-22T14:40:06.766033Z","iopub.status.idle":"2025-03-22T14:40:07.672858Z","shell.execute_reply.started":"2025-03-22T14:40:06.766007Z","shell.execute_reply":"2025-03-22T14:40:07.672021Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# initilize DataModule\n\ndatamodule = ImageClsDataModule(\n    hdf5_file=hdf5_file,\n    batch_size=batch_size,\n    ds_total_len=ds_multiplier,\n    initial_augment_prob=initial_augment_prob,\n    final_augment_prob=final_augment_prob,\n    num_epochs=num_epochs,\n    device=device,\n    num_workers=num_workers,\n    image_height=image_cls_height,\n    image_width=image_cls_width\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:40:10.331280Z","iopub.execute_input":"2025-03-22T14:40:10.331599Z","iopub.status.idle":"2025-03-22T14:40:10.335657Z","shell.execute_reply.started":"2025-03-22T14:40:10.331576Z","shell.execute_reply":"2025-03-22T14:40:10.334845Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ViTClassifier(L.LightningModule):\n    def __init__(self, pretrained_model, config, num_class, lr=None):\n        \"\"\"\n        Params:\n            pretrained_model: The loaded ViT model\n            config:(transformers.models.vit.configuration_vit.ViTConfig) the model's config file\n            num_class: the numer of classes\n            lr: learning rate\n        \"\"\"\n        super().__init__()\n        # saving the information in checkpoints and YAML files\n        self.save_hyperparameters()\n\n        self.model = pretrained_model\n        self.config = config\n        self.num_class = num_class\n        self.learning_rate = (\n            lr if lr is not None else self.hparams.get(\"learning_rate\", 1e-4)\n        )\n\n        # define what layers in the pretrained model should be fine-tuning\n        # fine tune the embedding (positional embedding) layer\n        for param in self.model.embeddings.parameters():\n            param.requires_grad = True\n\n        # freeze the first 9 layers(total 12)\n        for name, param in self.model.encoder.named_parameters():\n            if any(f\"layer.{i}\" in name for i in range(9)):\n                param.requires_grad = False\n\n        # fine tuning the last 3 layers\n        for name, param in self.model.encoder.named_parameters():\n            if any(f\"layer.{i}\" in name for i in range(9, 12)):\n                param.requires_grad = True\n\n        # fine tune pooling layer\n        for param in self.model.pooler.parameters():\n            param.requires_grad = True\n\n        # define the classfier head\n        self.classifier = nn.Linear(\n            in_features=config.hidden_size, out_features=self.num_class\n        )\n\n    def forward(self, x):\n        outputs = self.model(x)\n        # In VIT model, we rely on CLS token to do classification\n        cls_token = outputs.last_hidden_state[:, 0, :]\n        # classifier head\n        logits = self.classifier(cls_token)\n\n        return logits\n\n    def training_step(self, batch, batch_idx):\n        images, labels = batch\n        images = images.to(self.device)\n        labels = labels.to(self.device)\n\n        logits = self(images)\n        loss = F.cross_entropy(logits, labels)\n        self.log(\"train_loss\", loss, prog_bar=True, on_step=False, on_epoch=True)\n\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        images, labels = batch\n        images = images.to(self.device)\n        labels = labels.to(self.device)\n        logits = self(images)\n        val_loss = F.cross_entropy(logits, labels)\n        acc = (logits.argmax(dim=1) == labels).float().mean()\n        self.log(\"val_loss\", val_loss, prog_bar=True, on_step=False, on_epoch=True)\n        self.log(\"val_acc\", acc, prog_bar=True, on_step=False, on_epoch=True)\n\n        return {\"val_loss\": val_loss, \"val_acc\": acc}\n\n    def predict_step(self, batch, batch_idx):\n        images, _ = batch\n        logits = self(images)\n        preds = torch.argmax(logits, dim=1)\n\n        return preds\n\n    def on_train_epoch_end(self):\n\n        train_loss = self.trainer.callback_metrics.get(\"train_loss\")\n        if train_loss is not None:\n            print(f\"Train Epoch {self.current_epoch + 1} - Avg Loss: {train_loss:.4f}\")\n\n        # mlflow\n        self.logger.log_metrics(\n            {\n                \"manual_record_train_loss\": train_loss.item()\n            },\n            step=self.current_epoch,\n        )\n\n    def on_validation_epoch_end(self):\n\n        val_loss = self.trainer.callback_metrics.get(\"val_loss\")\n        val_acc = self.trainer.callback_metrics.get(\"val_acc\")\n        if val_loss is not None and val_acc is not None:\n            print(\n                f\"Val Epoch {self.current_epoch + 1} - Loss: {val_loss:.4f}, Acc: {val_acc:.4f}\"\n            )\n\n        # mlflow\n        self.logger.log_metrics(\n            {\n                \"manual_record_val_loss\": val_loss.item(),\n                \"manual_record_val_acc\": val_acc.item(),\n            },\n            step=self.current_epoch,\n        )\n\n    def on_fit_start(self):\n        print(f\"✅ Model is on device: {next(self.parameters()).device}\")\n        if self.logger:\n            self.logger.log_hyperparams(self.hparams)\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)\n\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            mode=\"min\",\n            factor=0.5,\n            patience=2, # trigger lr scheduler if without imporvement 3 times\n            verbose=True,\n        )\n\n        # decrease lr based off val loss automatically\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"monitor\": \"val_loss\",\n                \"interval\": \"epoch\",\n                \"frequency\": 1, # do step each 1 epoch\n            },\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:42:38.812149Z","iopub.execute_input":"2025-03-22T14:42:38.812472Z","iopub.status.idle":"2025-03-22T14:42:38.826261Z","shell.execute_reply.started":"2025-03-22T14:42:38.812447Z","shell.execute_reply":"2025-03-22T14:42:38.825393Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"VitModel=ViTClassifier(pretrained_model=model.to(device),config=config,num_class=len(id2label))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:42:39.529137Z","iopub.execute_input":"2025-03-22T14:42:39.529488Z","iopub.status.idle":"2025-03-22T14:42:39.563167Z","shell.execute_reply.started":"2025-03-22T14:42:39.529458Z","shell.execute_reply":"2025-03-22T14:42:39.562492Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ViTCallback(Callback):\n    def on_train_epoch_end(self, trainer, pl_module):\n        lr = trainer.optimizers[0].param_groups[0]['lr']\n        print(f\"📉 Learning Rate after Epoch {trainer.current_epoch + 1}: {lr:.6f}\")\n\n\n# checkpoint\ncheckpoint_callback = ModelCheckpoint(\n    monitor=\"val_acc\",\n    mode=\"max\",\n    save_top_k=1,\n    dirpath=f\"ViTCheckpoints/{timestamp}/\", \n    filename=\"vit-best-{epoch:02d}-{val_acc:.4f}\",\n    save_weights_only=True,\n    verbose=True,\n)\n\n# early stopping\nearly_stop_callback = EarlyStopping(\n    monitor=\"val_acc\",\n    patience=5, # stop the training if without improvement 5 epochs continusly\n    mode=\"max\",\n    verbose=True,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:42:40.288551Z","iopub.execute_input":"2025-03-22T14:42:40.288859Z","iopub.status.idle":"2025-03-22T14:42:40.294654Z","shell.execute_reply.started":"2025-03-22T14:42:40.288837Z","shell.execute_reply":"2025-03-22T14:42:40.293799Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer = L.Trainer(\n    max_epochs=num_epochs,\n    callbacks=[ViTCallback(), checkpoint_callback, early_stop_callback],\n    logger=mlflow_logger,\n    enable_checkpointing=True,\n    log_every_n_steps=1,\n    gradient_clip_val=1.0,\n    enable_model_summary=True,\n    accumulate_grad_batches=4\n)\n\n\ntrainer.fit(VitModel,datamodule)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-22T14:43:11.649108Z","iopub.execute_input":"2025-03-22T14:43:11.649455Z","execution_failed":"2025-03-23T02:24:22.676Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import shutil\n\n# 打包整个 /kaggle/working 到一个 zip 文件中（可以包含模型、日志等）\nshutil.make_archive(\"/kaggle/working/kaggle_output\", 'zip', \"/kaggle/working\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T02:44:41.420702Z","iopub.execute_input":"2025-03-23T02:44:41.420888Z","iopub.status.idle":"2025-03-23T02:45:15.445812Z","shell.execute_reply.started":"2025-03-23T02:44:41.420869Z","shell.execute_reply":"2025-03-23T02:45:15.444956Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}