feat: add GCN, DQN, SimCLR, YOLO (4 new models + notebooks)

README.md

@@ -49,11 +49,36 @@ Implement mainstream deep learning models from scratch.
 │   └── generate.py        # Sample generation + latent interpolation
 ├── ddpm/
 │   ├── __init__.py
-│   ├── config.yaml        # DDPM hyperparameters (T=1000, noise schedule, etc.)
+│   ├── config.yaml        # DDPM hyperparameters
 │   ├── model.py           # UNet + timestep embedding + DDPM forward/sample
 │   ├── data.py            # CIFAR-10 via HF datasets
 │   ├── train.py           # Noise prediction training
 │   └── generate.py        # Reverse diffusion sampling
+├── gcn/
+│   ├── __init__.py
+│   ├── config.yaml        # GCN hyperparameters
+│   ├── model.py           # Graph Convolution layers + 2-layer GCN
+│   ├── data.py            # Cora citation network loader
+│   ├── train.py           # Semi-supervised node classification
+│   └── eval.py            # Test accuracy evaluation
+├── dqn/
+│   ├── __init__.py
+│   ├── config.yaml        # DQN hyperparameters
+│   ├── dqn.py             # DQN, ReplayBuffer, train_episode helpers
+│   └── train.py           # CartPole RL training loop
+├── simclr/
+│   ├── __init__.py
+│   ├── config.yaml        # SimCLR hyperparameters
+│   ├── model.py           # ResNet18 encoder + Projector + NT-Xent loss
+│   ├── data.py            # CIFAR-10 with dual augmentation
+│   └── train.py           # Contrastive learning training
+├── yolo/
+│   ├── __init__.py
+│   ├── config.yaml        # YOLO hyperparameters
+│   ├── model.py           # CNN backbone + detection head
+│   ├── loss.py            # YOLO loss + NMS
+│   ├── data.py            # Pascal VOC dataset
+│   └── train.py           # Object detection training
 ├── dcgan/
 │   ├── __init__.py
 │   ├── config.yaml        # DCGAN hyperparameters
@@ -204,6 +229,42 @@ uv run python -m resnet18.train
 | Training | Noise prediction (MSE), T=1000, linear β schedule |
 | Sampling | Reverse diffusion (x_T → x_0), 1000 steps |
 
+## GCN
+
+| Item | Value |
+|---|---|
+| Model | 2-layer Graph Convolutional Network (23K params) |
+| Dataset | Cora via URL — 2708 nodes, 1433 features, 7 classes |
+| Architecture | GraphConv × 2: Â @ H @ W (spectral graph convolution) |
+| Training | Semi-supervised (20 labels/class), CrossEntropyLoss |
+
+## DQN
+
+| Item | Value |
+|---|---|
+| Model | Deep Q-Network (17K params) |
+| Environment | CartPole-v1 via Gymnasium — 4-dim state, 2 actions |
+| Architecture | 3-layer MLP (4→128→128→2) |
+| Training | Experience replay, target network, ε-greedy decay |
+
+## SimCLR
+
+| Item | Value |
+|---|---|
+| Model | SimCLR (11M params: ResNet18 encoder + MLP projector) |
+| Dataset | CIFAR-10 via HF datasets — self-supervised (no labels) |
+| Architecture | ResNet18 → Projector(512→256→128) → NT-Xent loss |
+| Training | 100 epoch, temperature=0.5, dual random augmentation |
+
+## YOLO
+
+| Item | Value |
+|---|---|
+| Model | Simplified YOLO (59M params) |
+| Dataset | Pascal VOC via HF datasets — 20 classes |
+| Architecture | CNN backbone → FC detection head → 7×7×30 output |
+| Training | YOLO loss (coord + obj + noobj + class), NMS at inference |
+
 ## DCGAN
 
 | Item | Value |
@@ -349,6 +410,10 @@ it demonstrates.
 | `nlp/bert/` | BERT mini | **Self-Attention** (semantic aggregation), **Masked Language Model** (entropy increase + denoising), LayerNorm, positional encoding |
 | `nlp/word2vec/` | Word2Vec | **Embedding lookup tables**, **Negative Sampling**, CBOW vs Skip-gram, subsampling frequent words, cosine similarity |
 | `nlp/lstm/` | LSTM | **Input/forget/output gates**, **cell state**, gradient flow through gating, sequential processing vs parallel attention |
+| `gcn/` | GCN | Graph convolution, message passing, semi-supervised node classification |
+| `dqn/` | DQN | Q-Learning, experience replay, target network, ε-greedy |
+| `simclr/` | SimCLR | Contrastive learning, NT-Xent loss, data augmentation |
+| `yolo/` | YOLO | Single-stage object detection, grid-based regression, NMS |
 | `nlp/gpt/` | GPT | **Causal Self-Attention**, **KV Cache**, autoregressive generation, word-level tokenizer, temperature + top-k sampling, bad-token blocking |
 
 ## Setup & Run
@@ -378,6 +443,19 @@ uv run python -m vae.generate
 uv run python -m nlp.seq2seq.train
 uv run python -m nlp.seq2seq.generate
 
+# Train / Evaluate GCN
+uv run python -m gcn.train
+uv run python -m gcn.eval
+
+# Train DQN
+uv run python -m dqn.train
+
+# Train SimCLR
+uv run python -m simclr.train
+
+# Train YOLO
+uv run python -m yolo.train
+
 # Train / Generate DDPM
 uv run python -m ddpm.train
 uv run python -m ddpm.generate
@@ -443,6 +521,10 @@ locally after training; paths are shown below for reference.
 | ResNet50 (40 attrs, 200K samples) | `resnet50/resnet50_celeba.pt` | ~90 MB |
 | VAE (CelebA, 64×64) | `vae/vae_celeba.pt` | 10 MB |
 | Seq2Seq Transformer (Multi30k) | `nlp/seq2seq/seq2seq_multi30k.pt` | 4 MB |
+| GCN (Cora) | `gcn/gcn_cora.pt` | 0.1 MB |
+| DQN (CartPole) | `dqn/dqn_cartpole.pt` | 0.07 MB |
+| SimCLR (CIFAR-10) | `simclr/simclr_cifar10.pt` | 22 MB |
+| YOLO (Pascal VOC) | `yolo/yolo_voc.pt` | 226 MB |
 | DDPM (CIFAR-10, 32×32) | `ddpm/ddpm_cifar10.pt` | 62 MB |
 | DCGAN (CelebA, 64×64) | `dcgan/dcgan_celeba.pt` | ~23 MB (G+D) |
 | ViT (CIFAR-10, 32×32) | `vit/vit_cifar10.pt` | 3.2 MB |

ROADMAP.md

@@ -4,25 +4,25 @@
 
 ## ✅ Completed
 
-**Models (26):**
+**Models (30):**
 - Basics (9): Logistic Regression, Linear Regression, K-Means, SVM (GD + SMO), Decision Tree, Naive Bayes, PCA, k-NN, Perceptron
 - Deep Learning (11): MLP (pure NumPy), SimpleCNN, ResNet18, ResNet34, ResNet50 (Bottleneck), BERT, Word2Vec (CBOW + Skip-gram), LSTM (hand-written gates), GPT (Causal Attention + KV Cache), Seq2Seq Transformer (cross-attention), DDPM (Denoising Diffusion)
 - Image Generation (1): DCGAN — CelebA 64×64
 - Variational AE (1): VAE — Encoder, reparameterization trick, KL divergence, CelebA 64×64
 - Vision Transformer (1): ViT — CIFAR-10
 - Segmentation (1): UNet — Oxford-IIIT Pet
+- Graph (1): GCN — Cora citation network, spectral graph convolution
+- Reinforcement Learning (1): DQN — CartPole-v1, experience replay
+- Self-Supervised (1): SimCLR — CIFAR-10, NT-Xent contrastive loss
+- Object Detection (1): YOLO — Pascal VOC, grid-based detection
 
 **Infrastructure:**
-- Config system (YAML)
-- TensorBoard logging
-- Reproducibility (seed + config saving)
+- Config system (YAML), TensorBoard logging, Reproducibility (seed + config saving)
 - Device auto-detection (`utils.device.get_device`: CUDA → MPS → CPU)
 
-**Jupyter Notebooks (15):**
-- P0: GPT, ViT, DCGAN, UNet
-- P1: CNN, BERT, ResNet18, ResNet34
-- P2: basics (9)
-- P3: Word2Vec, LSTM
+**Jupyter Notebooks (19):**
+- Original 15: GPT, ViT, DCGAN, UNet, CNN, BERT, ResNet18, ResNet34, basics (9), Word2Vec, LSTM
+- New 4: GCN, DQN, SimCLR, YOLO
 
 ---
 

dqn/__init__.py

@@ -0,0 +1,3 @@
+from .dqn import DQN, ReplayBuffer, train_episode
+
+__all__ = ["DQN", "ReplayBuffer", "train_episode"]

dqn/config.yaml

@@ -0,0 +1,12 @@
+seed: 42
+lr: 0.001
+gamma: 0.99
+epsilon_start: 1.0
+epsilon_end: 0.01
+epsilon_decay: 500
+batch_size: 64
+buffer_size: 50000
+target_update: 100
+num_episodes: 500
+hidden_dim: 128
+model_path: dqn/dqn_cartpole.pt

dqn/dqn.ipynb

