diff --git a/solutions/07_batchnorm_solution.ipynb b/solutions/07_batchnorm_solution.ipynb index 60d7c86..7ddca6a 100644 --- a/solutions/07_batchnorm_solution.ipynb +++ b/solutions/07_batchnorm_solution.ipynb @@ -2,11 +2,12 @@ "cells": [ { "cell_type": "markdown", + "id": "ffd42526", "metadata": {}, "source": [ "# 🟡 Solution: Implement BatchNorm\n", "\n", - "Reference solution for Batch Normalization (training mode)." + "Reference solution for Batch Normalization with both **training** and **inference** behavior, including running mean/variance updates." ] }, { @@ -27,9 +28,37 @@ "source": [ "# ✅ SOLUTION\n", "\n", - "def my_batch_norm(x, gamma, beta, eps=1e-5):\n", - " mean = x.mean(dim=0)\n", - " var = x.var(dim=0, unbiased=False)\n", + "import torch\n", + "\n", + "def my_batch_norm(\n", + " x,\n", + " gamma,\n", + " beta,\n", + " running_mean,\n", + " running_var,\n", + " eps=1e-5,\n", + " momentum=0.1,\n", + " training=True,\n", + "):\n", + " \"\"\"BatchNorm with train/eval behavior and running stats.\n", + "\n", + " - Training: use batch stats, update running_mean / running_var in-place.\n", + " - Inference: use running_mean / running_var as-is.\n", + " \"\"\"\n", + " if training:\n", + " batch_mean = x.mean(dim=0)\n", + " batch_var = x.var(dim=0, unbiased=False)\n", + "\n", + " # Update running statistics in-place. Detach to avoid tracking gradients.\n", + " running_mean.mul_(1 - momentum).add_(momentum * batch_mean.detach())\n", + " running_var.mul_(1 - momentum).add_(momentum * batch_var.detach())\n", + "\n", + " mean = batch_mean\n", + " var = batch_var\n", + " else:\n", + " mean = running_mean\n", + " var = running_var\n", + "\n", " x_norm = (x - mean) / torch.sqrt(var + eps)\n", " return gamma * x_norm + beta" ] @@ -37,14 +66,28 @@ { "cell_type": "code", "execution_count": null, + "id": "dbd7bb4e", "metadata": {}, "outputs": [], "source": [ "# Verify\n", "x = torch.randn(8, 4)\n", - "out = my_batch_norm(x, torch.ones(4), torch.zeros(4))\n", - "print(\"Column means:\", out.mean(dim=0))\n", - "print(\"Column stds: \", out.std(dim=0))" + "gamma = torch.ones(4)\n", + "beta = torch.zeros(4)\n", + "\n", + "running_mean = torch.zeros(4)\n", + "running_var = torch.ones(4)\n", + "\n", + "# Training behavior: normalize with batch stats and update running stats\n", + "out_train = my_batch_norm(x, gamma, beta, running_mean, running_var, training=True)\n", + "print(\"[Train] Column means:\", out_train.mean(dim=0))\n", + "print(\"[Train] Column stds: \", out_train.std(dim=0))\n", + "print(\"Updated running_mean:\", running_mean)\n", + "print(\"Updated running_var:\", running_var)\n", + "\n", + "# Inference behavior: use running_mean / running_var only\n", + "out_eval = my_batch_norm(x, gamma, beta, running_mean, running_var, training=False)\n", + "print(\"[Eval] Column means (using running stats):\", out_eval.mean(dim=0))" ] }, { diff --git a/templates/07_batchnorm.ipynb b/templates/07_batchnorm.ipynb index ff1218c..5307942 100644 --- a/templates/07_batchnorm.ipynb +++ b/templates/07_batchnorm.ipynb @@ -2,28 +2,45 @@ "cells": [ { "cell_type": "markdown", + "id": "89fd15cb", "metadata": {}, "source": [ "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb)\n", "\n", "# 🟡 Medium: Implement BatchNorm\n", "\n", - "Implement **Batch Normalization** (training mode) from scratch.\n", + "Implement **Batch Normalization** with both **training** and **inference** behavior.\n", + "\n", + "In training mode, use **batch statistics** and update running estimates:\n", "\n", "$$\\text{BN}(x) = \\gamma \\cdot \\frac{x - \\mu_B}{\\sqrt{\\sigma_B^2 + \\epsilon}} + \\beta$$\n", "\n", "where $\\mu_B$ and $\\sigma_B^2$ are the mean and variance computed **across the batch** (dim=0).\n", "\n", + "In inference mode, use the provided **running mean/var** instead of current batch stats.\n", + "\n", "### Signature\n", "```python\n", - "def my_batch_norm(x: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:\n", + "def my_batch_norm(\n", + " x: torch.Tensor,\n", + " gamma: torch.Tensor,\n", + " beta: torch.Tensor,\n", + " running_mean: torch.Tensor,\n", + " running_var: torch.Tensor,\n", + " eps: float = 1e-5,\n", + " momentum: float = 0.1,\n", + " training: bool = True,\n", + ") -> torch.Tensor:\n", " # x: (N, D) — normalize each feature across all samples in the batch\n", + " # running_mean, running_var: updated in-place during training; used as-is during inference\n", "```\n", "\n", "### Rules\n", "- Do **NOT** use `F.batch_norm`, `nn.BatchNorm1d`, etc.\n", - "- Compute mean and variance over `dim=0` with `unbiased=False`\n", - "- Must support autograd" + "- Compute batch mean and variance over `dim=0` with `unbiased=False`\n", + "- Update running stats like PyTorch: `running = (1 - momentum) * running + momentum * batch_stat`\n", + "- Use `running_mean` / `running_var` for inference when `training=False`\n", + "- Must support autograd w.r.t. `x`, `gamma`, `beta`（running statistics 应视作 buffer，而不是需要梯度的参数）" ] }, { @@ -44,13 +61,23 @@ "source": [ "# ✏️ YOUR IMPLEMENTATION HERE\n", "\n", - "def my_batch_norm(x, gamma, beta, eps=1e-5):\n", + "def my_batch_norm(\n", + " x,\n", + " gamma,\n", + " beta,\n", + " running_mean,\n", + " running_var,\n", + " eps=1e-5,\n", + " momentum=0.1,\n", + " training=True,\n", + "):\n", " pass # Replace this" ] }, { "cell_type": "code", "execution_count": null, + "id": "26b93e71", "metadata": {}, "outputs": [], "source": [ @@ -58,10 +85,22 @@ "x = torch.randn(8, 4)\n", "gamma = torch.ones(4)\n", "beta = torch.zeros(4)\n", - "out = my_batch_norm(x, gamma, beta)\n", - "print(\"Output shape:\", out.shape)\n", - "print(\"Column means:\", out.mean(dim=0)) # should be ~0\n", - "print(\"Column stds: \", out.std(dim=0)) # should be ~1" + "\n", + "# Running stats typically live on the same device and shape as features\n", + "running_mean = torch.zeros(4)\n", + "running_var = torch.ones(4)\n", + "\n", + "# Training mode: uses batch stats and updates running_mean / running_var\n", + "out_train = my_batch_norm(x, gamma, beta, running_mean, running_var, training=True)\n", + "print(\"[Train] Output shape:\", out_train.shape)\n", + "print(\"[Train] Column means:\", out_train.mean(dim=0)) # should be ~0\n", + "print(\"[Train] Column stds: \", out_train.std(dim=0)) # should be ~1\n", + "print(\"Updated running_mean:\", running_mean)\n", + "print(\"Updated running_var:\", running_var)\n", + "\n", + "# Inference mode: uses running_mean / running_var only\n", + "out_eval = my_batch_norm(x, gamma, beta, running_mean, running_var, training=False)\n", + "print(\"[Eval] Output shape:\", out_eval.shape)" ] }, { diff --git a/torch_judge/tasks/batchnorm.py b/torch_judge/tasks/batchnorm.py index 7ad6c15..74e972a 100644 --- a/torch_judge/tasks/batchnorm.py +++ b/torch_judge/tasks/batchnorm.py @@ -4,46 +4,85 @@ "title": "Implement BatchNorm", "difficulty": "Medium", "function_name": "my_batch_norm", - "hint": "Normalize each feature across the batch (dim=0): (x - mean) / sqrt(var + eps) * gamma + beta. Use unbiased=False for variance.", + "hint": ( + "Implement train/eval BatchNorm: in training, use batch stats over dim=0 " + "and update running_mean/running_var with momentum; in inference, normalize " + "using the running statistics only." + ), "tests": [ { - "name": "Basic behavior — zero mean per feature", + "name": "Training mode — zero mean per feature", "code": """ import torch x = torch.randn(8, 4) gamma = torch.ones(4) beta = torch.zeros(4) -out = {fn}(x, gamma, beta) +running_mean = torch.zeros(4) +running_var = torch.ones(4) +out = {fn}(x, gamma, beta, running_mean, running_var, training=True) assert out.shape == x.shape, f'Shape mismatch: {out.shape}' col_means = out.mean(dim=0) assert torch.allclose(col_means, torch.zeros(4), atol=1e-5), f'Column means not zero: {col_means}' """, }, { - "name": "Numerical correctness", + "name": "Training mode — numerical correctness and running stats update", "code": """ import torch +torch.manual_seed(0) x = torch.randn(16, 8) gamma = torch.randn(8) beta = torch.randn(8) -out = {fn}(x, gamma, beta) +running_mean = torch.zeros(8) +running_var = torch.ones(8) +momentum = 0.1 +out = {fn}(x, gamma, beta, running_mean, running_var, momentum=momentum, training=True) + +# Reference using batch stats mean = x.mean(dim=0) var = x.var(dim=0, unbiased=False) ref = gamma * (x - mean) / torch.sqrt(var + 1e-5) + beta assert torch.allclose(out, ref, atol=1e-4), 'Value mismatch' + +# Running stats should have moved toward batch stats +expected_mean = (1 - momentum) * torch.zeros_like(mean) + momentum * mean +expected_var = (1 - momentum) * torch.ones_like(var) + momentum * var +assert torch.allclose(running_mean, expected_mean, atol=1e-6), 'running_mean not updated correctly' +assert torch.allclose(running_var, expected_var, atol=1e-6), 'running_var not updated correctly' +""", + }, + { + "name": "Inference mode — uses running statistics", + "code": """ +import torch +torch.manual_seed(0) +x = torch.randn(4, 8) +gamma = torch.randn(8) +beta = torch.randn(8) + +# Pretend these came from previous training +running_mean = torch.randn(8) +running_var = torch.rand(8) + 0.5 # positive + +out = {fn}(x, gamma, beta, running_mean.clone(), running_var.clone(), training=False) +ref = gamma * (x - running_mean) / torch.sqrt(running_var + 1e-5) + beta +assert torch.allclose(out, ref, atol=1e-4), 'Inference should use running stats' """, }, { - "name": "Gradient flow", + "name": "Gradient flow w.r.t inputs and affine params", "code": """ import torch x = torch.randn(4, 8, requires_grad=True) gamma = torch.ones(8, requires_grad=True) beta = torch.zeros(8, requires_grad=True) -out = {fn}(x, gamma, beta) +running_mean = torch.zeros(8) +running_var = torch.ones(8) +out = {fn}(x, gamma, beta, running_mean, running_var, training=True) out.sum().backward() assert x.grad is not None, 'x.grad is None' assert gamma.grad is not None, 'gamma.grad is None' +assert beta.grad is not None, 'beta.grad is None' """, }, ],