Skip to content

fix: batchnorm add inference code. #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
duoan merged 1 commit into from
Mar 7, 2026
Merged

fix: batchnorm add inference code. #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 50 additions & 7 deletions solutions/07_batchnorm_solution.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."
]
},
{
Expand All @@ -27,24 +28,66 @@
"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"
]
},
{
"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))"
]
},
{
Expand Down
57 changes: 48 additions & 9 deletions templates/07_batchnorm.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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，而不是需要梯度的参数）"
]
},
{
Expand All @@ -44,24 +61,46 @@
"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": [
"# 🧪 Debug\n",
"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)"
]
},
{
Expand Down
53 changes: 46 additions & 7 deletions torch_judge/tasks/batchnorm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
""",
},
],
Expand Down