lars
593c5f4d34
CI / Format (ruff format) (push) Successful in 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 34s
CI / Type check (ty) (push) Successful in 37s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Type check (ty) (pull_request) Successful in 46s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Tests (pull_request) Successful in 4m5s
CI / Tests (push) Successful in 4m6s
The same two-layer classifier head (Linear(cond_out_dim, hidden_dim // 2)
-> SiLU -> Linear(hidden_dim // 2, out_dim)) was hand-rolled five times in
giant/model/models.py: Stage1Model.n_sec_head, Stage2OneShot.n_sec_head/
.type_head, and Stage2Autoregressive.n_sec_head/.type_head. The `// 2`
ratio and fixed 2-layer depth were undocumented magic numbers, and both
n_sec accuracy and secondary-species accuracy are known weak spots that
were untunable independently of the trunk they hang off.
Adds `build_mlp_head(in_dim, out_dim, hidden, depth, act)` to
giant/model/layers.py (depth=1 is a bare Linear; depth>=2 matches the old
hardcoded shape exactly), and a new `HeadConfig` (hidden_ratio, depth)
dataclass in giant/config.py, wired in as `stage1_model.heads.n_sec` and
`stage2_model.heads.{n_sec,type}` — split per head type (not one shared
block per stage) since n_sec and species prediction are called out as
separate weak spots that may want independent capacity. Defaults
(hidden_ratio=0.5, depth=2) reproduce the old hardcoded architecture
bit-for-bit, so every existing config.toml and migrated v0.2 checkpoint
is unaffected; no changes were needed to migrate_config or the legacy
migration surfaces. No new CLI flags, matching how other nested
sub-config (router.*, trunk.*) is set via config.toml rather than
per-field flags.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from giant.model.layers import build_mlp_head
|
|
|
|
|
|
def test_build_mlp_head_depth_1_is_bare_linear():
|
|
head = build_mlp_head(8, 4, hidden=16, depth=1)
|
|
assert len(head) == 1
|
|
assert isinstance(head[0], torch.nn.Linear)
|
|
assert head[0].in_features == 8
|
|
assert head[0].out_features == 4
|
|
out = head(torch.randn(3, 8))
|
|
assert out.shape == (3, 4)
|
|
|
|
|
|
def test_build_mlp_head_depth_2_matches_pre_gitea_36_shape():
|
|
head = build_mlp_head(8, 4, hidden=16, depth=2)
|
|
assert len(head) == 3
|
|
assert isinstance(head[0], torch.nn.Linear)
|
|
assert head[0].in_features == 8
|
|
assert head[0].out_features == 16
|
|
assert isinstance(head[1], torch.nn.SiLU)
|
|
assert isinstance(head[2], torch.nn.Linear)
|
|
assert head[2].in_features == 16
|
|
assert head[2].out_features == 4
|
|
out = head(torch.randn(5, 8))
|
|
assert out.shape == (5, 4)
|
|
|
|
|
|
def test_build_mlp_head_depth_3_has_extra_hidden_layer():
|
|
head = build_mlp_head(8, 4, hidden=16, depth=3)
|
|
assert len(head) == 5
|
|
widths = [(m.in_features, m.out_features) for m in head if isinstance(m, torch.nn.Linear)]
|
|
assert widths == [(8, 16), (16, 16), (16, 4)]
|
|
out = head(torch.randn(2, 8))
|
|
assert out.shape == (2, 4)
|
|
|
|
|
|
def test_build_mlp_head_depth_0_raises():
|
|
with pytest.raises(ValueError, match="depth"):
|
|
build_mlp_head(8, 4, hidden=16, depth=0)
|