fix(kaggle): resolve CUDA OOM, KeyError tensor(32) in ONNX export, and checkpoint export path

models/CodeFormer/basicsr/archs/codeformer_arch.py

@@ -227,7 +227,7 @@ class CodeFormer(VQAutoEncoder):
         for i, block in enumerate(self.encoder.blocks):
             x = block(x) 
             if i in out_list:
-                enc_feat_dict[str(x.shape[-1])] = x.clone()
+                enc_feat_dict[str(int(x.shape[-1]))] = x.clone()
 
         lq_feat = x
         # ################# Transformer ###################
@@ -272,7 +272,7 @@ class CodeFormer(VQAutoEncoder):
         for i, block in enumerate(self.generator.blocks):
             x = block(x) 
             if i in fuse_list: # fuse after i-th block
-                f_size = str(x.shape[-1])
+                f_size = str(int(x.shape[-1]))
                 # Always execute to avoid data-dependent control flow in ONNX
                 x = self.fuse_convs_dict[f_size](enc_feat_dict[f_size].detach(), x, w)
         out = x

models/CodeFormer/options/CodeFormer_stage3_custom.yml

@@ -56,8 +56,8 @@ datasets:
     - 5
     - 50
     latent_gt_path: null
-    num_worker_per_gpu: 4
-    batch_size_per_gpu: 4
+    num_worker_per_gpu: 2
+    batch_size_per_gpu: 2
     dataset_enlarge_ratio: 5
     prefetch_mode: cuda
     pin_memory: true

train_custom.py

@@ -130,8 +130,9 @@ def main():
                 # workers which cause MemoryError/segfaults on Windows (Task 8).
                 dataset["prefetch_mode"] = None
             else:
-                dataset["num_worker_per_gpu"] = 4
-                dataset["batch_size_per_gpu"] = 4
+                dataset["num_worker_per_gpu"] = 2
+                dataset["batch_size_per_gpu"] = 2
+                dataset["prefetch_mode"] = "cuda"
                 
     # Update weights path if they are in the project weights folder
     project_weights_path = os.path.join(project_dir, "weights", "CodeFormer", "codeformer.pth")
@@ -212,6 +213,7 @@ def main():
         else:
             env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
             env["CUDA_VISIBLE_DEVICES"] = "0"
+            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
         env["PYTHONPATH"] = os.path.pathsep.join([codeformer_dir, env.get("PYTHONPATH", "")])
 
         print("\nStarting training process. Command:")

train_kaggle.ipynb

@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🚀 Custom AI Enhancer — Kaggle Cloud GPU Training Pipeline\n",
+    "# \ud83d\ude80 Custom AI Enhancer \u2014 Kaggle Cloud GPU Training Pipeline\n",
     "- Model: CodeFormer Stage III CFT Fine-Tuning with ArcFace Identity Loss\n",
     "- Hardware: Kaggle Nvidia T4 / P100 GPU\n",
     "- Target: 2,000 Iterations + ONNX Export"
@@ -48,7 +48,7 @@
     "!git clone https://github.com/supli6669/Enhance-Image.git custom-ai-enhancer\n",
     "%cd /kaggle/working/custom-ai-enhancer\n",
     "\n",
-    "# 2. Install dependencies — pin onnxscript to a version compatible with\n",
+    "# 2. Install dependencies \u2014 pin onnxscript to a version compatible with\n",
     "#    this PyTorch (the latest onnxscript dropped ParamSchema which torch.onnx needs)\n",
     "!pip install -q facexlib lpips gdown onnx 'onnxscript==0.1.0.dev20231023' onnxruntime-gpu pyyaml opencv-python scikit-image\n",
     "!python tools/patch_and_install_basicsr.py\n",
