File size: 6,724 Bytes
e6f80f3 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | import pytest
import torch
from fastfill.v2.batch import TinyTokenizer, collate_samples
from fastfill.v2.model import ModelConfig, build_model, load_model, model_inputs
from fastfill.v2.tests.test_batch import sample
def config():
return ModelConfig(backbone="tiny", decoder_dim=32, decoder_heads=4,
decoder_layers=2, tiny_hidden_size=32, max_objects=16,
tiny_layers=1, dropout=0.0)
@pytest.mark.parametrize("kwargs", [
{"decoder_dim": True}, {"decoder_layers": 1.5}, {"decoder_ffn_multiplier": 0},
{"dropout": float("nan")}, {"lora_dropout": float("inf")}, {"lora_rank": -1},
{"train_backbone": "false", "lora_rank": 0}, {"local_files_only": "false"},
{"size_reference": (True, 1., 1.)}, {"size_reference": (1e38, 1., 1.)},
{"size_reference": (1e-42, 1., 1.)}, {"backbone": ""}, {"lora_target_modules": [None]},
])
def test_model_config_rejects_ambiguous_types_and_nonfinite_geometry(kwargs):
with pytest.raises(ValueError):
ModelConfig(**kwargs)
def test_positive_size_bounds_survive_fp16_autocast():
cfg = ModelConfig(backbone="tiny", decoder_dim=32, decoder_heads=4, decoder_layers=1,
tiny_hidden_size=32, size_log_limit=30.)
model = build_model(cfg)
with torch.no_grad():
model.size_head.weight.zero_()
model.size_head.bias.copy_(torch.tensor([30., -30., 0.]))
batch = collate_samples([sample(1)], TinyTokenizer())
with torch.autocast("cpu", dtype=torch.float16):
sizes = model(**model_inputs(batch))["size"]
assert sizes.dtype == torch.float32
assert torch.isfinite(sizes).all() and (sizes > 0).all()
sizes.log().sum().backward()
assert torch.isfinite(model.size_head.bias.grad).all()
def test_exact_slots_positive_size_and_gradients_to_condition_backbone():
model = build_model(config())
b = collate_samples([sample(2), sample(1)], TinyTokenizer())
out = model(**model_inputs(b))
assert out["position_normalized"].shape == (2, 2, 3)
assert out["yaw_logits"].shape == (2, 2, 12)
assert out["slot_mask"].tolist() == [[True, True], [True, False]]
assert (out["size"][out["slot_mask"]] > 0).all()
assert torch.equal(out["position_normalized"][1, 1], torch.zeros(3))
loss = sum(out[k].square().sum() for k in ("position_normalized", "size", "yaw_logits", "yaw_residuals"))
loss.backward()
assert model.backbone.embedding.weight.grad.abs().sum() > 0
assert model.position_head.weight.grad.abs().sum() > 0
def test_batch_padding_does_not_change_valid_object_predictions():
torch.manual_seed(2)
model = build_model(config()).eval()
one = collate_samples([sample(1)], TinyTokenizer())
padded = collate_samples([sample(1), sample(3)], TinyTokenizer())
with torch.no_grad():
a = model(**model_inputs(one))
b = model(**model_inputs(padded))
for key in ("position_normalized", "size", "yaw_logits", "yaw_residuals"):
assert torch.allclose(a[key][0, 0], b[key][0, 0], atol=2e-6), key
def test_full_condition_after_object_changes_predictions():
model = build_model(config()).eval()
a = sample(1)
b = sample(1)
b["condition"]["constraints"] = [{"type": "faces_direction", "object_id": "chair_0", "direction_xy": [0, -1]}]
with torch.no_grad():
x = model(**model_inputs(collate_samples([a], TinyTokenizer())))
y = model(**model_inputs(collate_samples([b], TinyTokenizer())))
assert not torch.allclose(x["position_normalized"], y["position_normalized"])
def test_fixed_size_and_ground_height_override_head():
s = sample(1)
s["condition"]["objects"][0].update(fixed_size_local_m=[1.1, .7, 1.2], support_parent="floor")
s["target"]["objects"][0]["target_size_local_m"] = [1.1, .7, 1.2]
model = build_model(config())
b = collate_samples([s], TinyTokenizer())
out = model(**model_inputs(b))
assert torch.equal(out["size"][0, 0], torch.tensor([1.1, .7, 1.2]))
assert out["position_normalized"][0, 0, 2] == 0
def test_checkpoint_round_trip(tmp_path):
model = build_model(config()).eval()
b = collate_samples([sample(1)], TinyTokenizer())
model.save_pretrained(tmp_path)
loaded = load_model(tmp_path).eval()
with torch.no_grad():
assert torch.equal(model(**model_inputs(b))["size"], loaded(**model_inputs(b))["size"])
def test_empty_scene_and_mixed_empty_batch_do_not_create_nan():
model = build_model(config())
empty = collate_samples([sample(0)], TinyTokenizer())
out = model(**model_inputs(empty))
assert out["size"].shape == (1, 0, 3)
mixed = collate_samples([sample(0), sample(2)], TinyTokenizer())
out = model(**model_inputs(mixed))
assert torch.isfinite(out["size"]).all()
assert torch.equal(out["size"][0], torch.zeros(2, 3))
def test_bidirectional_object_interaction_changes_first_slot():
torch.manual_seed(7)
model = build_model(config()).eval()
a = sample(2)
b = sample(2)
b["condition"]["objects"][1]["description"] = "large reading chair with an attached lamp"
with torch.no_grad():
x = model(**model_inputs(collate_samples([a], TinyTokenizer())))
y = model(**model_inputs(collate_samples([b], TinyTokenizer())))
assert not torch.allclose(x["size"][0, 0], y["size"][0, 0])
@pytest.mark.parametrize("training", ["lora", "full", "frozen"])
def test_local_random_qwen_api_gradients_and_checkpoint(tmp_path, training):
transformers = pytest.importorskip("transformers")
if training == "lora":
pytest.importorskip("peft")
source = tmp_path / "qwen"
transformers.Qwen2Model(transformers.Qwen2Config(vocab_size=258, hidden_size=32,
intermediate_size=64, num_hidden_layers=1, num_attention_heads=4,
num_key_value_heads=2, max_position_embeddings=4096)).save_pretrained(source)
cfg = ModelConfig(backbone=str(source), decoder_dim=32, decoder_heads=4, decoder_layers=1,
lora_rank=2 if training == "lora" else 0, lora_alpha=4,
train_backbone=training == "full", local_files_only=True)
model = build_model(cfg)
b = collate_samples([sample(1)], TinyTokenizer())
model(**model_inputs(b))["size"].sum().backward()
gradients = [p.grad for p in model.backbone.parameters() if p.grad is not None]
assert bool(gradients) == (training != "frozen")
if gradients:
assert sum(float(grad.abs().sum()) for grad in gradients) > 0
checkpoint = tmp_path / "checkpoint"
model.eval().save_pretrained(checkpoint)
loaded = load_model(checkpoint, local_files_only=True).eval()
with torch.no_grad():
assert torch.equal(model(**model_inputs(b))["size"], loaded(**model_inputs(b))["size"])
|