{"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":"none","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nECG Digitization - ROBUST Training Pipeline (NaN Loss Fixed)\n\nFIXES FOR NaN LOSS:\n1. Signal normalization - ECG signals can have extreme values\n2. Gradient clipping - Prevent exploding gradients\n3. Learning rate warmup - Start with small LR\n4. Loss masking validation - Check for invalid values\n5. Input validation - Check data quality\n6. Stable initialization - Proper weight initialization\n\"\"\"\n\nimport os\nimport gc\nimport warnings\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nimport torchvision.models as models\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import train_test_split\n\nwarnings.filterwarnings('ignore')\n\n# ===================== CONFIGURATION =====================\nclass RobustConfig:\n    # Paths\n    DATA_DIR = Path('/kaggle/input/physionet-ecg-image-digitization')\n    OUTPUT_DIR = Path('/kaggle/working')\n    MODEL_DIR = OUTPUT_DIR / 'models'\n    \n    # Image settings\n    IMG_SIZE = (384, 768)\n    BATCH_SIZE = 4  # Slightly larger for stability\n    GRAD_ACCUM_STEPS = 4\n    \n    # Model\n    MODEL_NAME = 'efficientnet_b0'\n    \n    # Training - CRITICAL FIXES\n    EPOCHS = 20\n    LR = 1e-4  # Lower LR to prevent instability\n    WEIGHT_DECAY = 1e-5\n    WARMUP_EPOCHS = 2  # Gradual LR increase\n    MAX_GRAD_NORM = 1.0  # Gradient clipping\n    \n    # Numerical stability\n    SIGNAL_NORMALIZATION = True  # Normalize ECG signals\n    USE_MIXED_PRECISION = True\n    LOSS_SCALE = 1.0  # Can adjust if needed\n    \n    # Data\n    NUM_WORKERS = 2\n    PIN_MEMORY = True\n    \n    # Debug\n    DEBUG_MODE = True  # Enable extensive logging\n    CHECK_NAN_FREQUENCY = 10  # Check every N batches\n    \n    # ECG settings\n    LEADS = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n    MAX_SIGNAL_LENGTH = 5000\n    \n    # Device\n    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n    \n    def __init__(self):\n        self.MODEL_DIR.mkdir(parents=True, exist_ok=True)\n\nconfig = RobustConfig()\n\nprint(f\"Robust Configuration:\")\nprint(f\"  Signal Normalization: {config.SIGNAL_NORMALIZATION}\")\nprint(f\"  Learning Rate: {config.LR}\")\nprint(f\"  Warmup Epochs: {config.WARMUP_EPOCHS}\")\nprint(f\"  Gradient Clipping: {config.MAX_GRAD_NORM}\")\nprint(f\"  Debug Mode: {config.DEBUG_MODE}\")\n\n# ===================== SIGNAL NORMALIZATION =====================\nclass SignalNormalizer:\n    \"\"\"Normalize ECG signals to prevent numerical issues\"\"\"\n    \n    @staticmethod\n    def normalize(signals):\n        \"\"\"Normalize signals to [-1, 1] range\"\"\"\n        # Typical ECG range is -5mV to +5mV\n        # Clip extreme outliers first\n        signals = np.clip(signals, -10, 10)\n        \n        # Normalize per lead\n        normalized = np.zeros_like(signals)\n        for i in range(signals.shape[0]):\n            lead = signals[i]\n            # Robust normalization using percentiles\n            p5, p95 = np.percentile(lead, [5, 95])\n            if p95 - p5 > 0.01:  # Avoid division by zero\n                normalized[i] = (lead - np.mean(lead)) / (p95 - p5 + 1e-8)\n            else:\n                normalized[i] = lead\n        \n        # Final clip\n        normalized = np.clip(normalized, -5, 5)\n        \n        return normalized.astype(np.float32)\n    \n    @staticmethod\n    def denormalize(signals, stats):\n        \"\"\"Denormalize back to original scale (for inference)\"\"\"\n        # For now, keep normalized (model trained on normalized)\n        return signals\n\n# ===================== ROBUST DATASET =====================\nclass RobustECGDataset(Dataset):\n    \"\"\"Dataset with extensive validation and normalization\"\"\"\n    \n    def __init__(self, df, image_dir, transform=None, is_training=True):\n        self.df = df.reset_index(drop=True)\n        self.image_dir = image_dir\n        self.transform = transform\n        self.is_training = is_training\n        self.normalizer = SignalNormalizer()\n        \n        # Pre-validate paths\n        if config.DEBUG_MODE:\n            self.validate_paths()\n    \n    def validate_paths(self):\n        \"\"\"Check if files exist\"\"\"\n        print(\"Validating dataset paths...\")\n        missing = 0\n        for idx in range(min(10, len(self.df))):\n            row = self.df.iloc[idx]\n            ecg_id = row['id']\n            \n            # Check image\n            img_path = self.image_dir / str(ecg_id) / f\"{ecg_id}-0001.png\"\n            if not img_path.exists():\n                missing += 1\n            \n            # Check signal\n            signal_path = self.image_dir / str(ecg_id) / f\"{ecg_id}.csv\"\n            if not signal_path.exists():\n                missing += 1\n        \n        if missing > 0:\n            print(f\"⚠ Warning: {missing} files missing in first 10 samples\")\n        else:\n            print(\"✓ All paths valid\")\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        try:\n            row = self.df.iloc[idx]\n            ecg_id = row['id']\n            \n            # Load image\n            if self.is_training:\n                segment = np.random.choice(['0001', '0003', '0004'])\n            else:\n                segment = '0001'\n            \n            img_path = self.image_dir / str(ecg_id) / f\"{ecg_id}-{segment}.png\"\n            img = cv2.imread(str(img_path))\n            \n            if img is None:\n                # Fallback to zeros\n                img = np.ones((*config.IMG_SIZE, 3), dtype=np.uint8) * 128\n            else:\n                img = cv2.resize(img, (config.IMG_SIZE[1], config.IMG_SIZE[0]))\n                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            \n            # Load signal\n            signal_path = self.image_dir / str(ecg_id) / f\"{ecg_id}.csv\"\n            signal_df = pd.read_csv(signal_path, usecols=config.LEADS)\n            \n            # Process signals\n            signals = []\n            for lead in config.LEADS:\n                lead_signal = signal_df[lead].values.astype(np.float32)\n                \n                # Handle length\n                if len(lead_signal) > config.MAX_SIGNAL_LENGTH:\n                    lead_signal = lead_signal[:config.MAX_SIGNAL_LENGTH]\n                elif len(lead_signal) < config.MAX_SIGNAL_LENGTH:\n                    lead_signal = np.pad(lead_signal, (0, config.MAX_SIGNAL_LENGTH - len(lead_signal)))\n                \n                signals.append(lead_signal)\n            \n            signals = np.stack(signals, axis=0)\n            \n            # CRITICAL: Normalize signals\n            if config.SIGNAL_NORMALIZATION:\n                signals = self.normalizer.normalize(signals)\n            \n            # Check for NaN/Inf in signals\n            if not np.isfinite(signals).all():\n                print(f\"Warning: Non-finite values in signals for {ecg_id}\")\n                signals = np.nan_to_num(signals, nan=0.0, posinf=5.0, neginf=-5.0)\n            \n            # Transform image\n            if self.transform:\n                augmented = self.transform(image=img)\n                img = augmented['image']\n            \n            # Create mask\n            actual_len = min(row['sig_len'], config.MAX_SIGNAL_LENGTH)\n            mask = np.zeros((12, config.MAX_SIGNAL_LENGTH), dtype=np.float32)\n            mask[:, :actual_len] = 1.0\n            \n            return {\n                'image': img,\n                'signals': torch.from_numpy(signals),\n                'mask': torch.from_numpy(mask)\n            }\n        \n        except Exception as e:\n            print(f\"Error loading sample {idx}: {e}\")\n            # Return dummy data\n            return {\n                'image': torch.zeros(3, config.IMG_SIZE[0], config.IMG_SIZE[1]),\n                'signals': torch.zeros(12, config.MAX_SIGNAL_LENGTH),\n                'mask': torch.ones(12, config.MAX_SIGNAL_LENGTH)\n            }\n\n# ===================== TRANSFORMS =====================\ndef get_robust_transforms(is_training):\n    if is_training:\n        return A.Compose([\n            A.GaussNoise(var_limit=(5, 15), p=0.3),\n            A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.3),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2(),\n        ])\n    else:\n        return A.Compose([\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2(),\n        ])\n\n# ===================== ROBUST MODEL =====================\nclass RobustECGModel(nn.Module):\n    \"\"\"Model with proper initialization and stability\"\"\"\n    \n    def __init__(self):\n        super().__init__()\n        \n        # Backbone\n        self.backbone = models.efficientnet_b0(pretrained=True)\n        self.feature_dim = 1280\n        self.backbone.classifier = nn.Identity()\n        \n        # Decoder with proper initialization\n        self.decoder = nn.Sequential(\n            nn.Linear(self.feature_dim, 512),\n            nn.LayerNorm(512),  # Add normalization\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(512, 256),\n            nn.LayerNorm(256),\n            nn.ReLU(),\n            nn.Dropout(0.1),\n            nn.Linear(256, 12 * config.MAX_SIGNAL_LENGTH)\n        )\n        \n        # Initialize decoder weights properly\n        self._initialize_weights()\n    \n    def _initialize_weights(self):\n        \"\"\"Proper weight initialization\"\"\"\n        for m in self.decoder.modules():\n            if isinstance(m, nn.Linear):\n                nn.init.xavier_uniform_(m.weight, gain=0.1)  # Small gain for stability\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n    \n    def forward(self, x):\n        # Extract features\n        features = self.backbone(x)\n        \n        # Check for NaN\n        if not torch.isfinite(features).all():\n            print(\"Warning: NaN in backbone features!\")\n            features = torch.nan_to_num(features, nan=0.0, posinf=1.0, neginf=-1.0)\n        \n        # Decode\n        out = self.decoder(features)\n        out = out.view(-1, 12, config.MAX_SIGNAL_LENGTH)\n        \n        # Tanh activation to bound outputs\n        out = torch.tanh(out) * 5.0  # Output range: [-5, 5]\n        \n        return out\n\n# ===================== ROBUST LOSS =====================\nclass RobustLoss(nn.Module):\n    \"\"\"Loss function with NaN protection\"\"\"\n    \n    def __init__(self):\n        super().__init__()\n    \n    def forward(self, pred, target, mask):\n        # Check inputs\n        if not torch.isfinite(pred).all():\n            print(\"⚠ NaN/Inf in predictions!\")\n            pred = torch.nan_to_num(pred, nan=0.0, posinf=5.0, neginf=-5.0)\n        \n        if not torch.isfinite(target).all():\n            print(\"⚠ NaN/Inf in targets!\")\n            target = torch.nan_to_num(target, nan=0.0, posinf=5.0, neginf=-5.0)\n        \n        # Apply mask\n        pred = pred * mask\n        target = target * mask\n        \n        # Compute loss\n        loss = F.mse_loss(pred, target, reduction='none')\n        loss = (loss * mask).sum() / (mask.sum() + 1e-8)\n        \n        # Check loss\n        if not torch.isfinite(loss):\n            print(\"⚠ NaN/Inf in loss!\")\n            loss = torch.tensor(1.0, device=loss.device)\n        \n        return loss\n\n# ===================== ROBUST TRAINER =====================\nclass RobustTrainer:\n    \"\"\"Trainer with extensive debugging and stability features\"\"\"\n    \n    def __init__(self, model, train_loader, val_loader):\n        self.model = model.to(config.DEVICE)\n        self.train_loader = train_loader\n        self.val_loader = val_loader\n        \n        # Optimizer with lower LR\n        self.optimizer = torch.optim.AdamW(\n            model.parameters(),\n            lr=config.LR,\n            weight_decay=config.WEIGHT_DECAY,\n            eps=1e-8  # Numerical stability\n        )\n        \n        # Learning rate scheduler with warmup\n        self.scheduler = self.get_scheduler()\n        \n        # Mixed precision\n        self.scaler = GradScaler(enabled=config.USE_MIXED_PRECISION)\n        \n        # Loss\n        self.criterion = RobustLoss()\n        \n        # Tracking\n        self.best_loss = float('inf')\n        self.global_step = 0\n    \n    def get_scheduler(self):\n        \"\"\"Scheduler with warmup\"\"\"\n        from torch.optim.lr_scheduler import LambdaLR\n        \n        def lr_lambda(epoch):\n            if epoch < config.WARMUP_EPOCHS:\n                # Linear warmup\n                return (epoch + 1) / config.WARMUP_EPOCHS\n            else:\n                # Cosine decay\n                progress = (epoch - config.WARMUP_EPOCHS) / (config.EPOCHS - config.WARMUP_EPOCHS)\n                return 0.5 * (1 + np.cos(np.pi * progress))\n        \n        return LambdaLR(self.optimizer, lr_lambda)\n    \n    def check_model_health(self):\n        \"\"\"Check for NaN in model parameters\"\"\"\n        for name, param in self.model.named_parameters():\n            if not torch.isfinite(param).all():\n                print(f\"⚠ NaN/Inf in parameter: {name}\")\n                return False\n        return True\n    \n    def train_epoch(self, epoch):\n        self.model.train()\n        total_loss = 0\n        num_batches = 0\n        \n        self.optimizer.zero_grad()\n        \n        pbar = tqdm(self.train_loader, desc=f'Epoch {epoch+1}/{config.EPOCHS}')\n        \n        for step, batch in enumerate(pbar):\n            try:\n                images = batch['image'].to(config.DEVICE, non_blocking=True)\n                signals = batch['signals'].to(config.DEVICE, non_blocking=True)\n                mask = batch['mask'].to(config.DEVICE, non_blocking=True)\n                \n                # Check inputs\n                if config.DEBUG_MODE and step % config.CHECK_NAN_FREQUENCY == 0:\n                    if not torch.isfinite(images).all():\n                        print(f\"⚠ NaN in images at step {step}\")\n                    if not torch.isfinite(signals).all():\n                        print(f\"⚠ NaN in signals at step {step}\")\n                \n                # Forward pass\n                with autocast(enabled=config.USE_MIXED_PRECISION):\n                    pred = self.model(images)\n                    loss = self.criterion(pred, signals, mask)\n                    loss = loss / config.GRAD_ACCUM_STEPS\n                \n                # Check loss\n                if not torch.isfinite(loss):\n                    print(f\"⚠ NaN loss at step {step}, skipping batch\")\n                    self.optimizer.zero_grad()\n                    continue\n                \n                # Backward\n                self.scaler.scale(loss).backward()\n                \n                # Update weights\n                if (step + 1) % config.GRAD_ACCUM_STEPS == 0:\n                    # Unscale gradients\n                    self.scaler.unscale_(self.optimizer)\n                    \n                    # Clip gradients\n                    grad_norm = torch.nn.utils.clip_grad_norm_(\n                        self.model.parameters(), \n                        config.MAX_GRAD_NORM\n                    )\n                    \n                    if config.DEBUG_MODE and grad_norm > 10:\n                        print(f\"⚠ Large gradient norm: {grad_norm:.2f}\")\n                    \n                    # Optimizer step\n                    self.scaler.step(self.optimizer)\n                    self.scaler.update()\n                    self.optimizer.zero_grad()\n                \n                total_loss += loss.item() * config.GRAD_ACCUM_STEPS\n                num_batches += 1\n                \n                # Update progress bar\n                avg_loss = total_loss / num_batches\n                pbar.set_postfix({\n                    'loss': f'{avg_loss:.6f}',\n                    'lr': f'{self.optimizer.param_groups[0][\"lr\"]:.2e}'\n                })\n                \n                self.global_step += 1\n                \n            except RuntimeError as e:\n                print(f\"Error at step {step}: {e}\")\n                self.optimizer.zero_grad()\n                continue\n        \n        if num_batches == 0:\n            return float('nan')\n        \n        return total_loss / num_batches\n    \n    def validate(self):\n        self.model.eval()\n        total_loss = 0\n        num_batches = 0\n        \n        with torch.no_grad():\n            for batch in tqdm(self.val_loader, desc='Validation'):\n                try:\n                    images = batch['image'].to(config.DEVICE, non_blocking=True)\n                    signals = batch['signals'].to(config.DEVICE, non_blocking=True)\n                    mask = batch['mask'].to(config.DEVICE, non_blocking=True)\n                    \n                    with autocast(enabled=config.USE_MIXED_PRECISION):\n                        pred = self.model(images)\n                        loss = self.criterion(pred, signals, mask)\n                    \n                    if torch.isfinite(loss):\n                        total_loss += loss.item()\n                        num_batches += 1\n                \n                except RuntimeError as e:\n                    print(f\"Error in validation: {e}\")\n                    continue\n        \n        if num_batches == 0:\n            return float('nan')\n        \n        return total_loss / num_batches\n    \n    def fit(self):\n        print(\"\\nStarting training...\")\n        print(f\"Total parameters: {sum(p.numel() for p in self.model.parameters()):,}\")\n        \n        for epoch in range(config.EPOCHS):\n            # Check model health\n            if not self.check_model_health():\n                print(\"⚠ Model has NaN parameters, stopping training\")\n                break\n            \n            # Train\n            train_loss = self.train_epoch(epoch)\n            \n            # Validate\n            val_loss = self.validate()\n            \n            # Print\n            print(f'Epoch {epoch+1}: Train={train_loss:.6f}, Val={val_loss:.6f}')\n            \n            # Save best\n            if np.isfinite(val_loss) and val_loss < self.best_loss:\n                self.best_loss = val_loss\n                self.save_model('best_model.pth')\n                print(f'✓ Best model saved (val={val_loss:.6f})')\n            \n            # Step scheduler\n            self.scheduler.step()\n            \n            # Memory cleanup\n            if epoch % 2 == 0:\n                gc.collect()\n                torch.cuda.empty_cache()\n        \n        return self.best_loss\n    \n    def save_model(self, filename):\n        torch.save(\n            self.model.state_dict(),\n            config.MODEL_DIR / filename\n        )\n\n# ===================== MAIN =====================\ndef main():\n    print(\"=\"*60)\n    print(\"ROBUST TRAINING PIPELINE (NaN Fix)\")\n    print(\"=\"*60)\n    \n    # Load metadata\n    print(\"\\nLoading metadata...\")\n    train_df = pd.read_csv(config.DATA_DIR / 'train.csv')\n    print(f\"Total samples: {len(train_df)}\")\n    \n    # Quick sample statistics\n    print(f\"\\nDataset statistics:\")\n    print(f\"  Sampling frequencies: {train_df['fs'].unique()}\")\n    print(f\"  Signal lengths: {train_df['sig_len'].describe()}\")\n    \n    # Split\n    train_ids, val_ids = train_test_split(\n        train_df['id'].unique(),\n        test_size=0.15,\n        random_state=42\n    )\n    \n    train_subset = train_df[train_df['id'].isin(train_ids)].reset_index(drop=True)\n    val_subset = train_df[train_df['id'].isin(val_ids)].reset_index(drop=True)\n    \n    # Use subset for initial testing\n    if config.DEBUG_MODE:\n        train_subset = train_subset.sample(min(200, len(train_subset))).reset_index(drop=True)\n        val_subset = val_subset.sample(min(50, len(val_subset))).reset_index(drop=True)\n    \n    print(f\"Train: {len(train_subset)}, Val: {len(val_subset)}\")\n    \n    # Create datasets\n    print(\"\\nCreating datasets...\")\n    train_dataset = RobustECGDataset(\n        train_subset,\n        config.DATA_DIR / 'train',\n        transform=get_robust_transforms(True),\n        is_training=True\n    )\n    \n    val_dataset = RobustECGDataset(\n        val_subset,\n        config.DATA_DIR / 'train',\n        transform=get_robust_transforms(False),\n        is_training=False\n    )\n    \n    # Create dataloaders\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=config.PIN_MEMORY,\n        drop_last=True  # Drop incomplete batches\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=config.PIN_MEMORY\n    )\n    \n    print(\"✓ Data loaders ready\")\n    \n    # Create model\n    print(\"\\nCreating model...\")\n    model = RobustECGModel()\n    \n    # Train\n    trainer = RobustTrainer(model, train_loader, val_loader)\n    best_loss = trainer.fit()\n    \n    print(f\"\\n✓ Training complete! Best val loss: {best_loss:.6f}\")\n    \n    # Cleanup\n    del model, trainer\n    gc.collect()\n    torch.cuda.empty_cache()\n\nif __name__ == '__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-24T05:31:30.184885Z","iopub.execute_input":"2025-10-24T05:31:30.185767Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}