@@ -0,0 +1,212 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d18ad3b8",
+   "metadata": {},
+   "source": [
+    "# DQN: Deep Q-Network\n",
+    "\n",
+    "Reinforcement learning with experience replay on CartPole-v1."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6088c114",
+   "metadata": {},
+   "source": [
+    "## 背景\n",
+    "\n",
+    "DQN（Mnih et al. 2015）将深度学习与 Q-Learning 结合，首次在 Atari 游戏上达到人类水平。\n",
+    "核心创新：\n",
+    "\n",
+    "- **经验回放（Experience Replay）**：存储过去经验，随机采样训练，打破数据相关性\n",
+    "- **目标网络（Target Network）**：固定 Q-target，稳定训练\n",
+    "\n",
+    "环境：**CartPole-v1** — 控制小车左右移动，保持杆子不倒。\n",
+    "状态：4 维（位置、速度、角度、角速度），动作：2 维（左、右）。\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c5a83de",
+   "metadata": {},
+   "source": [
+    "## 数学原理\n",
+    "\n",
+    "### Q-Learning\n",
+    "\n",
+    "$$Q(s, a) \\leftarrow Q(s, a) + \\alpha \\left(r + \\gamma \\max_{a'} Q(s', a') - Q(s, a)\\right)$$\n",
+    "\n",
+    "### DQN Loss\n",
+    "\n",
+    "$$\\mathcal{L} = \\mathbb{E}_{(s,a,r,s') \\sim \\mathcal{D}} \\left[\\left(r + \\gamma \\max_{a'} Q_{\\theta^-}(s', a') - Q_\\theta(s, a)\\right)^2\\right]$$\n",
+    "\n",
+    "- $\\mathcal{D}$: 经验回放缓冲区\n",
+    "- $\\theta$: 在线网络参数\n",
+    "- $\\theta^-$: 目标网络参数（每隔 $C$ 步复制一次）\n",
+    "\n",
+    "### ε-greedy 探索\n",
+    "\n",
+    "$$a = \\begin{cases} \\text{random}, & \\text{概率 } \\varepsilon \\\\ \\arg\\max_a Q(s, a), & \\text{概率 } 1-\\varepsilon \\end{cases}$$\n",
+    "\n",
+    "ε 随时间指数衰减。\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a82a0cbd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gymnasium as gym\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "\n",
+    "from dqn.dqn import DQN, ReplayBuffer, train_episode, epsilon_by_episode\n",
+    "from utils.config import load_config\n",
+    "\n",
+    "device = torch.device(\"mps\" if torch.backends.mps.is_available() else \"cpu\")\n",
+    "print(f\"Device: {device}\")\n",
+    "\n",
+    "env = gym.make(\"CartPole-v1\")\n",
+    "state_dim = env.observation_space.shape[0]\n",
+    "action_dim = env.action_space.n\n",
+    "print(f\"State: {state_dim}  Action: {action_dim}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "707631bd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = DQN(state_dim, action_dim, hidden_dim=128).to(device)\n",
+    "target = DQN(state_dim, action_dim, hidden_dim=128).to(device)\n",
+    "target.load_state_dict(model.state_dict())\n",
+    "target.eval()\n",
+    "print(f\"Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "872fd75b",
+   "metadata": {},
+   "source": [
+    "## 训练\n",
+    "\n",
+    "> ⏱ 预估耗时：**500 episode × ~0.2s ≈ 2 分钟**（M4 Max）\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2efa79de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_EPISODES = 500\n",
+    "LR = 0.001\n",
+    "GAMMA = 0.99\n",
+    "BATCH_SIZE = 64\n",
+    "BUFFER_SIZE = 50000\n",
+    "TARGET_UPDATE = 100\n",
+    "EPSILON_START = 1.0\n",
+    "EPSILON_END = 0.01\n",
+    "EPSILON_DECAY = 500\n",
+    "\n",
+    "optimizer = optim.Adam(model.parameters(), lr=LR)\n",
+    "buffer = ReplayBuffer(BUFFER_SIZE)\n",
+    "rewards = []\n",
+    "\n",
+    "for episode in range(1, NUM_EPISODES + 1):\n",
+    "    state, _ = env.reset()\n",
+    "    episode_reward = 0\n",
+    "    eps = epsilon_by_episode(episode, EPSILON_START, EPSILON_END, EPSILON_DECAY)\n",
+    "\n",
+    "    while True:\n",
+    "        if np.random.random() < eps:\n",
+    "            action = env.action_space.sample()\n",
+    "        else:\n",
+    "            with torch.no_grad():\n",
+    "                q = model(torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0))\n",
+    "                action = q.argmax().item()\n",
+    "\n",
+    "        next_state, reward, terminated, truncated, _ = env.step(action)\n",
+    "        done = terminated or truncated\n",
+    "        buffer.push(state, action, reward, next_state, done)\n",
+    "        state = next_state\n",
+    "        episode_reward += reward\n",
+    "        _ = train_episode(model, target, optimizer, buffer, BATCH_SIZE, GAMMA)\n",
+    "        if done:\n",
+    "            break\n",
+    "\n",
+    "    if episode % TARGET_UPDATE == 0:\n",
+    "        target.load_state_dict(model.state_dict())\n",
+    "\n",
+    "    rewards.append(episode_reward)\n",
+    "    if episode % 50 == 0:\n",
+    "        avg = np.mean(rewards[-50:])\n",
+    "        print(f\"Episode [{episode:3d}/{NUM_EPISODES}]  Reward: {episode_reward:.0f}  Avg(50): {avg:.1f}  ε: {eps:.3f}\")\n",
+    "\n",
+    "env.close()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cff4e844",
+   "metadata": {},
+   "source": [
+    "## Reward 曲线"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "562279a0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.figure(figsize=(8, 4))\n",
+    "plt.plot(rewards)\n",
+    "plt.xlabel(\"Episode\"); plt.ylabel(\"Total Reward\"); plt.title(\"DQN Training on CartPole\")\n",
+    "plt.grid(True)\n",
+    "plt.axhline(y=195, color='r', linestyle='--', label='Solved (195)')\n",
+    "plt.legend()\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "84349669",
+   "metadata": {},
+   "source": [
+    "## 思考题\n",
+    "\n",
+    "1. 为什么需要经验回放（Experience Replay）？在线学习会有什么问题？\n",
+    "2. 目标网络（Target Network）解决了什么？如果不固定 target，Q 值会发散吗？\n",
+    "3. ϵ-greedy 中的 ϵ 从 1.0 开始衰减有什么含义？\n",
+    "4. 如果把 `hidden_dim` 从 128 改到 32，训练速度会怎样？收敛难度呢？\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

dqn/dqn.py

@@ -0,0 +1,78 @@
+"""DQN: Deep Q-Network for CartPole-v1."""
+
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DQN(nn.Module):
+    """Q-network: state → action values."""
+
+    def __init__(self, state_dim, action_dim, hidden_dim=128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, action_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class ReplayBuffer:
+    """Experience replay buffer."""
+
+    def __init__(self, capacity=50000):
+        self.capacity = capacity
+        self.buffer = []
+        self.pos = 0
+
+    def push(self, state, action, reward, next_state, done):
+        if len(self.buffer) < self.capacity:
+            self.buffer.append(None)
+        self.buffer[self.pos] = (state, action, reward, next_state, done)
+        self.pos = (self.pos + 1) % self.capacity
+
+    def sample(self, batch_size):
+        batch = random.sample(self.buffer, batch_size)
+        states, actions, rewards, next_states, dones = zip(*batch)
+        return (torch.tensor(np.array(states), dtype=torch.float32),
+                torch.tensor(actions, dtype=torch.long),
+                torch.tensor(rewards, dtype=torch.float32),
+                torch.tensor(np.array(next_states), dtype=torch.float32),
+                torch.tensor(dones, dtype=torch.float32))
+
+    def __len__(self):
+        return len(self.buffer)
+
+
+def train_episode(model, target_model, optimizer, replay_buffer, batch_size, gamma):
+    """Train DQN on one batch from replay buffer."""
+    if len(replay_buffer) < batch_size:
+        return 0.0
+
+    states, actions, rewards, next_states, dones = replay_buffer.sample(batch_size)
+
+    # Q(s, a)
+    q_values = model(states).gather(1, actions.unsqueeze(1)).squeeze(1)
+
+    # Target: r + γ * max Q'(s', a')
+    with torch.no_grad():
+        max_next_q = target_model(next_states).max(dim=1)[0]
+        target = rewards + gamma * max_next_q * (1 - dones)
+
+    loss = F.mse_loss(q_values, target)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    return loss.item()
+
+
+def epsilon_by_episode(episode, epsilon_start, epsilon_end, epsilon_decay):
+    """Annealed ε-greedy."""
+    return epsilon_end + (epsilon_start - epsilon_end) * np.exp(-1.0 * episode / epsilon_decay)

dqn/train.py

@@ -0,0 +1,93 @@
+"""DQN training on CartPole-v1."""
+
+import numpy as np
+import torch
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+
+from dqn.dqn import DQN, ReplayBuffer, train_episode, epsilon_by_episode
+from utils.config import load_config, save_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+
+def train():
+    cfg = load_config("dqn/config.yaml")
+    set_seed(cfg["seed"])
+
+    try:
+        import gymnasium as gym
+    except ImportError:
+        print("Installing gymnasium...")
+        import subprocess, sys
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "gymnasium"])
+        import gymnasium as gym
+
+    device = get_device()
+    print(f"Device: {device}")
+
+    env = gym.make("CartPole-v1")
+    state_dim = env.observation_space.shape[0]
+    action_dim = env.action_space.n
+    print(f"CartPole: state={state_dim}, action={action_dim}")
+
+    model = DQN(state_dim, action_dim, cfg["hidden_dim"]).to(device)
+    target_model = DQN(state_dim, action_dim, cfg["hidden_dim"]).to(device)
+    target_model.load_state_dict(model.state_dict())
+    target_model.eval()
+
+    optimizer = optim.Adam(model.parameters(), lr=cfg["lr"])
+    replay_buffer = ReplayBuffer(cfg["buffer_size"])
+
+    num_episodes = cfg["num_episodes"]
+    writer = SummaryWriter(log_dir="runs/dqn")
+
+    episode_rewards = []
+
+    for episode in range(1, num_episodes + 1):
+        state, _ = env.reset()
+        episode_reward = 0
+        epsilon = epsilon_by_episode(episode, cfg["epsilon_start"], cfg["epsilon_end"], cfg["epsilon_decay"])
+
+        while True:
+            # ε-greedy action selection.
+            if np.random.random() < epsilon:
+                action = env.action_space.sample()
+            else:
+                with torch.no_grad():
+                    q = model(torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0))
+                    action = q.argmax().item()
+
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+            replay_buffer.push(state, action, reward, next_state, done)
+
+            state = next_state
+            episode_reward += reward
+
+            loss = train_episode(model, target_model, optimizer, replay_buffer, cfg["batch_size"], cfg["gamma"])
+
+            if done:
+                break
+
+        # Update target network.
+        if episode % cfg["target_update"] == 0:
+            target_model.load_state_dict(model.state_dict())
+
+        episode_rewards.append(episode_reward)
+        writer.add_scalar("train/reward", episode_reward, episode)
+        writer.add_scalar("train/epsilon", epsilon, episode)
+
+        if episode % 50 == 0:
+            avg_reward = np.mean(episode_rewards[-50:])
+            print(f"Episode [{episode:4d}/{num_episodes}]  Reward: {episode_reward:.0f}  Avg(50): {avg_reward:.1f}  ε: {epsilon:.3f}")
+
+    writer.close()
+    env.close()
+    torch.save(model.state_dict(), cfg["model_path"])
+    save_config(cfg, cfg["model_path"].replace(".pt", "_config.yaml"))
+    print(f"\nModel saved to {cfg['model_path']}")
+
+
+if __name__ == "__main__":
+    train()

gcn/__init__.py

@@ -0,0 +1,3 @@
+from .model import GCN
+
+__all__ = ["GCN"]

gcn/config.yaml

@@ -0,0 +1,7 @@
+seed: 42
+lr: 0.01
+weight_decay: 5e-4
+hidden_dim: 16
+num_epochs: 200
+dropout: 0.5
+model_path: gcn/gcn_cora.pt

gcn/data.py

@@ -0,0 +1,71 @@
+"""Cora citation network dataset loader."""
+
+import os
+import urllib.request
+import numpy as np
+import torch
+
+
+CORA_URL = "https://raw.githubusercontent.com/tkipf/pygcn/master/data/cora/"
+
+
+def _download(raw_dir="gcn/raw"):
+    os.makedirs(raw_dir, exist_ok=True)
+    for fname in ["cora.content", "cora.cites"]:
+        path = os.path.join(raw_dir, fname)
+        if not os.path.exists(path):
+            urllib.request.urlretrieve(CORA_URL + fname, path)
+
+
+def _encode_onehot(labels):
+    classes = sorted(set(labels))
+    mapping = {c: i for i, c in enumerate(classes)}
+    return torch.eye(len(classes))[torch.tensor([mapping[l] for l in labels])], mapping
+
+
+def load_cora(raw_dir="gcn/raw"):
+    _download(raw_dir)
+
+    # Parse content: node_id features label
+    content = np.genfromtxt(os.path.join(raw_dir, "cora.content"), dtype=np.dtype(str))
+    node_ids = content[:, 0].astype(int)
+    features = torch.tensor(content[:, 1:-1].astype(np.float32))
+    labels_onehot, class_mapping = _encode_onehot(content[:, -1])
+    labels = labels_onehot.argmax(dim=1)
+
+    n = len(node_ids)
+    id_to_idx = {int(nid): i for i, nid in enumerate(node_ids)}
+
+    # Parse cites: source target
+    cites = np.genfromtxt(os.path.join(raw_dir, "cora.cites"), dtype=np.int32)
+    adj = torch.zeros((n, n), dtype=torch.float32)
+    for src, tgt in cites:
+        if src in id_to_idx and tgt in id_to_idx:
+            adj[id_to_idx[src], id_to_idx[tgt]] = 1.0
+            adj[id_to_idx[tgt], id_to_idx[src]] = 1.0  # undirected
+
+    # Normalize: D^{-1/2} @ A @ D^{-1/2}
+    rowsum = adj.sum(dim=1).clamp(min=1e-8)
+    d_inv_sqrt = torch.diag(rowsum ** -0.5)
+    adj_norm = d_inv_sqrt @ adj @ d_inv_sqrt
+
+    # Standard splits (20 nodes/class for train, 500 val, 1000 test).
+    n_classes = len(class_mapping)
+    idx_per_class = [torch.where(labels == c)[0] for c in range(n_classes)]
+
+    train_idx = torch.cat([idx[:20] for idx in idx_per_class])
+    rest = torch.cat([idx[20:] for idx in idx_per_class])
+    val_idx = rest[:500]
+    test_idx = rest[500:1500]
+
+    train_mask = torch.zeros(n, dtype=torch.bool)
+    val_mask = torch.zeros(n, dtype=torch.bool)
+    test_mask = torch.zeros(n, dtype=torch.bool)
+    train_mask[train_idx] = True
+    val_mask[val_idx] = True
+    test_mask[test_idx] = True
+
+    print(f"Cora: {n} nodes, {features.size(1)} features, {n_classes} classes")
+    print(f"  Train: {train_idx.size(0)}  Val: {val_idx.size(0)}  Test: {test_idx.size(0)}")
+
+    return features, adj_norm, labels, train_mask, val_mask, test_mask, class_mapping

gcn/eval.py

@@ -0,0 +1,48 @@
+import torch
+
+from gcn.model import GCN
+from gcn.data import load_cora
+from utils.config import load_config
+
+
+def evaluate():
+    cfg = load_config("gcn/config.yaml")
+
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+
+    features, adj_norm, labels, _, _, test_mask, class_mapping = load_cora()
+    features, adj_norm, labels = features.to(device), adj_norm.to(device), labels.to(device)
+    test_mask = test_mask.to(device)
+
+    model = GCN(
+        in_features=features.size(1),
+        hidden_dim=cfg["hidden_dim"],
+        num_classes=labels.max().item() + 1,
+        dropout=cfg["dropout"],
+    )
+    model.load_state_dict(torch.load(cfg["model_path"], map_location=device, weights_only=True))
+    model = model.to(device)
+    model.eval()
+    print(f"Loaded model from {cfg['model_path']}")
+
+    with torch.no_grad():
+        output = model(features, adj_norm)
+        pred = output[test_mask].argmax(dim=1)
+        correct = (pred == labels[test_mask]).sum().item()
+        total = test_mask.sum().item()
+        acc = correct / total
+
+    print(f"Test Accuracy: {acc:.2%} ({correct}/{total})")
+
+    id_to_name = {v: k for k, v in class_mapping.items()}
+    print(f"\nPer-class accuracy:")
+    pred_all = output.argmax(dim=1)
+    for c, name in sorted(id_to_name.items()):
+        mask = test_mask & (labels == c)
+        if mask.sum() > 0:
+            acc_c = (pred_all[mask] == labels[mask]).float().mean().item()
+            print(f"  {name:<15} {acc_c:.2%}")
+
+
+if __name__ == "__main__":
+    evaluate()

gcn/gcn.ipynb

@@ -0,0 +1,217 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f71e99ba",
+   "metadata": {},
+   "source": [
+    "# GCN: Graph Convolutional Network\n",
+    "\n",
+    "Node classification on citation graphs using spectral graph convolution.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "47197e28",
+   "metadata": {},
+   "source": [
+    "## 背景\n",
+    "\n",
+    "GCN（Kipf & Welling, 2017）将卷积操作推广到图结构数据。核心思想：\n",
+    "每个节点的特征由其邻居节点加权聚合而来。\n",
+    "\n",
+    "与图像 CNN 的区别：\n",
+    "- CNN：固定网格结构，卷积核在空间上滑动\n",
+    "- GCN：任意图结构，卷积由邻接矩阵定义的消息传递实现\n",
+    "\n",
+    "数据集：**Cora** — 2708 篇论文，每篇用 1433 维词袋向量表示，分为 7 类。边表示引用关系。\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f75792e0",
+   "metadata": {},
+   "source": [
+    "## 数学原理\n",
+    "\n",
+    "### 图卷积层\n",
+    "\n",
+    "$$H^{(l+1)} = \\sigma\\left(\\hat{A} H^{(l)} W^{(l)}\\right)$$\n",
+    "\n",
+    "其中 $\\hat{A} = D^{-1/2} A D^{-1/2}$ 是归一化邻接矩阵。\n",
+    "\n",
+    "- $A$: 邻接矩阵（加自环后）\n",
+    "- $D$: 度矩阵 $D_{ii} = \\sum_j A_{ij}$\n",
+    "- $H^{(l)}$: 第 $l$ 层的节点表示\n",
+    "- $W^{(l)}$: 可学习的权重矩阵\n",
+    "\n",
+    "### 2 层 GCN\n",
+    "\n",
+    "$$Z = \\text{softmax}\\left(\\hat{A}\\ \\text{ReLU}\\left(\\hat{A} X W^{(0)}\\right) W^{(1)}\\right)$$\n",
+    "\n",
+    "半监督学习：只用少量标注节点（每类 20 个）训练，模型通过图结构传播标签信息到未标注节点。\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "32984352",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "\n",
+    "from gcn.model import GCN\n",
+    "from gcn.data import load_cora\n",
+    "from utils.config import load_config\n",
+    "from utils.seed import set_seed\n",
+    "from utils.device import get_device\n",
+    "\n",
+    "device = get_device()\n",
+    "print(f\"Device: {device}\")\n",
+    "\n",
+    "features, adj_norm, labels, train_mask, val_mask, test_mask, classes = load_cora()\n",
+    "features = features.to(device)\n",
+    "adj_norm = adj_norm.to(device)\n",
+    "labels = labels.to(device)\n",
+    "train_mask = train_mask.to(device)\n",
+    "val_mask = val_mask.to(device)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "419e0b87",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = GCN(\n",
+    "    in_features=features.size(1),\n",
+    "    hidden_dim=16,\n",
+    "    num_classes=labels.max().item() + 1,\n",
+    "    dropout=0.5,\n",
+    ").to(device)\n",
+    "print(f\"Parameters: {model.num_params():,}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ce56f902",
+   "metadata": {},
+   "source": [
+    "## 训练\n",
+    "\n",
+    "> ⏱ 预估耗时：**200 epoch × ~0.1s/epoch ≈ 20 秒**（CPU 即可完成）\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e9179b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_EPOCHS = 200\n",
+    "LR = 0.01\n",
+    "WEIGHT_DECAY = 5e-4\n",
+    "\n",
+    "criterion = nn.CrossEntropyLoss()\n",
+    "optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n",
+    "\n",
+    "loss_hist, acc_hist = [], []\n",
+    "\n",
+    "for epoch in range(1, NUM_EPOCHS + 1):\n",
+    "    model.train()\n",
+    "    optimizer.zero_grad()\n",
+    "    output = model(features, adj_norm)\n",
+    "    loss = criterion(output[train_mask], labels[train_mask])\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    model.eval()\n",
+    "    with torch.no_grad():\n",
+    "        output = model(features, adj_norm)\n",
+    "        val_acc = (output[val_mask].argmax(dim=1) == labels[val_mask]).float().mean().item()\n",
+    "        loss_hist.append(loss.item())\n",
+    "        acc_hist.append(val_acc)\n",
+    "\n",
+    "    if epoch % 20 == 0 or epoch == 1:\n",
+    "        print(f\"Epoch [{epoch:3d}/{NUM_EPOCHS}]  Loss: {loss.item():.4f}  Val Acc: {val_acc:.2%}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "769f23ab",
+   "metadata": {},
+   "source": [
+    "## Loss 曲线"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad7e46b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))\n",
+    "ax1.plot(loss_hist); ax1.set_xlabel(\"Epoch\"); ax1.set_ylabel(\"Loss\"); ax1.set_title(\"Training Loss\"); ax1.grid(True)\n",
+    "ax2.plot(acc_hist, color='green'); ax2.set_xlabel(\"Epoch\"); ax2.set_ylabel(\"Val Acc\"); ax2.set_title(\"Validation Accuracy\"); ax2.grid(True)\n",
+    "plt.tight_layout(); plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bf205210",
+   "metadata": {},
+   "source": [
+    "## 测试准确率"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59cc2eac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    output = model(features, adj_norm)\n",
+    "    pred = output[test_mask].argmax(dim=1)\n",
+    "    test_acc = (pred == labels[test_mask]).float().mean().item()\n",
+    "print(f\"Test Accuracy: {test_acc:.2%}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a77fac89",
+   "metadata": {},
+   "source": [
+    "## 思考题\n",
+    "\n",
+    "1. 为什么 GCN 的归一化用 $D^{-1/2} A D^{-1/2}$ 而不是 $D^{-1} A$？\n",
+    "2. GCN 能处理归纳式（inductive）任务吗？还是只能直推式（transductive）？\n",
+    "3. 如果不用邻接矩阵只用节点特征，准确率会降到多少？\n",
+    "4. GCN 层数加深为什么会导致性能下降？（提示：过平滑问题）\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

gcn/model.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class GraphConv(nn.Module):
+    """Graph Convolution: H' = σ(Â @ H @ W)"""
+
+    def __init__(self, in_features, out_features):
+        super().__init__()
+        self.W = nn.Parameter(torch.randn(in_features, out_features) * 0.01)
+
+    def forward(self, x, adj_norm):
+        # x: (N, in_features), adj_norm: (N, N)
+        return adj_norm @ x @ self.W
+
+
+class GCN(nn.Module):
+    """2-layer Graph Convolutional Network (Kipf & Welling, 2017)."""
+
+    def __init__(self, in_features, hidden_dim, num_classes, dropout=0.5):
+        super().__init__()
+        self.conv1 = GraphConv(in_features, hidden_dim)
+        self.conv2 = GraphConv(hidden_dim, num_classes)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, adj_norm):
+        x = self.conv1(x, adj_norm)
+        x = F.relu(x)
+        x = self.dropout(x)
+        x = self.conv2(x, adj_norm)
+        return x
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

gcn/train.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+
+from gcn.model import GCN
+from gcn.data import load_cora
+from utils.config import load_config, save_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+
+def train():
+    cfg = load_config("gcn/config.yaml")
+    set_seed(cfg["seed"])
+
+    device = get_device()
+    print(f"Device: {device}")
+
+    features, adj_norm, labels, train_mask, val_mask, test_mask, _ = load_cora()
+    features, adj_norm, labels = features.to(device), adj_norm.to(device), labels.to(device)
+    train_mask, val_mask = train_mask.to(device), val_mask.to(device)
+
+    model = GCN(
+        in_features=features.size(1),
+        hidden_dim=cfg["hidden_dim"],
+        num_classes=labels.max().item() + 1,
+        dropout=cfg["dropout"],
+    ).to(device)
+    print(f"Parameters: {model.num_params():,}")
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
+
+    num_epochs = cfg["num_epochs"]
+    writer = SummaryWriter(log_dir="runs/gcn")
+
+    for epoch in range(1, num_epochs + 1):
+        model.train()
+        optimizer.zero_grad()
+        output = model(features, adj_norm)
+        loss = criterion(output[train_mask], labels[train_mask])
+        loss.backward()
+        optimizer.step()
+
+        with torch.no_grad():
+            model.eval()
+            output = model(features, adj_norm)
+            val_loss = criterion(output[val_mask], labels[val_mask])
+            val_acc = (output[val_mask].argmax(dim=1) == labels[val_mask]).float().mean().item()
+
+        writer.add_scalar("train/loss", loss.item(), epoch)
+        writer.add_scalar("val/loss", val_loss.item(), epoch)
+        writer.add_scalar("val/acc", val_acc, epoch)
+
+        if epoch % 10 == 0 or epoch == 1:
+            print(f"Epoch [{epoch:3d}/{num_epochs}]  Loss: {loss.item():.4f}  Val Acc: {val_acc:.2%}")
+
+    writer.close()
+    save_path = cfg["model_path"]
+    torch.save(model.state_dict(), save_path)
+    save_config(cfg, save_path.replace(".pt", "_config.yaml"))
+    print(f"\nModel saved to {save_path}")
+
+
+if __name__ == "__main__":
+    train()

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "scikit-learn>=1.0",
     "nbformat>=5.0",
     "matplotlib>=3.0",
+    "gymnasium>=0.29",
     "jupyterlab>=4.0",
     "ipykernel>=6.0",
 ]

scripts/gen_dqn_notebook.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+"""Generate DQN notebook."""
+
+import nbformat as nbf
+
+nb = nbf.v4.new_notebook()
+nb.metadata = {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},"language_info": {"name": "python", "version": "3.12.0"}}
+
+cells = []
+def md(s): cells.append(nbf.v4.new_markdown_cell(s))
+def code(s): cells.append(nbf.v4.new_code_cell(s))
+
+md("# DQN: Deep Q-Network\n\nReinforcement learning with experience replay on CartPole-v1.")
+
+md("""## 背景
+
+DQN（Mnih et al. 2015）将深度学习与 Q-Learning 结合，首次在 Atari 游戏上达到人类水平。
+核心创新：
+
+- **经验回放（Experience Replay）**：存储过去经验，随机采样训练，打破数据相关性
+- **目标网络（Target Network）**：固定 Q-target，稳定训练
+
+环境：**CartPole-v1** — 控制小车左右移动，保持杆子不倒。
+状态：4 维（位置、速度、角度、角速度），动作：2 维（左、右）。
+""")
+
+md("""## 数学原理
+
+### Q-Learning
+
+$$Q(s, a) \\leftarrow Q(s, a) + \\alpha \\left(r + \\gamma \\max_{a'} Q(s', a') - Q(s, a)\\right)$$
+
+### DQN Loss
+
+$$\\mathcal{L} = \\mathbb{E}_{(s,a,r,s') \\sim \\mathcal{D}} \\left[\\left(r + \\gamma \\max_{a'} Q_{\\theta^-}(s', a') - Q_\\theta(s, a)\\right)^2\\right]$$
+
+- $\\mathcal{D}$: 经验回放缓冲区
+- $\\theta$: 在线网络参数
+- $\\theta^-$: 目标网络参数（每隔 $C$ 步复制一次）
+
+### ε-greedy 探索
+
+$$a = \\begin{cases} \\text{random}, & \\text{概率 } \\varepsilon \\\\ \\arg\\max_a Q(s, a), & \\text{概率 } 1-\\varepsilon \\end{cases}$$
+
+ε 随时间指数衰减。
+""")
+
+code("""\
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from dqn.dqn import DQN, ReplayBuffer, train_episode, epsilon_by_episode
+from utils.config import load_config
+
+device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+print(f"Device: {device}")
+
+env = gym.make("CartPole-v1")
+state_dim = env.observation_space.shape[0]
+action_dim = env.action_space.n
+print(f"State: {state_dim}  Action: {action_dim}")
+""")
+
+code("""\
+model = DQN(state_dim, action_dim, hidden_dim=128).to(device)
+target = DQN(state_dim, action_dim, hidden_dim=128).to(device)
+target.load_state_dict(model.state_dict())
+target.eval()
+print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
+""")
+
+md("""## 训练
+
+> ⏱ 预估耗时：**500 episode × ~0.2s ≈ 2 分钟**（M4 Max）
+""")
+
+code("""\
+NUM_EPISODES = 500
+LR = 0.001
+GAMMA = 0.99
+BATCH_SIZE = 64
+BUFFER_SIZE = 50000
+TARGET_UPDATE = 100
+EPSILON_START = 1.0
+EPSILON_END = 0.01
+EPSILON_DECAY = 500
+
+optimizer = optim.Adam(model.parameters(), lr=LR)
+buffer = ReplayBuffer(BUFFER_SIZE)
+rewards = []
+
+for episode in range(1, NUM_EPISODES + 1):
+    state, _ = env.reset()
+    episode_reward = 0
+    eps = epsilon_by_episode(episode, EPSILON_START, EPSILON_END, EPSILON_DECAY)
+
+    while True:
+        if np.random.random() < eps:
+            action = env.action_space.sample()
+        else:
+            with torch.no_grad():
+                q = model(torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0))
+                action = q.argmax().item()
+
+        next_state, reward, terminated, truncated, _ = env.step(action)
+        done = terminated or truncated
+        buffer.push(state, action, reward, next_state, done)
+        state = next_state
+        episode_reward += reward
+        _ = train_episode(model, target, optimizer, buffer, BATCH_SIZE, GAMMA)
+        if done:
+            break
+
+    if episode % TARGET_UPDATE == 0:
+        target.load_state_dict(model.state_dict())
+
+    rewards.append(episode_reward)
+    if episode % 50 == 0:
+        avg = np.mean(rewards[-50:])
+        print(f"Episode [{episode:3d}/{NUM_EPISODES}]  Reward: {episode_reward:.0f}  Avg(50): {avg:.1f}  ε: {eps:.3f}")
+
+env.close()
+""")
+
+md("""## Reward 曲线""")
+
+code("""\
+import matplotlib.pyplot as plt
+
+plt.figure(figsize=(8, 4))
+plt.plot(rewards)
+plt.xlabel("Episode"); plt.ylabel("Total Reward"); plt.title("DQN Training on CartPole")
+plt.grid(True)
+plt.axhline(y=195, color='r', linestyle='--', label='Solved (195)')
+plt.legend()
+plt.show()
+""")
+
+md("""\
+## 思考题
+
+1. 为什么需要经验回放（Experience Replay）？在线学习会有什么问题？
+2. 目标网络（Target Network）解决了什么？如果不固定 target，Q 值会发散吗？
+3. ϵ-greedy 中的 ϵ 从 1.0 开始衰减有什么含义？
+4. 如果把 `hidden_dim` 从 128 改到 32，训练速度会怎样？收敛难度呢？
+""")
+
+nb.cells = cells
+with open("dqn/dqn.ipynb", "w") as f:
+    nbf.write(nb, f)
+print("Generated dqn/dqn.ipynb")

scripts/gen_gcn_notebook.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+"""Generate GCN notebook."""
+
+import nbformat as nbf
+
+nb = nbf.v4.new_notebook()
+nb.metadata = {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},"language_info": {"name": "python", "version": "3.12.0"}}
+
+cells = []
+def md(s): cells.append(nbf.v4.new_markdown_cell(s))
+def code(s): cells.append(nbf.v4.new_code_cell(s))
+
+md("# GCN: Graph Convolutional Network\n\nNode classification on citation graphs using spectral graph convolution.\n")
+
+md("""## 背景
+
+GCN（Kipf & Welling, 2017）将卷积操作推广到图结构数据。核心思想：
+每个节点的特征由其邻居节点加权聚合而来。
+
+与图像 CNN 的区别：
+- CNN：固定网格结构，卷积核在空间上滑动
+- GCN：任意图结构，卷积由邻接矩阵定义的消息传递实现
+
+数据集：**Cora** — 2708 篇论文，每篇用 1433 维词袋向量表示，分为 7 类。边表示引用关系。
+""")
+
+md("""## 数学原理
+
+### 图卷积层
+
+$$H^{(l+1)} = \\sigma\\left(\\hat{A} H^{(l)} W^{(l)}\\right)$$
+
+其中 $\\hat{A} = D^{-1/2} A D^{-1/2}$ 是归一化邻接矩阵。
+
+- $A$: 邻接矩阵（加自环后）
+- $D$: 度矩阵 $D_{ii} = \\sum_j A_{ij}$
+- $H^{(l)}$: 第 $l$ 层的节点表示
+- $W^{(l)}$: 可学习的权重矩阵
+
+### 2 层 GCN
+
+$$Z = \\text{softmax}\\left(\\hat{A}\\ \\text{ReLU}\\left(\\hat{A} X W^{(0)}\\right) W^{(1)}\\right)$$
+
+半监督学习：只用少量标注节点（每类 20 个）训练，模型通过图结构传播标签信息到未标注节点。
+""")
+
+code("""\
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from gcn.model import GCN
+from gcn.data import load_cora
+from utils.config import load_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+device = get_device()
+print(f"Device: {device}")
+
+features, adj_norm, labels, train_mask, val_mask, test_mask, classes = load_cora()
+features = features.to(device)
+adj_norm = adj_norm.to(device)
+labels = labels.to(device)
+train_mask = train_mask.to(device)
+val_mask = val_mask.to(device)
+""")
+
+code("""\
+model = GCN(
+    in_features=features.size(1),
+    hidden_dim=16,
+    num_classes=labels.max().item() + 1,
+    dropout=0.5,
+).to(device)
+print(f"Parameters: {model.num_params():,}")
+""")
+
+md("""## 训练
+
+> ⏱ 预估耗时：**200 epoch × ~0.1s/epoch ≈ 20 秒**（CPU 即可完成）
+""")
+
+code("""\
+NUM_EPOCHS = 200
+LR = 0.01
+WEIGHT_DECAY = 5e-4
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
+
+loss_hist, acc_hist = [], []
+
+for epoch in range(1, NUM_EPOCHS + 1):
+    model.train()
+    optimizer.zero_grad()
+    output = model(features, adj_norm)
+    loss = criterion(output[train_mask], labels[train_mask])
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.no_grad():
+        output = model(features, adj_norm)
+        val_acc = (output[val_mask].argmax(dim=1) == labels[val_mask]).float().mean().item()
+        loss_hist.append(loss.item())
+        acc_hist.append(val_acc)
+
+    if epoch % 20 == 0 or epoch == 1:
+        print(f"Epoch [{epoch:3d}/{NUM_EPOCHS}]  Loss: {loss.item():.4f}  Val Acc: {val_acc:.2%}")
+""")
+
+md("""## Loss 曲线""")
+
+code("""\
+import matplotlib.pyplot as plt
+
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
+ax1.plot(loss_hist); ax1.set_xlabel("Epoch"); ax1.set_ylabel("Loss"); ax1.set_title("Training Loss"); ax1.grid(True)
+ax2.plot(acc_hist, color='green'); ax2.set_xlabel("Epoch"); ax2.set_ylabel("Val Acc"); ax2.set_title("Validation Accuracy"); ax2.grid(True)
+plt.tight_layout(); plt.show()
+""")
+
+md("""## 测试准确率""")
+
+code("""\
+model.eval()
+with torch.no_grad():
+    output = model(features, adj_norm)
+    pred = output[test_mask].argmax(dim=1)
+    test_acc = (pred == labels[test_mask]).float().mean().item()
+print(f"Test Accuracy: {test_acc:.2%}")
+""")
+
+md("""\
+## 思考题
+
+1. 为什么 GCN 的归一化用 $D^{-1/2} A D^{-1/2}$ 而不是 $D^{-1} A$？
+2. GCN 能处理归纳式（inductive）任务吗？还是只能直推式（transductive）？
+3. 如果不用邻接矩阵只用节点特征，准确率会降到多少？
+4. GCN 层数加深为什么会导致性能下降？（提示：过平滑问题）
+""")
+
+nb.cells = cells
+with open("gcn/gcn.ipynb", "w") as f:
+    nbf.write(nb, f)
+print("Generated gcn/gcn.ipynb")

scripts/gen_simclr_notebook.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Generate SimCLR notebook."""
+
+import nbformat as nbf
+
+nb = nbf.v4.new_notebook()
+nb.metadata = {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},"language_info": {"name": "python", "version": "3.12.0"}}
+
+cells = []
+def md(s): cells.append(nbf.v4.new_markdown_cell(s))
+def code(s): cells.append(nbf.v4.new_code_cell(s))
+
+md("# SimCLR: Contrastive Learning\n\nSelf-supervised representation learning with NT-Xent loss on CIFAR-10.")
+
+md("""## 背景
+
+SimCLR（Chen et al. 2020）通过**对比学习**在没有标签的情况下学习图像表示。
+核心思路：同一张图的不同增强视图应该得到相似的表示，不同图的视图应该不同。
+
+关键组件：
+- **数据增强**：随机裁剪、颜色抖动、高斯模糊、灰度化
+- **Encoder**：ResNet18（去掉最后一层）
+- **Projector**：MLP 将表示投影到对比空间
+- **NT-Xent Loss**：归一化温度标度的交叉熵损失
+
+训练完成后，encoder 可以迁移到下游分类任务，只需加一个线性分类器。
+""")
+
+md("""## 数学原理
+
+### NT-Xent Loss
+
+对每个批次 $N$ 张图，生成两个增强视图，共 $2N$ 个样本：
+
+$$\ell(i, j) = -\\log \\frac{\\exp(\\text{sim}(z_i, z_j) / \\tau)}{\\sum_{k=1}^{2N} \\mathbb{1}_{[k \\neq i]} \\exp(\\text{sim}(z_i, z_k) / \\tau)}$$
+
+其中 $\\text{sim}(u, v) = \\frac{u^\\top v}{\\|u\\|\\|v\\|}$ 是余弦相似度，$(i, j)$ 是一对正样本（同一图的两种增强）。
+""")
+
+code("""\
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+
+from simclr.model import SimCLR
+from utils.config import load_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+device = get_device()
+print(f"Device: {device}")
+""")
+
+code("""\
+from simclr.data import SimCLRTransform, load_cifar10_simclr
+
+loader = load_cifar10_simclr(batch_size=256, num_workers=4)
+print(f"Batches per epoch: {len(loader)}")
+""")
+
+code("""\
+model = SimCLR(project_dim=128, temperature=0.5).to(device)
+print(f"Parameters: {model.num_params():,}")
+
+# Count encoder vs projector params
+enc = sum(p.numel() for p in model.encoder.parameters() if p.requires_grad)
+proj = sum(p.numel() for p in model.projector.parameters() if p.requires_grad)
+print(f"  Encoder (ResNet18): {enc:,}")
+print(f"  Projector (MLP):    {proj:,}")
+""")
+
+md("""## 训练
+
+> ⏱ 预估耗时：**100 epoch × ~40s/epoch ≈ 1 小时**（M4 Max, batch_size=256）
+> 如果太久，把下面 `NUM_EPOCHS` 改到 10 先看 loss 趋势。
+""")
+
+code("""\
+NUM_EPOCHS = 100
+LR = 0.0003
+
+optimizer = optim.Adam(model.parameters(), lr=LR)
+loss_hist = []
+
+for epoch in range(1, NUM_EPOCHS + 1):
+    model.train()
+    total_loss = 0.0
+    num_batches = 0
+    for batch in loader:
+        x1, x2 = batch["view1"].to(device), batch["view2"].to(device)
+        z1, z2 = model(x1), model(x2)
+        loss = model.nt_xent_loss(z1, z2)
+        optimizer.zero_grad(); loss.backward(); optimizer.step()
+        total_loss += loss.item(); num_batches += 1
+
+    avg = total_loss / num_batches
+    loss_hist.append(avg)
+    print(f"Epoch [{epoch:2d}/{NUM_EPOCHS}]  Loss: {avg:.4f}")
+""")
+
+md("""## Loss 曲线""")
+
+code("""\
+import matplotlib.pyplot as plt
+plt.plot(loss_hist)
+plt.xlabel("Epoch"); plt.ylabel("Loss"); plt.title("SimCLR Contrastive Loss"); plt.grid(True)
+plt.show()
+""")
+
+md("""\
+## 思考题
+
+1. SimCLR 为什么需要 Projector？直接用 encoder 的输出做对比学习效果会差吗？
+2. 数据增强的质量对对比学习有多重要？如果只用翻转，loss 会怎样？
+3. NT-Xent 中的 temperature $\\tau$ 起什么作用？增大/减小各有什么影响？
+4. SimCLR 为什么需要大 batch size？（提示：负样本数量）
+""")
+
+nb.cells = cells
+with open("simclr/simclr.ipynb", "w") as f:
+    nbf.write(nb, f)
+print("Generated simclr/simclr.ipynb")

scripts/gen_yolo_notebook.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""Generate YOLO notebook."""
+
+import nbformat as nbf
+
+nb = nbf.v4.new_notebook()
+nb.metadata = {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},"language_info": {"name": "python", "version": "3.12.0"}}
+
+cells = []
+def md(s): cells.append(nbf.v4.new_markdown_cell(s))
+def code(s): cells.append(nbf.v4.new_code_cell(s))
+
+md("# YOLO: You Only Look Once\n\nSimplified object detection with grid-based bounding box regression on Pascal VOC.")
+
+md("""## 背景
+
+YOLO（Redmon et al. 2016）是首个单阶段目标检测器，将检测视为回归问题。
+一张图通过 CNN 一次前向传播，直接输出边界框和类别概率。
+
+核心思想：将图像分成 $S \\times S$ 网格，每个网格预测 $B$ 个边界框和 $C$ 个类别的概率。
+
+与两阶段检测器（Faster R-CNN）的区别：
+- YOLO：一次前向 → 端到端，速度快但精度略低
+- Faster R-CNN：候选区域 → 分类，精度高但速度慢
+
+数据集：**Pascal VOC** — 20 类物体，含边界框标注。
+""")
+
+md("""## 数学原理
+
+### 输出表示
+
+每个网格单元预测 $B$ 个边界框，每个框 5 个值：
+
+$$(x, y, w, h, \\text{confidence})$$
+
+- $x, y$: 框中心相对于网格单元的偏移（0~1）
+- $w, h$: 框宽高相对于图像尺寸的比例
+- $\\text{confidence}$: $P(\\text{object}) \\times \\text{IoU}_{\\text{pred}}^{\\text{truth}}$
+
+再加上 $C$ 个类别概率 $P(\\text{class}_i \\mid \\text{object})$
+
+输出张量：$S \\times S \\times (B \\times 5 + C)$
+
+### 损失函数
+
+$$\\mathcal{L} = \\lambda_{\\text{coord}} \\sum \\mathbb{1}_{ij}^{\\text{obj}} [(x - \\hat{x})^2 + (y - \\hat{y})^2 + (\\sqrt{w} - \\sqrt{\\hat{w}})^2 + (\\sqrt{h} - \\sqrt{\\hat{h}})^2] + \\sum \\mathbb{1}_{ij}^{\\text{obj}} (C - \\hat{C})^2 + \\lambda_{\\text{noobj}} \\sum \\mathbb{1}_{ij}^{\\text{noobj}} (C - \\hat{C})^2 + \\sum \\mathbb{1}_{i}^{\\text{obj}} \\sum_{c=1}^C (p_i(c) - \\hat{p}_i(c))^2$$
+
+### 非极大值抑制（NMS）
+
+对同一类别的重叠框，保留得分最高的，移除与其 IoU 超过阈值的框。
+""")
+
+code("""\
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+
+from yolo.model import YOLO
+from yolo.loss import yolo_loss
+from utils.config import load_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+device = get_device()
+print(f"Device: {device}")
+""")
+
+code("""\
+from yolo.data import load_voc, VOC_CLASSES
+
+train_loader, test_loader = load_voc(
+    batch_size=32, image_size=224, S=7, B=2, C=20, num_workers=4,
+)
+print(f"Classes ({len(VOC_CLASSES)}): {VOC_CLASSES}")
+print(f"Train batches: {len(train_loader)}")
+""")
+
+code("""\
+model = YOLO(S=7, B=2, C=20).to(device)
+print(f"Parameters: {model.num_params():,}")
+""")
+
+md("""## 训练
+
+> ⏱ 预估耗时：**50 epoch × ~120s/epoch ≈ 1.5 小时**（M4 Max, batch_size=32）
+> 如果太久，把下面 `NUM_EPOCHS` 改到 5 先看 loss 趋势。
+""")
+
+code("""\
+NUM_EPOCHS = 50
+LR = 0.0001
+
+optimizer = optim.Adam(model.parameters(), lr=LR)
+loss_hist = []
+
+for epoch in range(1, NUM_EPOCHS + 1):
+    model.train()
+    total_loss = 0.0
+    for images, targets in train_loader:
+        images, targets = images.to(device), targets.to(device)
+        pred = model(images)
+        loss = yolo_loss(pred, targets, S=7, B=2, C=20, coord_scale=5, noobj_scale=0.5)
+        optimizer.zero_grad(); loss.backward(); optimizer.step()
+        total_loss += loss.item()
+
+    avg = total_loss / len(train_loader)
+    loss_hist.append(avg)
+    print(f"Epoch [{epoch:2d}/{NUM_EPOCHS}]  Loss: {avg:.4f}")
+""")
+
+md("""## Loss 曲线""")
+
+code("""\
+import matplotlib.pyplot as plt
+plt.plot(loss_hist)
+plt.xlabel("Epoch"); plt.ylabel("Loss"); plt.title("YOLO Training Loss"); plt.grid(True)
+plt.show()
+""")
+
+md("""\
+## 思考题
+
+1. YOLO 的 $S \\times S$ 网格中，一个网格只能预测一个物体（每个类）。这对检测小物体有什么影响？
+2. 为什么边界框的 $w, h$ 用平方根而不是直接用？这有什么物理意义？
+3. NMS 中的 IoU 阈值高低各有什么影响？
+4. YOLO 和两阶段检测器（Faster R-CNN）的核心区别是什么？
+""")
+
+nb.cells = cells
+with open("yolo/yolo.ipynb", "w") as f:
+    nbf.write(nb, f)
+print("Generated yolo/yolo.ipynb")

simclr/__init__.py

@@ -0,0 +1,3 @@
+from .model import SimCLR
+
+__all__ = ["SimCLR"]

simclr/config.yaml

@@ -0,0 +1,8 @@
+seed: 42
+batch_size: 256
+lr: 0.0003
+num_epochs: 100
+num_workers: 4
+temperature: 0.5
+project_dim: 128
+model_path: simclr/simclr_cifar10.pt

simclr/data.py

@@ -0,0 +1,41 @@
+"""CIFAR-10 with SimCLR augmentations (two random views per image)."""
+
+from torchvision import transforms
+from torch.utils.data import Dataset
+from datasets import load_dataset
+
+
+class SimCLRTransform:
+    """Random augmentation for SimCLR (crop + color jitter + flip + grayscale)."""
+
+    def __init__(self, image_size=32):
+        self.transform = transforms.Compose([
+            transforms.RandomResizedCrop(image_size, scale=(0.2, 1.0)),
+            transforms.RandomHorizontalFlip(),
+            transforms.ColorJitter(0.8, 0.8, 0.8, 0.2),
+            transforms.RandomGrayscale(p=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616]),
+        ])
+
+    def __call__(self, x):
+        return self.transform(x), self.transform(x)
+
+
+def load_cifar10_simclr(batch_size=256, num_workers=4):
+    transform = SimCLRTransform()
+
+    def transform_batch(batch):
+        images = [img.convert("RGB") for img in batch["img"]]
+        views1, views2 = zip(*[transform(img) for img in images])
+        batch["view1"] = list(views1)
+        batch["view2"] = list(views2)
+        return batch
+
+    ds = load_dataset("uoft-cs/cifar10", split="train")
+    ds.set_transform(transform_batch)
+
+    from torch.utils.data import DataLoader
+    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+
+    return loader

simclr/model.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from resnet18.model import ResNet, BasicBlock
+
+
+class Projector(nn.Module):
+    """MLP projection head."""
+
+    def __init__(self, in_dim=512, hidden_dim=256, out_dim=128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.BatchNorm1d(hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class SimCLR(nn.Module):
+    """SimCLR: contrastive learning with NT-Xent loss."""
+
+    def __init__(self, project_dim=128, temperature=0.5):
+        super().__init__()
+        self.temperature = temperature
+        # Encoder: ResNet18 without FC layer.
+        self.encoder = ResNet(BasicBlock, [2, 2, 2, 2])
+        # Remove the FC layer (avgpool is kept).
+        self.encoder.fc = nn.Identity()
+        self.projector = Projector(in_dim=512, hidden_dim=256, out_dim=project_dim)
+
+    def forward(self, x):
+        h = self.encoder(x)  # (B, 512)
+        return self.projector(h)  # (B, project_dim)
+
+    def nt_xent_loss(self, z1, z2):
+        """NT-Xent loss between two augmentation views.
+
+        z1, z2: (B, D) embeddings of two views.
+        """
+        B = z1.size(0)
+        z = torch.cat([z1, z2], dim=0)  # (2B, D)
+        z = F.normalize(z, dim=1)
+
+        # Cosine similarity matrix: (2B, 2B)
+        sim = z @ z.T / self.temperature
+
+        # Mask out self-similarity.
+        mask = torch.eye(2 * B, device=z.device, dtype=torch.bool)
+        sim = sim.masked_fill(mask, float("-inf"))
+
+        # Positive pairs: (i, i+B) and (i+B, i)
+        pos_mask = torch.zeros(2 * B, 2 * B, device=z.device, dtype=torch.bool)
+        for i in range(B):
+            pos_mask[i, i + B] = True
+            pos_mask[i + B, i] = True
+
+        # Compute loss for all 2B samples.
+        pos = sim[pos_mask].view(2 * B, 1)
+        neg = sim.masked_fill(pos_mask, float("-inf"))
+        logits = torch.cat([pos, neg], dim=1)
+        labels = torch.zeros(2 * B, device=z.device, dtype=torch.long)
+
+        return F.cross_entropy(logits, labels)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

simclr/simclr.ipynb

@@ -0,0 +1,188 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "403b8e0a",
+   "metadata": {},
+   "source": [
+    "# SimCLR: Contrastive Learning\n",
+    "\n",
+    "Self-supervised representation learning with NT-Xent loss on CIFAR-10."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0b4c57e2",
+   "metadata": {},
+   "source": [
+    "## 背景\n",
+    "\n",
+    "SimCLR（Chen et al. 2020）通过**对比学习**在没有标签的情况下学习图像表示。\n",
+    "核心思路：同一张图的不同增强视图应该得到相似的表示，不同图的视图应该不同。\n",
+    "\n",
+    "关键组件：\n",
+    "- **数据增强**：随机裁剪、颜色抖动、高斯模糊、灰度化\n",
+    "- **Encoder**：ResNet18（去掉最后一层）\n",
+    "- **Projector**：MLP 将表示投影到对比空间\n",
+    "- **NT-Xent Loss**：归一化温度标度的交叉熵损失\n",
+    "\n",
+    "训练完成后，encoder 可以迁移到下游分类任务，只需加一个线性分类器。\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f760a161",
+   "metadata": {},
+   "source": [
+    "## 数学原理\n",
+    "\n",
+    "### NT-Xent Loss\n",
+    "\n",
+    "对每个批次 $N$ 张图，生成两个增强视图，共 $2N$ 个样本：\n",
+    "\n",
+    "$$\\ell(i, j) = -\\log \\frac{\\exp(\\text{sim}(z_i, z_j) / \\tau)}{\\sum_{k=1}^{2N} \\mathbb{1}_{[k \\neq i]} \\exp(\\text{sim}(z_i, z_k) / \\tau)}$$\n",
+    "\n",
+    "其中 $\\text{sim}(u, v) = \\frac{u^\\top v}{\\|u\\|\\|v\\|}$ 是余弦相似度，$(i, j)$ 是一对正样本（同一图的两种增强）。\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9ed9383",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torchvision import transforms\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "from simclr.model import SimCLR\n",
+    "from utils.config import load_config\n",
+    "from utils.seed import set_seed\n",
+    "from utils.device import get_device\n",
+    "\n",
+    "device = get_device()\n",
+    "print(f\"Device: {device}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "745f1ad9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from simclr.data import SimCLRTransform, load_cifar10_simclr\n",
+    "\n",
+    "loader = load_cifar10_simclr(batch_size=256, num_workers=4)\n",
+    "print(f\"Batches per epoch: {len(loader)}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8684ccbb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = SimCLR(project_dim=128, temperature=0.5).to(device)\n",
+    "print(f\"Parameters: {model.num_params():,}\")\n",
+    "\n",
+    "# Count encoder vs projector params\n",
+    "enc = sum(p.numel() for p in model.encoder.parameters() if p.requires_grad)\n",
+    "proj = sum(p.numel() for p in model.projector.parameters() if p.requires_grad)\n",
+    "print(f\"  Encoder (ResNet18): {enc:,}\")\n",
+    "print(f\"  Projector (MLP):    {proj:,}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c04f9912",
+   "metadata": {},
+   "source": [
+    "## 训练\n",
+    "\n",
+    "> ⏱ 预估耗时：**100 epoch × ~40s/epoch ≈ 1 小时**（M4 Max, batch_size=256）\n",
+    "> 如果太久，把下面 `NUM_EPOCHS` 改到 10 先看 loss 趋势。\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e7d15e7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_EPOCHS = 100\n",
+    "LR = 0.0003\n",
+    "\n",
+    "optimizer = optim.Adam(model.parameters(), lr=LR)\n",
+    "loss_hist = []\n",
+    "\n",
+    "for epoch in range(1, NUM_EPOCHS + 1):\n",
+    "    model.train()\n",
+    "    total_loss = 0.0\n",
+    "    num_batches = 0\n",
+    "    for batch in loader:\n",
+    "        x1, x2 = batch[\"view1\"].to(device), batch[\"view2\"].to(device)\n",
+    "        z1, z2 = model(x1), model(x2)\n",
+    "        loss = model.nt_xent_loss(z1, z2)\n",
+    "        optimizer.zero_grad(); loss.backward(); optimizer.step()\n",
+    "        total_loss += loss.item(); num_batches += 1\n",
+    "\n",
+    "    avg = total_loss / num_batches\n",
+    "    loss_hist.append(avg)\n",
+    "    print(f\"Epoch [{epoch:2d}/{NUM_EPOCHS}]  Loss: {avg:.4f}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4c81184a",
+   "metadata": {},
+   "source": [
+    "## Loss 曲线"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8d146211",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "plt.plot(loss_hist)\n",
+    "plt.xlabel(\"Epoch\"); plt.ylabel(\"Loss\"); plt.title(\"SimCLR Contrastive Loss\"); plt.grid(True)\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "66c5af41",
+   "metadata": {},
+   "source": [
+    "## 思考题\n",
+    "\n",
+    "1. SimCLR 为什么需要 Projector？直接用 encoder 的输出做对比学习效果会差吗？\n",
+    "2. 数据增强的质量对对比学习有多重要？如果只用翻转，loss 会怎样？\n",
+    "3. NT-Xent 中的 temperature $\\tau$ 起什么作用？增大/减小各有什么影响？\n",
+    "4. SimCLR 为什么需要大 batch size？（提示：负样本数量）\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

simclr/train.py

@@ -0,0 +1,63 @@
+import torch
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+
+from simclr.model import SimCLR
+from simclr.data import load_cifar10_simclr
+from utils.config import load_config, save_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+
+def train():
+    cfg = load_config("simclr/config.yaml")
+    set_seed(cfg["seed"])
+
+    device = get_device()
+    print(f"Device: {device}")
+    torch.set_num_threads(4)
+
+    loader = load_cifar10_simclr(batch_size=cfg["batch_size"], num_workers=cfg["num_workers"])
+    print(f"Dataset: 50,000 CIFAR-10 images (with augmentations)")
+
+    model = SimCLR(project_dim=cfg["project_dim"], temperature=cfg["temperature"]).to(device)
+    print(f"Parameters: {model.num_params():,}")
+
+    optimizer = optim.Adam(model.parameters(), lr=cfg["lr"])
+
+    num_epochs = cfg["num_epochs"]
+    writer = SummaryWriter(log_dir="runs/simclr")
+
+    for epoch in range(1, num_epochs + 1):
+        model.train()
+        total_loss = 0.0
+        num_batches = 0
+
+        for batch in loader:
+            x1 = batch["view1"].to(device)
+            x2 = batch["view2"].to(device)
+
+            z1 = model(x1)
+            z2 = model(x2)
+            loss = model.nt_xent_loss(z1, z2)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            num_batches += 1
+
+        avg_loss = total_loss / num_batches
+        writer.add_scalar("train/loss", avg_loss, epoch)
+        print(f"Epoch [{epoch:2d}/{num_epochs}]  Loss: {avg_loss:.4f}")
+
+    writer.close()
+    save_path = cfg["model_path"]
+    torch.save(model.state_dict(), save_path)
+    save_config(cfg, save_path.replace(".pt", "_config.yaml"))
+    print(f"\nModel saved to {save_path}")
+
+
+if __name__ == "__main__":
+    train()

uv.lock

@@ -445,6 +445,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" },
 ]
 
+[[package]]
+name = "cloudpickle"
+version = "3.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" },
+]
+
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -534,7 +543,7 @@ name = "cuda-bindings"
 version = "13.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "cuda-pathfinder" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" },
@@ -565,34 +574,34 @@ wheels = [
 
 [package.optional-dependencies]
 cudart = [
-    { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-runtime" },
 ]
 cufft = [
-    { name = "nvidia-cufft", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cufft" },
 ]
 cufile = [
-    { name = "nvidia-cufile", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cufile" },
 ]
 cupti = [
-    { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-cupti" },
 ]
 curand = [
-    { name = "nvidia-curand", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-curand" },
 ]
 cusolver = [
-    { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cusolver" },
 ]
 cusparse = [
-    { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cusparse" },
 ]
 nvjitlink = [
-    { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink" },
 ]
 nvrtc = [
-    { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cuda-nvrtc" },
 ]
 nvtx = [
-    { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvtx" },
 ]
 
 [[package]]
@@ -683,6 +692,7 @@ version = "0.1.0"
 source = { virtual = "." }
 dependencies = [
     { name = "datasets" },
+    { name = "gymnasium" },
     { name = "ipykernel" },
     { name = "jupyterlab" },
     { name = "matplotlib" },
@@ -699,6 +709,7 @@ dependencies = [
 [package.metadata]
 requires-dist = [
     { name = "datasets", specifier = ">=5.0.0" },
+    { name = "gymnasium", specifier = ">=0.29" },
     { name = "ipykernel", specifier = ">=6.0" },
     { name = "jupyterlab", specifier = ">=4.0" },
     { name = "matplotlib", specifier = ">=3.0" },
@@ -721,6 +732,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017", size = 28317, upload-time = "2025-09-01T09:48:08.5Z" },
 ]
 
+[[package]]
+name = "farama-notifications"
+version = "0.0.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ec/91/14397890dde30adc4bee6462158933806207bc5dd10d7b4d09d5c33845cf/farama_notifications-0.0.6.tar.gz", hash = "sha256:b19acac4bb41d76e59e03394b5dd165f4761c86fa327f56307a35cbee3b60158", size = 2517, upload-time = "2026-04-24T08:43:57.603Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl", hash = "sha256:f84839188efa1ce5bb361c2a84881b2dc2c0d0d7fb661ff00421820170930935", size = 2897, upload-time = "2026-04-24T08:43:56.785Z" },
+]
+
 [[package]]
 name = "fastjsonschema"
 version = "2.21.2"
@@ -933,6 +953,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ff/b0/4af731ff7492c68a96e4c71bfd0f4590acde92b31c6fe4894e6465c10ff6/grpcio-1.81.1-cp314-cp314-win_amd64.whl", hash = "sha256:3768a5ff1b2125e6f552e561b6b2dca0e64982d8949689b4df145cf8b98d7821", size = 5070275, upload-time = "2026-06-11T12:46:48.486Z" },
 ]
 
+[[package]]
+name = "gymnasium"
+version = "1.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cloudpickle" },
+    { name = "farama-notifications" },
+    { name = "numpy" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/4d/ff/14b6880d703dfaca204490979d3254ccd280c99550798993319902873658/gymnasium-1.3.0.tar.gz", hash = "sha256:6939e86e835d6b71b6ba6bfd360487420876deafc79bfb7bacba83a7c446bcf3", size = 830646, upload-time = "2026-04-22T13:47:14.155Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e9/73/fda6a25f3beeb5e49d74330b44092b9e5a547395ccd478d1103ddcbff1fc/gymnasium-1.3.0-py3-none-any.whl", hash = "sha256:6b8c159a8540dcbcb221722d7efda24d78ebbcbc3bd2ea1c2611aa2a34471fc2", size = 953904, upload-time = "2026-04-22T13:47:12.13Z" },
+]
+
 [[package]]
 name = "h11"
 version = "0.16.0"
@@ -1898,7 +1933,7 @@ name = "nvidia-cublas"
 version = "13.1.1.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-cuda-nvrtc" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" },
@@ -1937,7 +1972,7 @@ name = "nvidia-cudnn-cu13"
 version = "9.20.0.48"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-cublas" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" },
@@ -1949,7 +1984,7 @@ name = "nvidia-cufft"
 version = "12.0.0.61"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-nvjitlink" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
@@ -1979,9 +2014,9 @@ name = "nvidia-cusolver"
 version = "12.0.4.66"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
-    { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
-    { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-cublas" },
+    { name = "nvidia-cusparse" },
+    { name = "nvidia-nvjitlink" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
@@ -1993,7 +2028,7 @@ name = "nvidia-cusparse"
 version = "12.6.3.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-nvjitlink" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
@@ -2129,7 +2164,7 @@ name = "pexpect"
 version = "4.9.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "ptyprocess" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
 wheels = [

yolo/__init__.py

@@ -0,0 +1,4 @@
+from .model import YOLO
+from .loss import yolo_loss, nms
+
+__all__ = ["YOLO", "yolo_loss", "nms"]

yolo/config.yaml

@@ -0,0 +1,12 @@
+seed: 42
+batch_size: 32
+lr: 0.0001
+num_epochs: 50
+num_workers: 4
+image_size: 224
+S: 7            # grid size
+B: 2            # boxes per cell
+C: 20           # Pascal VOC classes
+coord_scale: 5
+noobj_scale: 0.5
+model_path: yolo/yolo_voc.pt

yolo/data.py

@@ -0,0 +1,107 @@
+"""Pascal VOC dataset for YOLO."""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+
+
+VOC_CLASSES = [
+    "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car",
+    "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike",
+    "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor",
+]
+
+
+def _build_transform(image_size=224):
+    return transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+    ])
+
+
+def _voc_to_yolo(annotation, S, B, C, image_size=224):
+    """Convert VOC annotation to YOLO target tensor (S, S, B*5+C)."""
+    target = torch.zeros(S, S, B * 5 + C)
+
+    for obj in annotation:
+        try:
+            label = obj["category"]
+            bbox = obj["bbox"]  # [xmin, ymin, xmax, ymax]
+            xmin, ymin, xmax, ymax = bbox
+            w = xmax - xmin
+            h = ymax - ymin
+
+            # Skip invalid boxes.
+            if w <= 0 or h <= 0:
+                continue
+
+            # Center and normalize.
+            x_center = (xmin + xmax) / 2 / image_size
+            y_center = (ymin + ymax) / 2 / image_size
+            w_norm = w / image_size
+            h_norm = h / image_size
+
+            # Grid cell.
+            col = int(x_center * S)
+            row = int(y_center * S)
+            if col >= S or row >= S:
+                continue
+
+            # Relative position within cell.
+            x_cell = x_center * S - col
+            y_cell = y_center * S - row
+
+            class_idx = label
+            if class_idx < 0 or class_idx >= C:
+                continue
+
+            # Assign to first available box.
+            for b in range(B):
+                box_start = b * 5
+                if target[row, col, box_start + 4] == 0:  # confidence unused
+                    target[row, col, box_start:box_start + 5] = torch.tensor([
+                        x_cell, y_cell, w_norm, h_norm, 1.0,
+                    ])
+                    target[row, col, B * 5 + class_idx] = 1.0
+                    break
+        except (KeyError, TypeError, IndexError):
+            continue
+
+    return target
+
+
+class VOCDataset(Dataset):
+    def __init__(self, split="train", image_size=224, S=7, B=2, C=20):
+        self.S, self.B, self.C = S, B, C
+        self.image_size = image_size
+        self.transform = _build_transform(image_size)
+
+        ds = load_dataset("widerface/pascal_voc", split="train")
+        self.ds = ds
+
+    def __len__(self):
+        return len(self.ds)
+
+    def __getitem__(self, idx):
+        item = self.ds[idx]
+        image = item["image"].convert("RGB")
+        image = self.transform(image)
+
+        try:
+            objects = item["objects"]
+        except (KeyError, TypeError):
+            objects = []
+
+        target = _voc_to_yolo(objects, self.S, self.B, self.C, self.image_size)
+        return image, target
+
+
+def load_voc(batch_size=32, image_size=224, S=7, B=2, C=20, num_workers=4):
+    train_dataset = VOCDataset(split="train", image_size=image_size, S=S, B=B, C=C)
+    test_dataset = VOCDataset(split="test", image_size=image_size, S=S, B=B, C=C)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+    return train_loader, test_loader

yolo/loss.py

@@ -0,0 +1,104 @@
+"""YOLO loss and NMS utilities."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def yolo_loss(pred, target, S=7, B=2, C=20, coord_scale=5, noobj_scale=0.5):
+    """YOLO loss: coord + obj + noobj + class."""
+    pred = pred.view(-1, S, S, B * 5 + C)
+    target = target.view(-1, S, S, B * 5 + C)
+
+    obj_mask = target[..., 4] > 0  # (N, S, S) — first box has obj
+
+    # ── Box coordinates loss (only for obj cells, best box) ──
+    coord_loss = torch.tensor(0.0, device=pred.device)
+    if obj_mask.any():
+        pred_box = pred[obj_mask]
+        target_box = target[obj_mask]
+        # Take first box for simplicity (would IoU match in real YOLO).
+        coord_loss = F.mse_loss(pred_box[:, :4], target_box[:, :4], reduction="sum")
+
+    # ── Confidence loss ──
+    obj_conf = pred[..., 4]
+    target_conf = target[..., 4]
+    obj_loss = F.mse_loss(obj_conf[obj_mask], target_conf[obj_mask], reduction="sum")
+    noobj_loss = F.mse_loss(obj_conf[~obj_mask], target_conf[~obj_mask], reduction="sum")
+
+    # ── Class loss ──
+    class_loss = torch.tensor(0.0, device=pred.device)
+    if obj_mask.any():
+        class_loss = F.mse_loss(
+            pred[obj_mask][:, B * 5:],
+            target[obj_mask][:, B * 5:],
+            reduction="sum",
+        )
+
+    N = obj_mask.numel()
+    return (coord_scale * coord_loss + obj_loss + noobj_scale * noobj_loss + class_loss) / N
+
+
+def nms(predictions, conf_threshold=0.3, iou_threshold=0.5):
+    """Non-Maximum Suppression — simplified for single image."""
+    boxes, scores, labels = [], [], []
+    S, B, C = 7, 2, 20
+
+    for row in range(S):
+        for col in range(S):
+            for b in range(B):
+                box = predictions[row, col, b * 5:(b + 1) * 5]
+                conf = box[4].item()
+                if conf < conf_threshold:
+                    continue
+                x, y, w, h = box[:4].tolist()
+                # Convert to absolute coordinates.
+                x_abs = (col + x) / S * 224
+                y_abs = (row + y) / S * 224
+                w_abs = w * 224
+                h_abs = h * 224
+                x1 = x_abs - w_abs / 2
+                y1 = y_abs - h_abs / 2
+                x2 = x_abs + w_abs / 2
+                y2 = y_abs + h_abs / 2
+
+                class_probs = predictions[row, col, B * 5:].softmax(dim=0)
+                class_score, class_idx = class_probs.max(dim=0)
+
+                score = conf * class_score.item()
+                boxes.append([x1, y1, x2, y2])
+                scores.append(score)
+                labels.append(class_idx.item())
+
+    if not boxes:
+        return [], [], []
+
+    boxes = torch.tensor(boxes)
+    scores = torch.tensor(scores)
+    labels = torch.tensor(labels)
+
+    # Sort by score.
+    _, order = scores.sort(descending=True)
+    boxes, scores, labels = boxes[order], scores[order], labels[order]
+
+    keep = []
+    while boxes.size(0) > 0:
+        keep.append(0)
+        if boxes.size(0) == 1:
+            break
+        ious = _compute_iou(boxes[0:1], boxes[1:])
+        mask = ious < iou_threshold
+        boxes, scores, labels = boxes[1:][mask], scores[1:][mask], labels[1:][mask]
+
+    return boxes[keep], scores[keep], labels[keep]
+
+
+def _compute_iou(box1, box2):
+    x1 = torch.max(box1[:, 0], box2[:, 0])
+    y1 = torch.max(box1[:, 1], box2[:, 1])
+    x2 = torch.min(box1[:, 2], box2[:, 2])
+    y2 = torch.min(box1[:, 3], box2[:, 3])
+    inter = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)
+    area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
+    area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
+    return inter / (area1 + area2 - inter).clamp(min=1e-8)

yolo/model.py

@@ -0,0 +1,39 @@
+"""Simplified YOLO: CNN backbone + detection head."""
+
+import torch
+import torch.nn as nn
+
+
+class YOLO(nn.Module):
+    """Simplified YOLO detector (like YOLOv1)."""
+
+    def __init__(self, S=7, B=2, C=20):
+        super().__init__()
+        self.S, self.B, self.C = S, B, C
+
+        # CNN backbone: 224 → 112 → 56 → 28 → 14 → 7.
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.LeakyReLU(0.1), nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, 1, 1), nn.BatchNorm2d(128), nn.LeakyReLU(0.1), nn.MaxPool2d(2),
+            nn.Conv2d(128, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.LeakyReLU(0.1), nn.MaxPool2d(2),
+            nn.Conv2d(256, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.LeakyReLU(0.1), nn.MaxPool2d(2),
+            nn.Conv2d(512, 1024, 3, 1, 1), nn.BatchNorm2d(1024), nn.LeakyReLU(0.1), nn.MaxPool2d(2),
+            nn.AdaptiveAvgPool2d((S, S)),
+        )
+
+        # Detection head.
+        self.det_head = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(1024 * S * S, 1024),
+            nn.LeakyReLU(0.1),
+            nn.Linear(1024, S * S * (B * 5 + C)),
+        )
+
+    def forward(self, x):
+        B = x.size(0)
+        x = self.features(x)
+        x = self.det_head(x)
+        return x.view(B, self.S, self.S, self.B * 5 + self.C)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

yolo/train.py

@@ -0,0 +1,65 @@
+import torch
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+
+from yolo.model import YOLO
+from yolo.loss import yolo_loss
+from yolo.data import load_voc, VOC_CLASSES
+from utils.config import load_config, save_config
+from utils.seed import set_seed
+from utils.device import get_device
+
+
+def train():
+    cfg = load_config("yolo/config.yaml")
+    set_seed(cfg["seed"])
+
+    device = get_device()
+    print(f"Device: {device}")
+    torch.set_num_threads(4)
+
+    train_loader, test_loader = load_voc(
+        batch_size=cfg["batch_size"], image_size=cfg["image_size"],
+        S=cfg["S"], B=cfg["B"], C=cfg["C"], num_workers=cfg["num_workers"],
+    )
+    print(f"Train batches: {len(train_loader)}, Test batches: {len(test_loader)}")
+
+    model = YOLO(S=cfg["S"], B=cfg["B"], C=cfg["C"]).to(device)
+    print(f"Parameters: {model.num_params():,}")
+
+    optimizer = optim.Adam(model.parameters(), lr=cfg["lr"])
+
+    num_epochs = cfg["num_epochs"]
+    writer = SummaryWriter(log_dir="runs/yolo")
+
+    for epoch in range(1, num_epochs + 1):
+        model.train()
+        total_loss = 0.0
+        num_batches = 0
+
+        for images, targets in train_loader:
+            images, targets = images.to(device), targets.to(device)
+            pred = model(images)
+            loss = yolo_loss(pred, targets, cfg["S"], cfg["B"], cfg["C"],
+                             cfg["coord_scale"], cfg["noobj_scale"])
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            num_batches += 1
+
+        avg_loss = total_loss / num_batches
+        writer.add_scalar("train/loss", avg_loss, epoch)
+        print(f"Epoch [{epoch:2d}/{num_epochs}]  Loss: {avg_loss:.4f}")
+
+    writer.close()
+    save_path = cfg["model_path"]
+    torch.save(model.state_dict(), save_path)
+    save_config(cfg, save_path.replace(".pt", "_config.yaml"))
+    print(f"\nModel saved to {save_path}")
+
+
+if __name__ == "__main__":
+    train()

yolo/yolo.ipynb

@@ -0,0 +1,199 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e6ec2ac3",
+   "metadata": {},
+   "source": [
+    "# YOLO: You Only Look Once\n",
+    "\n",
+    "Simplified object detection with grid-based bounding box regression on Pascal VOC."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f3030108",
+   "metadata": {},
+   "source": [
+    "## 背景\n",
+    "\n",
+    "YOLO（Redmon et al. 2016）是首个单阶段目标检测器，将检测视为回归问题。\n",
+    "一张图通过 CNN 一次前向传播，直接输出边界框和类别概率。\n",
+    "\n",
+    "核心思想：将图像分成 $S \\times S$ 网格，每个网格预测 $B$ 个边界框和 $C$ 个类别的概率。\n",
+    "\n",
+    "与两阶段检测器（Faster R-CNN）的区别：\n",
+    "- YOLO：一次前向 → 端到端，速度快但精度略低\n",
+    "- Faster R-CNN：候选区域 → 分类，精度高但速度慢\n",
+    "\n",
+    "数据集：**Pascal VOC** — 20 类物体，含边界框标注。\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "88a406f0",
+   "metadata": {},
+   "source": [
+    "## 数学原理\n",
+    "\n",
+    "### 输出表示\n",
+    "\n",
+    "每个网格单元预测 $B$ 个边界框，每个框 5 个值：\n",
+    "\n",
+    "$$(x, y, w, h, \\text{confidence})$$\n",
+    "\n",
+    "- $x, y$: 框中心相对于网格单元的偏移（0~1）\n",
+    "- $w, h$: 框宽高相对于图像尺寸的比例\n",
+    "- $\\text{confidence}$: $P(\\text{object}) \\times \\text{IoU}_{\\text{pred}}^{\\text{truth}}$\n",
+    "\n",
+    "再加上 $C$ 个类别概率 $P(\\text{class}_i \\mid \\text{object})$\n",
+    "\n",
+    "输出张量：$S \\times S \\times (B \\times 5 + C)$\n",
+    "\n",
+    "### 损失函数\n",
+    "\n",
+    "$$\\mathcal{L} = \\lambda_{\\text{coord}} \\sum \\mathbb{1}_{ij}^{\\text{obj}} [(x - \\hat{x})^2 + (y - \\hat{y})^2 + (\\sqrt{w} - \\sqrt{\\hat{w}})^2 + (\\sqrt{h} - \\sqrt{\\hat{h}})^2] + \\sum \\mathbb{1}_{ij}^{\\text{obj}} (C - \\hat{C})^2 + \\lambda_{\\text{noobj}} \\sum \\mathbb{1}_{ij}^{\\text{noobj}} (C - \\hat{C})^2 + \\sum \\mathbb{1}_{i}^{\\text{obj}} \\sum_{c=1}^C (p_i(c) - \\hat{p}_i(c))^2$$\n",
+    "\n",
+    "### 非极大值抑制（NMS）\n",
+    "\n",
+    "对同一类别的重叠框，保留得分最高的，移除与其 IoU 超过阈值的框。\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9e3b157",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torchvision import transforms\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "from yolo.model import YOLO\n",
+    "from yolo.loss import yolo_loss\n",
+    "from utils.config import load_config\n",
+    "from utils.seed import set_seed\n",
+    "from utils.device import get_device\n",
+    "\n",
+    "device = get_device()\n",
+    "print(f\"Device: {device}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fcf1335f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from yolo.data import load_voc, VOC_CLASSES\n",
+    "\n",
+    "train_loader, test_loader = load_voc(\n",
+    "    batch_size=32, image_size=224, S=7, B=2, C=20, num_workers=4,\n",
+    ")\n",
+    "print(f\"Classes ({len(VOC_CLASSES)}): {VOC_CLASSES}\")\n",
+    "print(f\"Train batches: {len(train_loader)}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0c4c1b08",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = YOLO(S=7, B=2, C=20).to(device)\n",
+    "print(f\"Parameters: {model.num_params():,}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b23a06f5",
+   "metadata": {},
+   "source": [
+    "## 训练\n",
+    "\n",
+    "> ⏱ 预估耗时：**50 epoch × ~120s/epoch ≈ 1.5 小时**（M4 Max, batch_size=32）\n",
+    "> 如果太久，把下面 `NUM_EPOCHS` 改到 5 先看 loss 趋势。\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a0696ca0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_EPOCHS = 50\n",
+    "LR = 0.0001\n",
+    "\n",
+    "optimizer = optim.Adam(model.parameters(), lr=LR)\n",
+    "loss_hist = []\n",
+    "\n",
+    "for epoch in range(1, NUM_EPOCHS + 1):\n",
+    "    model.train()\n",
+    "    total_loss = 0.0\n",
+    "    for images, targets in train_loader:\n",
+    "        images, targets = images.to(device), targets.to(device)\n",
+    "        pred = model(images)\n",
+    "        loss = yolo_loss(pred, targets, S=7, B=2, C=20, coord_scale=5, noobj_scale=0.5)\n",
+    "        optimizer.zero_grad(); loss.backward(); optimizer.step()\n",
+    "        total_loss += loss.item()\n",
+    "\n",
+    "    avg = total_loss / len(train_loader)\n",
+    "    loss_hist.append(avg)\n",
+    "    print(f\"Epoch [{epoch:2d}/{NUM_EPOCHS}]  Loss: {avg:.4f}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b1965c8",
+   "metadata": {},
+   "source": [
+    "## Loss 曲线"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "29b8d6d2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "plt.plot(loss_hist)\n",
+    "plt.xlabel(\"Epoch\"); plt.ylabel(\"Loss\"); plt.title(\"YOLO Training Loss\"); plt.grid(True)\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "73e98cc8",
+   "metadata": {},
+   "source": [
+    "## 思考题\n",
+    "\n",
+    "1. YOLO 的 $S \\times S$ 网格中，一个网格只能预测一个物体（每个类）。这对检测小物体有什么影响？\n",
+    "2. 为什么边界框的 $w, h$ 用平方根而不是直接用？这有什么物理意义？\n",
+    "3. NMS 中的 IoU 阈值高低各有什么影响？\n",
+    "4. YOLO 和两阶段检测器（Faster R-CNN）的核心区别是什么？\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}