{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":94147,"databundleVersionId":11390004,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11535317,"sourceType":"datasetVersion","datasetId":7234715},{"sourceId":11592169,"sourceType":"datasetVersion","datasetId":7269146}],"dockerImageVersionId":31011,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport time\nimport random\nfrom pathlib import Path\n\nimport cv2\nimport h5py\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torchvision\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.models as models\nimport torchvision.transforms as T\nfrom sklearn.model_selection import train_test_split\nfrom scipy.stats import spearmanr\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport imgaug\nimport wandb  \nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n#I had the code rewritten by AI to make it more readable, it ended up being a bit long\n\n#ckpt file from https://github.com/ozanciga/self-supervised-histopathology/releases/tag/tenpercent","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:12.903972Z","iopub.execute_input":"2025-05-06T12:51:12.904680Z","iopub.status.idle":"2025-05-06T12:51:12.910483Z","shell.execute_reply.started":"2025-05-06T12:51:12.904650Z","shell.execute_reply":"2025-05-06T12:51:12.909694Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport random\nimport h5py\nimport numpy as np\nimport pandas as pd\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.models as models\nimport torchvision.transforms as T\n\nfrom scipy.stats import spearmanr\nfrom tqdm import tqdm\nimport imgaug\nimport wandb\n\nCONFIG = {\n    \"seed\": 42,\n    \"data_path\": \"/kaggle/input/el-hackathon-2025\",\n    \"output_dir\": \"/kaggle/working/\",\n    \"batch_size\": 32,\n    \"num_workers\": 6,\n    \"learning_rate\": 0.003,\n    \"weight_decay\": 1e-1,\n    \"scheduler_step_size\": 5,\n    \"scheduler_gamma\": 0.1,\n    \"num_classes\": 35,\n    \"image_size\": (162, 162),\n    \"patch_size\": 54,\n    \"max_epochs\": 5, # Best scoring epoch will be used for submission.csv. Epoch 4 performed the best among 10 epochs.\n    \"patience\": 5,  # Number of epochs to wait for improvement before early stopping\n    \"min_delta\": 0.001,  # Minimum change to qualify as improvement\n    \"save_best_only\": True,\n    \"checkpoint_epochs\": [], # epoch 3 performed best\n    \"use_wandb\": False,  # Set to False if you don't want to use wandb\n    \"model_type\": \"resnet18\",\n    \"mixed_precision\": False,  # Use mixed precision training\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:12.911642Z","iopub.execute_input":"2025-05-06T12:51:12.912009Z","iopub.status.idle":"2025-05-06T12:51:12.927666Z","shell.execute_reply.started":"2025-05-06T12:51:12.911988Z","shell.execute_reply":"2025-05-06T12:51:12.926930Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def set_seed(seed=42):\n    \"\"\"Set seeds for reproducibility\"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    imgaug.seed(seed)\n\n\ndef get_device():\n    \"\"\"Get the appropriate device\"\"\"\n    if torch.cuda.is_available():\n        return torch.device(\"cuda\")\n    elif torch.backends.mps.is_available():\n        return torch.device(\"mps\")  # For Apple Silicon\n    else:\n        return torch.device(\"cpu\")\n\n\ndef spearman_rank_correlation(x, y):\n    \"\"\"Calculate Spearman rank correlation\"\"\"\n    # Handle edge cases\n    if np.all(x == x[0]) or np.all(y == y[0]):\n        return 0.0\n    return spearmanr(x, y)[0]\n\ndef spearman_corr(preds, targets):\n    \"\"\"Calculate mean Spearman correlation across all samples\"\"\"\n    correlations = []\n    for i in range(len(preds)):\n        corr = spearman_rank_correlation(preds[i], targets[i])\n        # Only count valid correlations\n        if not np.isnan(corr):\n            correlations.append(corr)\n    return np.mean(correlations)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:12.928822Z","iopub.execute_input":"2025-05-06T12:51:12.929038Z","iopub.status.idle":"2025-05-06T12:51:12.936455Z","shell.execute_reply.started":"2025-05-06T12:51:12.929024Z","shell.execute_reply":"2025-05-06T12:51:12.935740Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class HackhathonDataset(Dataset):\n    \"\"\"Dataset class for the hackathon competition\"\"\"\n\n    def __init__(self, data_path, transform=None, mode=\"train\"):\n        self.data_path = data_path\n        self.materials = []\n        self.transform = transform\n\n        # Define slides for train and validation sets\n        train_slides = [\"S_1\",\"S_2\",\"S_3\", \"S_4\", \"S_5\"]\n        val_slide = [\"S_6\"]\n        test_slide = [\"S_7\"]\n        self.mode = mode\n\n        slide_list = train_slides if mode == \"train\" else val_slide if mode == \"val\" else test_slide\n\n        with h5py.File(f\"{self.data_path}/elucidata_ai_challenge_data.h5\", \"r\") as h5file:\n            images_group = \"images/Train\" if mode != \"test\" else \"images/Test\"\n            spots_group = \"spots/Train\" if mode != \"test\" else \"spots/Test\"\n\n            train_images = h5file[images_group]\n            train_spots = h5file[spots_group]\n\n            for slide_name in tqdm(slide_list, desc=f\"Loading {mode} data\"):\n                if slide_name in train_images.keys():\n                    image = np.array(train_images[slide_name])\n                    spots = np.array(train_spots[slide_name])\n                    df = pd.DataFrame(spots)\n                    self._split_into_patches(image, df, CONFIG[\"patch_size\"])\n\n        print(f\"{len(self.materials)} patches initialized for {mode} set\")\n\n    def __len__(self):\n        return len(self.materials)\n\n    def __getitem__(self, idx):\n        image, stats = self.materials[idx]\n\n        if self.transform:\n            image = self.transform(image)\n\n        stats = torch.tensor(stats[2:], dtype=torch.float32)\n\n        return image, stats\n\n    def _split_into_patches(self, arr, df, patch_size):\n        \"\"\"Split the image into patches centered on spot coordinates\"\"\"\n        h, w, c = arr.shape\n\n        for idx in range(len(df)):\n            row = df.iloc[idx]\n            x, y = int(row[\"x\"]), int(row[\"y\"])\n\n            half_size = patch_size // 2\n\n            # Ensure patches don't go outside image boundaries\n            y_min = max(y - half_size, 0)\n            y_max = min(y + half_size, h)\n            x_min = max(x - half_size, 0)\n            x_max = min(x + half_size, w)\n\n            patch = arr[y_min:y_max, x_min:x_max, :]\n\n            # Only include complete patches\n            if patch.shape[0] == patch_size and patch.shape[1] == patch_size:\n                self.materials.append([patch, row])\n            else:\n                # Handle incomplete patches by padding\n                padded_patch = np.zeros((patch_size, patch_size, c), dtype=patch.dtype)\n                padded_patch[:patch.shape[0], :patch.shape[1], :] = patch\n                self.materials.append([padded_patch, row])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:12.937165Z","iopub.execute_input":"2025-05-06T12:51:12.937399Z","iopub.status.idle":"2025-05-06T12:51:12.953961Z","shell.execute_reply.started":"2025-05-06T12:51:12.937383Z","shell.execute_reply":"2025-05-06T12:51:12.953343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DifferentiableSpearmanLoss(nn.Module):\n    \"\"\"Differentiable approximation of Spearman correlation loss\"\"\"\n\n    def __init__(self, regularization_strength=1.0):\n        super().__init__()\n        self.regularization_strength = regularization_strength\n\n    def forward(self, y_pred, y_true):\n        y_pred = y_pred.float()\n        y_true = y_true.float()\n\n        # Calculate soft ranks\n        pred_rank = self._soft_rank(y_pred)\n        true_rank = self._soft_rank(y_true)\n\n        # Normalize ranks\n        pred_rank = F.normalize(pred_rank, dim=1)\n        true_rank = F.normalize(true_rank, dim=1)\n\n        # Calculate correlation\n        spearman = torch.sum(pred_rank * true_rank, dim=1)\n        return 1 - spearman.mean()\n\n    def _soft_rank(self, x, regularization_strength=None):\n        if regularization_strength is None:\n            regularization_strength = self.regularization_strength\n\n        x = x.unsqueeze(-1)  # [batch, n, 1]\n        diff = x - x.transpose(-1, -2)  # [batch, n, n]\n        P = torch.sigmoid(-regularization_strength * diff)  # pairwise comparisons\n        ranks = P.sum(dim=-1)  # approximate ranks\n        return ranks\n\n\nclass CombinedLoss(nn.Module):\n    \"\"\"Combined loss function using L1 and Spearman correlation\"\"\"\n\n    def __init__(self, alpha=0.5, regularization_strength=1.0):\n        super().__init__()\n        self.l1 = nn.L1Loss()\n        self.spearman = DifferentiableSpearmanLoss(regularization_strength)\n        self.alpha = alpha\n\n    def forward(self, y_pred, y_true):\n        l1_loss = self.l1(y_pred, y_true)\n        spearman_loss = self.spearman(y_pred, y_true)\n        return l1_loss + self.alpha * spearman_loss\n\ndef create_model(model_type, num_classes):\n    \"\"\"Create and initialize the model\"\"\"\n    if model_type == \"resnet18\":\n        #model = models.wide_resnet101_2(weights=models.Wide_ResNet101_2_Weights.IMAGENET1K_V1)\n        #model.fc = nn.Linear(model.fc.in_features, num_classes)\n        MODEL_PATH = '/kaggle/input/ckpt-file/tenpercent_resnet18.ckpt'\n        RETURN_PREACTIVATION = False  # return features from the model, if false return classification logits\n        NUM_CLASSES = 35  # only used if RETURN_PREACTIVATION = False\n\n        def load_model_weights(model, weights):\n\n            model_dict = model.state_dict()\n            weights = {k: v for k, v in weights.items() if k in model_dict}\n            if weights == {}:\n                print('No weight could be loaded..')\n            model_dict.update(weights)\n            model.load_state_dict(model_dict)\n\n            return model\n\n        model = models.resnet18(pretrained=False)\n\n        state = torch.load(MODEL_PATH, map_location='cuda:0', weights_only=False)\n\n        state_dict = state['state_dict']\n        for key in list(state_dict.keys()):\n            state_dict[key.replace('model.', '').replace('resnet.', '')] = state_dict.pop(key)\n\n        model = load_model_weights(model, state_dict)\n\n        if RETURN_PREACTIVATION:\n            model.fc = torch.nn.Sequential()\n        else:\n            model.fc = torch.nn.Linear(model.fc.in_features, NUM_CLASSES)\n\n\n    elif model_type == \"DeepCMorph\":\n        pass\n        # Defining the model and specifying the number of target classes:\n        # 41 for combined datasets, 32 for TCGA, 9 for CRC\n        #model = DeepCMorph(num_classes=35)\n\n        # Loading model weights corresponding to the network trained on combined datasets\n        # Possible 'dataset' values: TCGA, TCGA_REGULARIZED, CRC, COMBINED\n        #model.load_weights(dataset=\"CRC\")\n    elif model_type == \"resnet50\":\n\n        model = models.resnet50(weights=None)\n\n        model = load_pretrained_weights(model, \"31\")\n\n        model.fc = nn.Linear(model.fc.in_features, num_classes)\n\n\n    elif model_type == \"swin\":\n\n        model = models.swin_v2_b(weights=models.Swin_V2_B_Weights.IMAGENET1K_V1)\n\n        model.head = nn.Linear(model.head.in_features, num_classes)\n\n        return model\n\n    elif model_type == \"efficientnet_b3\":\n\n        model = models.efficientnet_b7(weights=models.EfficientNet_B7_Weights.IMAGENET1K_V1)\n\n        model.classifier = nn.Linear(model.classifier[1].in_features, num_classes)\n\n    elif model_type == \"convnext_large\":\n\n        model = models.convnext_large(weights=models.ConvNeXt_Large_Weights.IMAGENET1K_V1)\n\n        in_features = model.classifier[2].in_features\n\n        model.classifier[2] = nn.Linear(in_features, num_classes)\n\n    else:\n\n        raise ValueError(f\"Unsupported model type: {model_type}\")\n\n    return model\n\n\ndef get_transforms():\n    \"\"\"Get data transformations for training and validation\"\"\"\n    train_transform = T.Compose([\n        T.ToTensor(),\n        T.Resize(CONFIG[\"image_size\"]),\n        T.RandomApply([T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1)], p=0.8),\n        T.RandomHorizontalFlip(p=0.5),\n        T.RandomVerticalFlip(p=0.5),\n        T.RandomRotation(degrees=45),\n        T.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),\n    ])\n\n    val_transform = T.Compose([\n        T.ToTensor(),\n        T.Resize(CONFIG[\"image_size\"]),\n    ])\n\n    return train_transform, val_transform","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:13.001991Z","iopub.execute_input":"2025-05-06T12:51:13.002617Z","iopub.status.idle":"2025-05-06T12:51:13.015710Z","shell.execute_reply.started":"2025-05-06T12:51:13.002600Z","shell.execute_reply":"2025-05-06T12:51:13.014933Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, scheduler, epoch, metrics, filename):\n    \"\"\"Save model checkpoint\"\"\"\n    checkpoint = {\n        'epoch': epoch,\n        'model_state_dict': model.state_dict(),\n        'optimizer_state_dict': optimizer.state_dict(),\n        'scheduler_state_dict': scheduler.state_dict(),\n        'metrics': metrics\n    }\n    torch.save(checkpoint, filename)\n\n\n\ndef train_one_epoch(model, dataloader, loss_fn, optimizer, device, scaler=None):\n    \"\"\"Train model for one epoch\"\"\"\n    model.train()\n    epoch_loss = 0\n    all_preds, all_labels = [], []\n\n    progress_bar = tqdm(dataloader, desc=\"Training\")\n\n    for images, labels in progress_bar:\n        images = images.to(device)\n        labels = labels.to(device)\n\n        if scaler:  # Using mixed precision\n            with torch.cuda.amp.autocast():\n\n                outputs = model(images)\n                loss = loss_fn(outputs, labels)\n\n            optimizer.zero_grad()\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            outputs = model(images)\n            loss = loss_fn(outputs, labels)\n\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n        epoch_loss += loss.item()\n\n        # Store predictions and labels for metrics calculation\n        all_preds.extend(outputs.detach().cpu().numpy())\n        all_labels.extend(labels.detach().cpu().numpy())\n\n        # Update progress bar\n        progress_bar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})\n\n    # Calculate metrics\n    avg_loss = epoch_loss / len(dataloader)\n    spearman_score = spearman_corr(all_preds, all_labels)\n\n    return avg_loss, spearman_score, all_preds, all_labels\n       \n\ndef validate(model, dataloader, loss_fn, device):\n    \"\"\"Validate model on validation set\"\"\"\n    model.eval()\n    val_loss = 0\n    all_preds, all_labels = [], []\n\n    with torch.no_grad():\n        for images, labels in tqdm(dataloader, desc=\"Validation\"):\n            images = images.to(device)\n            labels = labels.to(device)\n\n            outputs = model(images)\n            loss = loss_fn(outputs, labels)\n\n            val_loss += loss.item()\n\n            # Store predictions and labels for metrics calculation\n            all_preds.extend(outputs.cpu().numpy())\n            all_labels.extend(labels.cpu().numpy())\n\n    # Calculate metrics\n    avg_loss = val_loss / len(dataloader)\n    spearman_score = spearman_corr(all_preds, all_labels)\n\n    return avg_loss, spearman_score, all_preds, all_labels\n\n\ndef get_tta_transforms():\n    \"\"\"Get test-time augmentation transformations\"\"\"\n    tta_transforms = [\n        # Original image\n        T.Compose([\n            T.ToTensor(),\n            T.Resize(CONFIG[\"image_size\"]),\n        ]),\n        # Horizontal flip\n        T.Compose([\n            T.ToTensor(),\n            T.Resize(CONFIG[\"image_size\"]),\n            T.RandomHorizontalFlip(p=1.0),\n        ]),\n        # Vertical flip\n        T.Compose([\n            T.ToTensor(),\n            T.Resize(CONFIG[\"image_size\"]),\n            T.RandomVerticalFlip(p=1.0),\n        ]),\n        # 90 degree rotation\n        T.Compose([\n            T.ToTensor(),\n            T.Resize(CONFIG[\"image_size\"]),\n            T.RandomRotation(degrees=(90, 90)),\n        ]),\n        # 180 degree rotation\n        T.Compose([\n            T.ToTensor(),\n            T.Resize(CONFIG[\"image_size\"]),\n            T.RandomRotation(degrees=(180, 180)),\n        ]),\n        # Color jitter\n        T.Compose([\n            T.ToTensor(),\n            T.Resize(CONFIG[\"image_size\"]),\n            T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),\n        ]),\n\n    ]\n    return tta_transforms\n\n\n\ndef predict_test_set_with_tta(model, data_path, device):\n    \"\"\"Generate predictions for the test set using Test Time Augmentation\"\"\"\n    tta_transforms = get_tta_transforms()\n\n    with h5py.File(f\"{data_path}/elucidata_ai_challenge_data.h5\", \"r\") as f:\n        test_spots = f[\"spots/Test\"]\n        test_images = f[\"images/Test\"]\n        sample = 'S_7'  # Test sample\n        image = np.array(test_images[sample])\n        spots = np.array(test_spots[sample])\n        x, y = spots[\"x\"], spots[\"y\"]\n\n        outputs = []\n\n        with torch.inference_mode():\n            model.eval()\n            patch_size = CONFIG[\"patch_size\"]\n            for x_, y_ in tqdm(zip(x, y), desc=\"Generating predictions with TTA\", total=len(x)):\n                half_size = patch_size // 2\n                # Ensure indices are within bounds\n                y_min = max(y_ - half_size, 0)\n                y_max = min(y_ + half_size, image.shape[0])\n                x_min = max(x_ - half_size, 0)\n                x_max = min(x_ + half_size, image.shape[1])\n                patch = image[y_min:y_max, x_min:x_max, :]\n\n                # Handle incomplete patches\n                if patch.shape[0] != patch_size or patch.shape[1] != patch_size:\n                    padded_patch = np.zeros((patch_size, patch_size, 3), dtype=patch.dtype)\n                    padded_patch[:patch.shape[0], :patch.shape[1], :] = patch\n                    patch = padded_patch\n\n                # Apply TTA and get predictions for each augmentation\n                patch_predictions = []\n                for transform in tta_transforms:\n                    patch_tensor = transform(patch)\n                    patch_tensor = patch_tensor.to(device)\n                    with torch.no_grad():\n                     output = model(patch_tensor.unsqueeze(0)).cpu().numpy()\n                     patch_predictions.append(output[0])\n\n                # Average predictions from all augmentations\n                avg_prediction = np.mean(patch_predictions, axis=0)\n                outputs.append(avg_prediction)\n\n    return np.array(outputs), x, y\n\n\n\ndef save_submission(predictions):\n    \"\"\"Save predictions to submission file\"\"\"\n    example_df = pd.read_csv(\"/kaggle/input/sample-submission/submission (1).csv\")\n    ID = example_df[\"ID\"]\n    output_df = pd.DataFrame(predictions)\n    submission_df = pd.concat([ID, output_df], axis=1)\n    submission_df.columns = example_df.columns\n\n    output_file = \"submission.csv\"\n    submission_df.to_csv(output_file, index=False)\n    print(f\"Saved submission to {output_file}\")\n\n    return output_file\n\ndef save_submission(predictions, data_path, epoch, model_name):\n    \"\"\"Save predictions to submission file\"\"\"\n    example_df = pd.read_csv(f\"/kaggle/input/sample-submission/submission (1).csv\")\n    ID = example_df[\"ID\"]\n    output_df = pd.DataFrame(predictions)\n    submission_df = pd.concat([ID, output_df], axis=1)\n    submission_df.columns = example_df.columns\n\n    output_file = \"submission.csv\"\n    submission_df.to_csv(output_file, index=False)\n    print(f\"Saved submission to {output_file}\")\n\n    return output_file","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:13.016913Z","iopub.execute_input":"2025-05-06T12:51:13.017574Z","iopub.status.idle":"2025-05-06T12:51:13.036197Z","shell.execute_reply.started":"2025-05-06T12:51:13.017548Z","shell.execute_reply":"2025-05-06T12:51:13.035582Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n\n    set_seed(CONFIG[\"seed\"])\n    device = get_device()\n    print(f\"Using device: {device}\")\n\n    # Create data transformations, datasets, and dataloaders as before...\n    train_transform, val_transform = get_transforms()\n    train_dataset = HackhathonDataset(CONFIG[\"data_path\"], transform=train_transform, mode=\"train\")\n    val_dataset = HackhathonDataset(CONFIG[\"data_path\"], transform=val_transform, mode=\"val\")\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=CONFIG[\"batch_size\"],\n        shuffle=True,\n        num_workers=CONFIG[\"num_workers\"],\n        pin_memory=True\n    )\n\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=CONFIG[\"batch_size\"],\n        shuffle=False,\n        num_workers=CONFIG[\"num_workers\"],\n        pin_memory=True\n    )\n\n    # Create model, loss function, optimizer, and scheduler as before...\n    model = create_model(CONFIG[\"model_type\"], CONFIG[\"num_classes\"])\n    model = model.to(device)\n    loss_fn = CombinedLoss()\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=CONFIG[\"learning_rate\"],\n        weight_decay=CONFIG[\"weight_decay\"]\n    )\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer,\n        mode='min',\n        factor=CONFIG[\"scheduler_gamma\"],\n        patience=3,\n    )\n\n    # Initialize mixed precision scaler if needed\n    scaler = torch.cuda.amp.GradScaler() if CONFIG[\"mixed_precision\"] and device.type == \"cuda\" else None\n\n    # Initialize wandb if enabled\n    if CONFIG[\"use_wandb\"]:\n        wandb.init(\n            project=\"hackathon-gene-expression\",\n            config=CONFIG,\n            name=f\"{CONFIG['model_type']}_run\"\n        )\n        wandb.watch(model)\n\n    # Training loop with early stopping\n    best_val_spearman = -1\n    no_improvement_count = 0\n    best_model_path = f\"{CONFIG['output_dir']}/best_{CONFIG['model_type']}_model.pt\"\n\n    for epoch in range(CONFIG[\"max_epochs\"]):\n        print(f\"\\nEpoch {epoch + 1}/{CONFIG['max_epochs']}\")\n\n        # Train one epoch\n        train_loss, train_spearman, _, _ = train_one_epoch(\n            model, train_loader, loss_fn, optimizer, device, scaler\n        )\n\n        # Validate\n        val_loss, val_spearman, _, _ = validate(\n            model, val_loader, loss_fn, device\n        )\n\n        # Update scheduler\n        scheduler.step(val_loss)\n\n        # Log metrics\n        metrics = {\n            \"train_loss\": train_loss,\n            \"train_spearman\": train_spearman,\n            \"val_loss\": val_loss,\n            \"val_spearman\": val_spearman,\n            \"learning_rate\": optimizer.param_groups[0]['lr']\n        }\n\n        print(f\"Train Loss: {train_loss:.4f}, Train Spearman: {train_spearman:.4f}\")\n        print(f\"Val Loss: {val_loss:.4f}, Val Spearman: {val_spearman:.4f}\")\n\n        if CONFIG[\"use_wandb\"]:\n            wandb.log(metrics)\n\n        # Check if this is the best model so far\n        improved = val_spearman > best_val_spearman + CONFIG[\"min_delta\"]\n\n        if improved:\n            best_val_spearman = val_spearman\n            no_improvement_count = 0\n\n            # Save the best model\n            torch.save({\n                'epoch': epoch,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'scheduler_state_dict': scheduler.state_dict(),\n                'val_spearman': val_spearman,\n            }, best_model_path)\n\n            print(f\"Saved best model with validation Spearman: {val_spearman:.4f}\")\n\n            # Generate test predictions with the best model so far\n            test_preds, x, y = predict_test_set_with_tta(model, CONFIG[\"data_path\"], device)\n\n            # Save submission file\n            submission_file = save_submission(\n                test_preds, CONFIG[\"data_path\"], f\"best_epoch_{epoch}\", CONFIG[\"model_type\"]\n            )\n\n            if CONFIG[\"use_wandb\"]:\n                wandb.save(submission_file)\n        else:\n            no_improvement_count += 1\n\n        # Generate and save predictions at checkpoint epochs regardless of improvement\n        if epoch in CONFIG[\"checkpoint_epochs\"]:\n            # Generate test predictions with Test Time Augmentation\n            test_preds, _, _ = predict_test_set_with_tta(model, CONFIG[\"data_path\"], device)\n\n            # Save submission file\n            submission_file = save_submission(\n                test_preds, CONFIG[\"data_path\"], epoch, CONFIG[\"model_type\"]\n            )\n\n            if CONFIG[\"use_wandb\"]:\n                wandb.save(submission_file)\n\n        # Check for early stopping\n        if no_improvement_count >= CONFIG[\"patience\"]:\n            print(f\"Early stopping triggered after {epoch + 1} epochs\")\n            break\n\n    # Load best model for final evaluation\n    checkpoint = torch.load(best_model_path,weights_only=False)\n    model.load_state_dict(checkpoint['model_state_dict'])\n    print(f\"Loaded best model from epoch {checkpoint['epoch']} with validation Spearman: {checkpoint['val_spearman']:.4f}\")\n\n    # Final evaluation and prediction\n    final_preds, _, _ = predict_test_set_with_tta(model, CONFIG[\"data_path\"], device)\n    final_submission = save_submission(\n        final_preds, CONFIG[\"data_path\"], \"final\", CONFIG[\"model_type\"]\n    )\n\n    if CONFIG[\"use_wandb\"]:\n        wandb.save(final_submission)\n        wandb.finish()\n\n    print(\"Training completed!\")\n\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-06T12:51:13.037197Z","iopub.execute_input":"2025-05-06T12:51:13.037422Z","execution_failed":"2025-05-06T12:51:21.082Z"}},"outputs":[],"execution_count":null}]}