File size: 3,644 Bytes
ecc81b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """The models must actually train, not merely have correct gradients.
"No trainer in the library" is a scope decision. It must not quietly become
"nobody ever checked that it converges" — a block can pass gradcheck and still
never learn. These tests are the guard against that reading.
"""
import warnings
import pytest
import torch
import torch.nn as nn
import torch_dimensions as td
def factory(cls, plan=None):
def build(lat, d_model):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
return cls(d_model, len(lat.axis_names), lat, plan=plan)
return build
@pytest.mark.parametrize("cls", [td.LSTM, td.GRU])
def test_the_rnn_family_learns_a_task_that_needs_axial_mixing(cls):
result = td.testing.check_trainable(factory(cls), steps=150)
assert result["ratio"] >= 3.0, result
assert result["held_out"] < result["initial"]
def test_a_model_that_never_sweeps_the_needed_axis_cannot_learn_it():
"""The negative that makes the positive mean something. The task is a
cumulative sum along `w`; a plan that only sweeps time has no path for
information to travel along `w` at all."""
blind = factory(td.LSTM, plan=td.ScanPlan.from_list(["time", "time"]))
result = td.testing.check_trainable(blind, steps=150, raise_on_failure=False)
assert result["ratio"] < 3.0, f"blind model should not solve this: {result}"
def test_check_trainable_raises_with_an_actionable_message():
blind = factory(td.LSTM, plan=td.ScanPlan.from_list(["time", "time"]))
with pytest.raises(AssertionError, match="did not learn"):
td.testing.check_trainable(blind, steps=100)
def test_check_trainable_reports_the_numbers_it_judged_on():
result = td.testing.check_trainable(factory(td.LSTM), steps=50, min_ratio=0.0)
assert set(result) == {"initial", "final", "held_out", "ratio"}
assert all(isinstance(v, float) for v in result.values())
def test_a_sparse_lattice_model_still_learns():
"""Masking absent cells must not sever the gradient path for present ones."""
valid = torch.tensor(
[[True, True, True, False], [True, True, True, True], [True, False, True, True]]
)
lat = td.Lattice(shape=(3, 4), names=("h", "w"), valid=valid, time=True)
torch.manual_seed(0)
model, head = td.LSTM(16, 3, lat), nn.Linear(16, 1)
opt = torch.optim.Adam([*model.parameters(), *head.parameters()], lr=1e-2)
g = torch.Generator().manual_seed(0)
def draw():
x = torch.randn(8, 5, 3, 4, 16, generator=g)
return x, x[..., :1].cumsum(dim=3) * lat.mask().to(x.dtype)
first = None
for _ in range(150):
x, y = draw()
loss = (head(model(x)) * lat.mask().to(x.dtype) - y).pow(2).mean()
first = first if first is not None else loss.item()
opt.zero_grad()
loss.backward()
opt.step()
assert loss.item() < first / 3, f"{first:.4f} -> {loss.item():.4f}"
def test_training_updates_every_parameter():
"""A parameter that never moves is dead weight the gradient check would
not catch, since it can receive a gradient of exactly zero forever."""
lat = td.Lattice(shape=(2, 3), time=True)
torch.manual_seed(0)
model = td.LSTM(8, 3, lat)
before = {n: p.detach().clone() for n, p in model.named_parameters()}
opt = torch.optim.Adam(model.parameters(), lr=1e-2)
for _ in range(5):
loss = model(torch.randn(4, 3, 2, 3, 8)).pow(2).mean()
opt.zero_grad()
loss.backward()
opt.step()
unmoved = [n for n, p in model.named_parameters() if torch.equal(p, before[n])]
assert not unmoved, unmoved
|