@@ -67,7 +67,9 @@
    "source": [
     "# 4. Launch GPU Training with ArcFace Identity Loss\n",
     "%cd /kaggle/working/custom-ai-enhancer\n",
-    "!python train_custom.py"
+    "import os\n",
+    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n",
+    "!python train_custom.py\n"
    ]
   },
   {
@@ -77,54 +79,48 @@
    "outputs": [],
    "source": [
     "# 5. Export to ONNX & INT8 Quantization\n",
-    "# Wrapped in try/except — ONNX export may fail on some PyTorch/onnxscript\n",
-    "# version combos but the .pth checkpoint is always saved regardless.\n",
     "%cd /kaggle/working/custom-ai-enhancer\n",
+    "import os, sys, shutil, glob\n",
     "\n",
     "onnx_ok = False\n",
     "try:\n",
-    "    ret = os.system('python tools/export_onnx.py')\n",
-    "    onnx_ok = (ret == 0) and os.path.exists('weights/CodeFormer/codeformer.onnx')\n",
+    "    ret = os.system(\"python tools/export_onnx.py\")\n",
+    "    onnx_ok = (ret == 0) and os.path.exists(\"weights/CodeFormer/codeformer.onnx\")\n",
     "    if onnx_ok:\n",
-    "        print('[OK] ONNX export succeeded.')\n",
+    "        print(\"[OK] ONNX export succeeded.\")\n",
     "    else:\n",
-    "        print('[WARN] ONNX export exited non-zero. Skipping quantization.')\n",
+    "        print(\"[WARN] ONNX export exited non-zero. Skipping quantization.\")\n",
     "except Exception as e:\n",
-    "    print(f'[WARN] ONNX export failed: {e}. Skipping quantization.')\n",
+    "    print(f\"[WARN] ONNX export failed: {e}. Skipping quantization.\")\n",
     "\n",
     "if onnx_ok:\n",
     "    try:\n",
-    "        os.system('python tools/quantize_onnx_static.py')\n",
+    "        os.system(\"python tools/quantize_onnx_static.py\")\n",
     "    except Exception as e:\n",
-    "        print(f'[WARN] Quantization failed: {e}')\n",
+    "        print(f\"[WARN] Quantization failed: {e}\")\n",
     "\n",
-    "# Always copy whatever output files exist to /kaggle/working/ for download\n",
-    "copy_map = {\n",
-    "    'weights/CodeFormer/codeformer_int8_v2.onnx': '/kaggle/working/codeformer_int8_v3.onnx',\n",
-    "    'weights/CodeFormer/codeformer.onnx':          '/kaggle/working/codeformer_v3.onnx',\n",
-    "}\n",
-    "for src, dst in copy_map.items():\n",
-    "    if os.path.exists(src):\n",
-    "        shutil.copy(src, dst)\n",
-    "        print(f'Copied {src} -> {dst}')\n",
+    "# Copy all ONNX and weight files to /kaggle/working/ for direct download\n",
+    "for f in glob.glob(\"weights/CodeFormer/codeformer*.*\"):\n",
+    "    dst = os.path.join(\"/kaggle/working\", os.path.basename(f))\n",
+    "    shutil.copy(f, dst)\n",
+    "    print(f\"Copied {f} -> {dst}\")\n",
     "\n",
-    "# Always copy the best .pth checkpoint so it can be downloaded even without ONNX\n",
-    "import glob\n",
-    "pth_files = sorted(glob.glob('experiments/**/net_g_*.pth', recursive=True))\n",
+    "# Copy the best .pth checkpoint from models/CodeFormer/experiments\n",
+    "pth_files = sorted(glob.glob(\"models/CodeFormer/experiments/**/net_g_*.pth\", recursive=True) + glob.glob(\"experiments/**/net_g_*.pth\", recursive=True))\n",
     "if pth_files:\n",
     "    latest = pth_files[-1]\n",
-    "    shutil.copy(latest, '/kaggle/working/codeformer_finetuned_v3.pth')\n",
-    "    print(f'Checkpoint saved: {latest} -> /kaggle/working/codeformer_finetuned_v3.pth')\n",
+    "    shutil.copy(latest, \"/kaggle/working/codeformer_finetuned_v3.pth\")\n",
+    "    print(f\"Checkpoint saved: {latest} -> /kaggle/working/codeformer_finetuned_v3.pth\")\n",
     "else:\n",
-    "    print('[WARN] No .pth checkpoint found in experiments/')\n",
+    "    print(\"[WARN] No .pth checkpoint found in experiments/\")\n",
     "\n",
-    "print('\\nTraining and Export finished! Outputs saved in /kaggle/working/')\n",
-    "print('Files available for download:')\n",
-    "for f in os.listdir('/kaggle/working'):\n",
-    "    fp = f'/kaggle/working/{f}'\n",
+    "print(\"\\nTraining and Export finished! Outputs saved in /kaggle/working/\")\n",
+    "print(\"Files available for download:\")\n",
+    "for f in os.listdir(\"/kaggle/working\"):\n",
+    "    fp = f\"/kaggle/working/{f}\"\n",
     "    if os.path.isfile(fp):\n",
     "        size_mb = os.path.getsize(fp) / 1024 / 1024\n",
-    "        print(f'  {f}: {size_mb:.1f} MB')"
+    "        print(f\"  {f}: {size_mb:.1f} MB\")\n"
    ]
   }
  ],