File size: 3,527 Bytes
e465a2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | {
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# PixelModel v2 — Kaggle training and evaluation\n",
"\n",
"This notebook trains the 200,259-parameter v2 model on the same approximately 20K disjoint COCO caption/image pairs as v1, then evaluates it on the held-out 5K COCO rows using FID and CLIP Score.\n",
"\n",
"Before running: enable **Internet** and **GPU** in Kaggle. Add this repository as a Kaggle Dataset named `pixelmodel-v2`, so it is mounted at `/kaggle/input/pixelmodel-v2`. The notebook copies it to the writable working directory; the trained `model.png`, `model.safetensors`, and evaluation results will be in `/kaggle/working/pixelmodel-v2-output`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip -q install datasets safetensors 'torchmetrics[image]'\n",
"!pip -q install --upgrade 'torch-fidelity'\n",
"!nvidia-smi"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import shutil\n",
"from pathlib import Path\n",
"\n",
"SOURCE_DIR = Path('/kaggle/input/pixelmodel-v2')\n",
"PROJECT_DIR = Path('/kaggle/working/pixelmodel-v2')\n",
"WORK_DIR = Path('/kaggle/working/pm-work')\n",
"OUTPUT_DIR = Path('/kaggle/working/pixelmodel-v2-output')\n",
"\n",
"assert SOURCE_DIR.exists(), f'Attach the v2 repository as a Kaggle Dataset: {SOURCE_DIR}'\n",
"shutil.rmtree(PROJECT_DIR, ignore_errors=True)\n",
"shutil.copytree(SOURCE_DIR, PROJECT_DIR)\n",
"WORK_DIR.mkdir(parents=True, exist_ok=True)\n",
"OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
"os.chdir(PROJECT_DIR)\n",
"!python model.py\n",
"!cat config.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Creates the v1-compatible split: ~20K 64x64 training pairs and 5K held-out eval pairs.\n",
"# This downloads COCO through Hugging Face, so Kaggle Internet must stay enabled.\n",
"!python fetch_coco_subset.py --out {WORK_DIR}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Start fresh: do not pass --resume for the first run.\n",
"# The v2 model is larger than v1, so use a GPU and give it a longer run.\n",
"!python train.py --data {WORK_DIR / 'coco_train.npz'} --epochs 50 --batch 128 --pixels 1024 --lr 0.002\n",
"!python convert_to_safetensors.py --model model.png --out model.safetensors\n",
"# Persist the trained model immediately; Kaggle exposes /kaggle/working as notebook output.\n",
"!cp model.png model.safetensors config.json {OUTPUT_DIR}/\n",
"!ls -lh {OUTPUT_DIR}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Held-out evaluation: generate 5K images, then report FID (lower is better) and CLIP Score (higher is better).\n",
"!python eval/run_eval.py --work {WORK_DIR} --model model.png --n 5000\n",
"!cp {WORK_DIR / 'eval_results.json'} {OUTPUT_DIR}/\n",
"!ls -lh {OUTPUT_DIR}\n",
"!cat {WORK_DIR / 'eval_results.json'}"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.x"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
|