28 lines
953 B
Python
28 lines
953 B
Python
import torch
|
|
|
|
from models.laneformer import LaneFormer
|
|
|
|
|
|
def test_forward_shapes():
|
|
torch.manual_seed(0)
|
|
model = LaneFormer(pretrained=False, max_lanes=8)
|
|
x = torch.randn(2, 3, 360, 640)
|
|
out = model(x)
|
|
assert out["cls_logits"].shape == (2, 8, 2)
|
|
assert out["curve_coeffs"].shape == (2, 8, 4)
|
|
assert out["endpoints"].shape == (2, 8, 2)
|
|
assert torch.isfinite(out["cls_logits"]).all()
|
|
assert torch.isfinite(out["curve_coeffs"]).all()
|
|
assert (out["endpoints"] >= 0).all() and (out["endpoints"] <= 1).all()
|
|
|
|
|
|
def test_backward_runs():
|
|
torch.manual_seed(0)
|
|
model = LaneFormer(pretrained=False, max_lanes=8)
|
|
x = torch.randn(1, 3, 360, 640)
|
|
out = model(x)
|
|
loss = out["cls_logits"].sum() + out["curve_coeffs"].sum() + out["endpoints"].sum()
|
|
loss.backward()
|
|
grads = [p.grad for p in model.parameters() if p.requires_grad]
|
|
assert any(g is not None and torch.isfinite(g).all() for g in grads)
|