mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-07-16 23:24:55 -05:00
13972d6807
train_epoch applies scheduler.get_lr(self.epoch) at the start of the epoch when self.epoch == 0, then increments epoch to 1 at the end. The test was asserting against get_lr(1) which is off by one.
481 lines
20 KiB
Python
481 lines
20 KiB
Python
"""
|
|
Module 08: Training - Coverage Tests
|
|
======================================
|
|
|
|
Tests for the parts of Module 08 that are implemented but have no test coverage:
|
|
- CosineSchedule correctness
|
|
- clip_grad_norm behaviour
|
|
- Trainer.save_checkpoint / load_checkpoint round-trip
|
|
- Trainer.evaluate (loss and accuracy)
|
|
- Scheduler integration inside Trainer.train_epoch
|
|
- Gradient clipping integration inside Trainer.train_epoch
|
|
- Trainer train → eval mode switching
|
|
"""
|
|
|
|
import numpy as np
|
|
import os
|
|
import pickle
|
|
import tempfile
|
|
import pytest
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
|
|
|
from tinytorch.core.autograd import enable_autograd
|
|
|
|
enable_autograd()
|
|
|
|
from tinytorch.core.tensor import Tensor
|
|
from tinytorch.core.layers import Linear
|
|
from tinytorch.core.losses import MSELoss, CrossEntropyLoss
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.training import Trainer, CosineSchedule, clip_grad_norm
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Helpers
|
|
# ─────────────────────────────────────────────
|
|
|
|
def simple_model():
|
|
"""Linear(2→1) model with known initial weights for deterministic tests."""
|
|
layer = Linear(2, 1)
|
|
layer.weight.data = np.array([[0.5], [0.5]])
|
|
layer.bias.data = np.array([0.0])
|
|
return layer
|
|
|
|
|
|
def simple_trainer(lr=0.01, scheduler=None, grad_clip=None):
|
|
model = simple_model()
|
|
opt = SGD(model.parameters(), lr=lr)
|
|
return Trainer(model, opt, MSELoss(), scheduler=scheduler, grad_clip_norm=grad_clip), model
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# CosineSchedule
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestCosineSchedule:
|
|
"""CosineSchedule returns correct learning rates."""
|
|
|
|
def test_start_equals_max_lr(self):
|
|
s = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=100)
|
|
assert abs(s.get_lr(0) - 0.1) < 1e-9
|
|
|
|
def test_end_equals_min_lr(self):
|
|
s = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=100)
|
|
assert abs(s.get_lr(100) - 0.01) < 1e-9
|
|
|
|
def test_midpoint_is_between_min_and_max(self):
|
|
s = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=100)
|
|
mid = s.get_lr(50)
|
|
assert 0.01 < mid < 0.1
|
|
|
|
def test_midpoint_formula(self):
|
|
"""get_lr(50) == (max_lr + min_lr) / 2 for a 100-epoch schedule."""
|
|
s = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=100)
|
|
expected = (0.1 + 0.01) / 2
|
|
assert abs(s.get_lr(50) - expected) < 1e-6
|
|
|
|
def test_monotonically_decreasing(self):
|
|
s = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=100)
|
|
lrs = [s.get_lr(e) for e in range(101)]
|
|
for i in range(len(lrs) - 1):
|
|
assert lrs[i] >= lrs[i + 1], (
|
|
f"LR should be non-increasing: lr[{i}]={lrs[i]:.6f} > lr[{i+1}]={lrs[i+1]:.6f}"
|
|
)
|
|
|
|
def test_past_total_epochs_returns_min_lr(self):
|
|
s = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=50)
|
|
assert abs(s.get_lr(999) - 0.01) < 1e-9
|
|
|
|
def test_single_epoch_schedule(self):
|
|
"""Edge case: total_epochs=1."""
|
|
s = CosineSchedule(max_lr=0.5, min_lr=0.05, total_epochs=1)
|
|
assert abs(s.get_lr(0) - 0.5) < 1e-9
|
|
assert abs(s.get_lr(1) - 0.05) < 1e-9
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# clip_grad_norm
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestClipGradNorm:
|
|
"""clip_grad_norm clips gradient magnitudes and returns the original norm."""
|
|
|
|
def _params_with_grads(self, grad_values):
|
|
"""Create Tensor params with preset gradients."""
|
|
params = []
|
|
for v in grad_values:
|
|
p = Tensor(np.zeros_like(v), requires_grad=True)
|
|
p.grad = np.array(v, dtype=np.float64)
|
|
params.append(p)
|
|
return params
|
|
|
|
def test_returns_original_norm(self):
|
|
params = self._params_with_grads([[3.0, 4.0]]) # norm = 5
|
|
original_norm = clip_grad_norm(params, max_norm=10.0)
|
|
assert abs(original_norm - 5.0) < 1e-6
|
|
|
|
def test_clips_large_gradients(self):
|
|
params = self._params_with_grads([[3.0, 4.0]]) # norm = 5
|
|
clip_grad_norm(params, max_norm=1.0)
|
|
clipped_norm = np.linalg.norm(params[0].grad)
|
|
assert abs(clipped_norm - 1.0) < 1e-6
|
|
|
|
def test_does_not_clip_small_gradients(self):
|
|
params = self._params_with_grads([[0.1, 0.1]]) # norm ≈ 0.14
|
|
original_grad = params[0].grad.copy()
|
|
clip_grad_norm(params, max_norm=1.0)
|
|
np.testing.assert_allclose(params[0].grad, original_grad)
|
|
|
|
def test_clips_across_multiple_params(self):
|
|
"""Global norm is computed over all params together."""
|
|
params = self._params_with_grads([[3.0, 4.0], [0.0, 0.0]])
|
|
# global norm = 5; max_norm = 1 → scale = 0.2
|
|
clip_grad_norm(params, max_norm=1.0)
|
|
expected = np.array([3.0, 4.0]) * (1.0 / 5.0)
|
|
np.testing.assert_allclose(params[0].grad, expected, rtol=1e-5)
|
|
|
|
def test_direction_preserved_after_clipping(self):
|
|
"""Clipping scales magnitude but preserves gradient direction."""
|
|
params = self._params_with_grads([[3.0, 4.0]])
|
|
original_dir = params[0].grad / np.linalg.norm(params[0].grad)
|
|
clip_grad_norm(params, max_norm=1.0)
|
|
clipped_dir = params[0].grad / np.linalg.norm(params[0].grad)
|
|
np.testing.assert_allclose(clipped_dir, original_dir, atol=1e-6)
|
|
|
|
def test_zero_gradients_no_division_by_zero(self):
|
|
"""All-zero gradients should not cause division by zero."""
|
|
params = self._params_with_grads([[0.0, 0.0, 0.0]])
|
|
norm = clip_grad_norm(params, max_norm=1.0)
|
|
assert np.isfinite(norm)
|
|
np.testing.assert_allclose(params[0].grad, np.zeros(3))
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Checkpoint round-trip
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestCheckpointing:
|
|
"""save_checkpoint / load_checkpoint preserve all training state."""
|
|
|
|
def test_checkpoint_file_is_created(self):
|
|
trainer, _ = simple_trainer()
|
|
with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
|
|
path = f.name
|
|
try:
|
|
trainer.save_checkpoint(path)
|
|
assert os.path.exists(path)
|
|
finally:
|
|
os.remove(path)
|
|
|
|
def test_checkpoint_contains_required_keys(self):
|
|
trainer, _ = simple_trainer()
|
|
with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
|
|
path = f.name
|
|
try:
|
|
trainer.save_checkpoint(path)
|
|
with open(path, "rb") as f:
|
|
ckpt = pickle.load(f)
|
|
for key in ("epoch", "step", "model_state", "optimizer_state", "history"):
|
|
assert key in ckpt, f"Missing key: {key}"
|
|
finally:
|
|
os.remove(path)
|
|
|
|
def test_epoch_and_step_restored(self):
|
|
trainer, _ = simple_trainer()
|
|
trainer.epoch = 42
|
|
trainer.step = 1337
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
|
|
path = f.name
|
|
try:
|
|
trainer.save_checkpoint(path)
|
|
trainer.epoch = 0
|
|
trainer.step = 0
|
|
trainer.load_checkpoint(path)
|
|
assert trainer.epoch == 42
|
|
assert trainer.step == 1337
|
|
finally:
|
|
os.remove(path)
|
|
|
|
def test_history_restored(self):
|
|
trainer, _ = simple_trainer()
|
|
trainer.history["train_loss"] = [0.9, 0.7, 0.5]
|
|
trainer.history["eval_loss"] = [0.8, 0.6]
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
|
|
path = f.name
|
|
try:
|
|
trainer.save_checkpoint(path)
|
|
trainer.history = {"train_loss": [], "eval_loss": [], "learning_rates": []}
|
|
trainer.load_checkpoint(path)
|
|
assert trainer.history["train_loss"] == [0.9, 0.7, 0.5]
|
|
assert trainer.history["eval_loss"] == [0.8, 0.6]
|
|
finally:
|
|
os.remove(path)
|
|
|
|
def test_model_weights_restored(self):
|
|
trainer, model = simple_trainer()
|
|
original_weights = model.parameters()[0].data.copy()
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
|
|
path = f.name
|
|
try:
|
|
trainer.save_checkpoint(path)
|
|
# Corrupt the weights
|
|
model.parameters()[0].data[:] = 999.0
|
|
trainer.load_checkpoint(path)
|
|
np.testing.assert_allclose(model.parameters()[0].data, original_weights)
|
|
finally:
|
|
os.remove(path)
|
|
|
|
def test_training_continues_after_load(self):
|
|
"""Model can keep training after loading a checkpoint."""
|
|
trainer, model = simple_trainer(lr=0.01)
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
|
|
trainer.train_epoch(data)
|
|
weights_after_first_epoch = model.parameters()[0].data.copy()
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
|
|
path = f.name
|
|
try:
|
|
trainer.save_checkpoint(path)
|
|
trainer.load_checkpoint(path)
|
|
trainer.train_epoch(data)
|
|
weights_after_second_epoch = model.parameters()[0].data.copy()
|
|
# Weights should change after resuming
|
|
assert not np.allclose(weights_after_first_epoch, weights_after_second_epoch)
|
|
finally:
|
|
os.remove(path)
|
|
|
|
def test_checkpoint_creates_parent_directory(self):
|
|
"""save_checkpoint creates intermediate directories if needed."""
|
|
trainer, _ = simple_trainer()
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path = os.path.join(tmpdir, "subdir", "deep", "ckpt.pkl")
|
|
trainer.save_checkpoint(path)
|
|
assert os.path.exists(path)
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Trainer.evaluate
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestTrainerEvaluate:
|
|
"""Trainer.evaluate computes correct metrics without modifying the model."""
|
|
|
|
def test_returns_finite_loss(self):
|
|
trainer, _ = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
loss, _ = trainer.evaluate(data)
|
|
assert np.isfinite(loss), f"Expected finite loss, got {loss}"
|
|
|
|
def test_returns_float(self):
|
|
trainer, _ = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
loss, acc = trainer.evaluate(data)
|
|
assert isinstance(loss, (float, np.floating))
|
|
assert isinstance(acc, (float, np.floating))
|
|
|
|
def test_model_set_to_eval_mode(self):
|
|
trainer, model = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.evaluate(data)
|
|
assert model.training is False
|
|
assert trainer.training_mode is False
|
|
|
|
def test_weights_unchanged_after_evaluate(self):
|
|
trainer, model = simple_trainer()
|
|
weights_before = model.parameters()[0].data.copy()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))] * 5
|
|
trainer.evaluate(data)
|
|
np.testing.assert_array_equal(model.parameters()[0].data, weights_before)
|
|
|
|
def test_eval_loss_recorded_in_history(self):
|
|
trainer, _ = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.evaluate(data)
|
|
assert len(trainer.history["eval_loss"]) == 1
|
|
|
|
def test_eval_loss_recorded_each_call(self):
|
|
trainer, _ = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.evaluate(data)
|
|
trainer.evaluate(data)
|
|
assert len(trainer.history["eval_loss"]) == 2
|
|
|
|
def test_classification_accuracy(self):
|
|
"""Accuracy is 1.0 when argmax predictions match integer targets."""
|
|
class PerfectClassifier:
|
|
training = True
|
|
|
|
def forward(self, x):
|
|
# Always predicts class 0 with very high confidence
|
|
batch = x.data.shape[0]
|
|
logits = np.zeros((batch, 3))
|
|
logits[:, 0] = 10.0
|
|
return Tensor(logits)
|
|
|
|
def parameters(self):
|
|
return []
|
|
|
|
loss_fn = CrossEntropyLoss()
|
|
opt = SGD([], lr=0.01)
|
|
trainer = Trainer(PerfectClassifier(), opt, loss_fn)
|
|
|
|
data = [(Tensor([[1.0, 0.0]]), Tensor(np.array([0])))]
|
|
_, accuracy = trainer.evaluate(data)
|
|
assert accuracy == 1.0, f"Perfect classifier should have accuracy=1.0, got {accuracy}"
|
|
|
|
def test_zero_accuracy_for_wrong_predictions(self):
|
|
"""Accuracy is 0.0 when predictions are always wrong."""
|
|
class WrongClassifier:
|
|
training = True
|
|
|
|
def forward(self, x):
|
|
batch = x.data.shape[0]
|
|
logits = np.zeros((batch, 3))
|
|
logits[:, 1] = 10.0 # always predicts class 1
|
|
return Tensor(logits)
|
|
|
|
def parameters(self):
|
|
return []
|
|
|
|
loss_fn = CrossEntropyLoss()
|
|
opt = SGD([], lr=0.01)
|
|
trainer = Trainer(WrongClassifier(), opt, loss_fn)
|
|
|
|
data = [(Tensor([[1.0, 0.0]]), Tensor(np.array([0])))] # target is class 0
|
|
_, accuracy = trainer.evaluate(data)
|
|
assert accuracy == 0.0, f"Wrong classifier should have accuracy=0.0, got {accuracy}"
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Scheduler integration
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestSchedulerIntegration:
|
|
"""CosineSchedule is applied correctly during train_epoch."""
|
|
|
|
def test_lr_recorded_in_history_when_scheduler_present(self):
|
|
scheduler = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=10)
|
|
trainer, _ = simple_trainer(scheduler=scheduler)
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.train_epoch(data)
|
|
assert len(trainer.history["learning_rates"]) == 1
|
|
|
|
def test_no_lr_in_history_without_scheduler(self):
|
|
trainer, _ = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.train_epoch(data)
|
|
assert len(trainer.history["learning_rates"]) == 0
|
|
|
|
def test_optimizer_lr_updated_by_scheduler(self):
|
|
"""After train_epoch, optimizer lr should match the scheduler value applied that epoch."""
|
|
scheduler = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=10)
|
|
trainer, _ = simple_trainer(lr=0.1, scheduler=scheduler)
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
# train_epoch applies scheduler.get_lr(self.epoch) at the start of the epoch,
|
|
# when self.epoch == 0, then increments epoch to 1 at the end.
|
|
expected_lr = scheduler.get_lr(0)
|
|
trainer.train_epoch(data)
|
|
assert abs(trainer.optimizer.lr - expected_lr) < 1e-9
|
|
|
|
def test_lr_decreases_over_epochs(self):
|
|
scheduler = CosineSchedule(max_lr=0.1, min_lr=0.01, total_epochs=20)
|
|
trainer, _ = simple_trainer(lr=0.1, scheduler=scheduler)
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
|
|
for _ in range(5):
|
|
trainer.train_epoch(data)
|
|
|
|
lrs = trainer.history["learning_rates"]
|
|
for i in range(len(lrs) - 1):
|
|
assert lrs[i] >= lrs[i + 1], (
|
|
f"LR should decrease: lrs[{i}]={lrs[i]:.6f} > lrs[{i+1}]={lrs[i+1]:.6f}"
|
|
)
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Gradient clipping integration
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestGradientClippingIntegration:
|
|
"""Gradient clipping actually limits gradient norms during training."""
|
|
|
|
def test_training_completes_with_grad_clip(self):
|
|
"""Training should not crash when grad_clip_norm is set."""
|
|
trainer, _ = simple_trainer(lr=0.01, grad_clip=1.0)
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
loss = trainer.train_epoch(data)
|
|
assert np.isfinite(loss), f"Loss should be finite with grad clipping, got {loss}"
|
|
|
|
def test_weights_update_with_grad_clip(self):
|
|
"""Weights still change when clipping is active."""
|
|
trainer, model = simple_trainer(lr=0.1, grad_clip=0.01)
|
|
weights_before = model.parameters()[0].data.copy()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.train_epoch(data)
|
|
assert not np.allclose(model.parameters()[0].data, weights_before)
|
|
|
|
def test_very_tight_clip_limits_updates(self):
|
|
"""Extremely small max_norm keeps weight updates very small."""
|
|
trainer_clipped, model_clipped = simple_trainer(lr=0.1, grad_clip=1e-6)
|
|
trainer_free, model_free = simple_trainer(lr=0.1)
|
|
|
|
# Same initial weights
|
|
w0 = np.array([[0.5], [0.5]])
|
|
model_clipped.parameters()[0].data[:] = w0
|
|
model_clipped.parameters()[1].data[:] = 0.0
|
|
model_free.parameters()[0].data[:] = w0
|
|
model_free.parameters()[1].data[:] = 0.0
|
|
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer_clipped.train_epoch(data)
|
|
trainer_free.train_epoch(data)
|
|
|
|
update_clipped = np.abs(model_clipped.parameters()[0].data - w0).max()
|
|
update_free = np.abs(model_free.parameters()[0].data - w0).max()
|
|
assert update_clipped < update_free, (
|
|
"Tightly clipped update should be smaller than unclipped update"
|
|
)
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Train / eval mode switching
|
|
# ─────────────────────────────────────────────
|
|
|
|
class TestTrainEvalMode:
|
|
"""Trainer correctly switches model between train and eval mode."""
|
|
|
|
def test_model_in_train_mode_during_train_epoch(self):
|
|
"""model.training should be True at the end of train_epoch."""
|
|
trainer, model = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.train_epoch(data)
|
|
assert model.training is True
|
|
assert trainer.training_mode is True
|
|
|
|
def test_model_in_eval_mode_during_evaluate(self):
|
|
trainer, model = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.evaluate(data)
|
|
assert model.training is False
|
|
assert trainer.training_mode is False
|
|
|
|
def test_train_after_eval_restores_train_mode(self):
|
|
"""Calling train_epoch after evaluate re-enables training mode."""
|
|
trainer, model = simple_trainer()
|
|
data = [(Tensor([[1.0, 0.5]]), Tensor([[2.0]]))]
|
|
trainer.evaluate(data)
|
|
assert model.training is False
|
|
trainer.train_epoch(data)
|
|
assert model.training is True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|