Huang Qi
c8fa5db7d3
Coverage: - test_model: SimpleNet forward (parametrized over batch sizes and both unsqueezed and flat input shapes), layer dimensions, differentiability, and ONNX round-trip - test_inference: load_model resolution order (bundled, cwd override, explicit path, missing path), and predict shape/dtype/determinism plus endpoint sanity across 8 cardinal/diagonal targets - test_train: _load_csv parsing, TrajectoryDataset indexing, full train() pipeline producing a single-file ONNX, plus a smoke test against the real data shipped under data/ - test_cli: --help for the three console scripts and a real run of mouse-visualize via both the entry point and python -m Wire up pytest via dependency-groups and tool.pytest.ini_options.
75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from mouse_control import SimpleNet
|
|
|
|
|
|
@pytest.fixture
|
|
def model() -> SimpleNet:
|
|
return SimpleNet()
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
|
|
def test_forward_with_unsqueezed_input(model: SimpleNet, batch_size: int) -> None:
|
|
out = model(torch.randn(batch_size, 1, 2))
|
|
assert out.shape == (batch_size, 10, 2)
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [1, 8, 32])
|
|
def test_forward_with_flat_input(model: SimpleNet, batch_size: int) -> None:
|
|
"""SimpleNet's first layer flattens, so (N, 2) is also valid."""
|
|
out = model(torch.randn(batch_size, 2))
|
|
assert out.shape == (batch_size, 10, 2)
|
|
|
|
|
|
def test_layers_have_expected_dimensions(model: SimpleNet) -> None:
|
|
assert model.fc1.in_features == 2
|
|
assert model.fc1.out_features == 64
|
|
assert model.fc2.in_features == 64
|
|
assert model.fc2.out_features == 32
|
|
assert model.fc3.in_features == 32
|
|
assert model.fc3.out_features == 20
|
|
|
|
|
|
def test_model_has_trainable_parameters(model: SimpleNet) -> None:
|
|
params = list(model.parameters())
|
|
assert params, "SimpleNet should expose parameters"
|
|
assert all(p.requires_grad for p in params)
|
|
|
|
|
|
def test_forward_is_differentiable(model: SimpleNet) -> None:
|
|
inp = torch.randn(2, 1, 2, requires_grad=True)
|
|
out = model(inp)
|
|
out.sum().backward()
|
|
assert inp.grad is not None
|
|
assert inp.grad.shape == inp.shape
|
|
|
|
|
|
def test_onnx_export_round_trips(tmp_path, model: SimpleNet) -> None:
|
|
"""Trained or not, the model architecture must round-trip via ONNX."""
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
onnx_path = tmp_path / "round_trip.onnx"
|
|
dummy = torch.randn(1, 1, 2)
|
|
model.eval()
|
|
torch.onnx.export(
|
|
model, dummy, str(onnx_path),
|
|
input_names=["input"], output_names=["output"],
|
|
external_data=False,
|
|
)
|
|
assert onnx_path.exists()
|
|
assert not onnx_path.with_suffix(onnx_path.suffix + ".data").exists(), \
|
|
"external_data=False must produce a single-file ONNX"
|
|
|
|
session = ort.InferenceSession(str(onnx_path))
|
|
inp = np.array([[[42.0, -17.0]]], dtype=np.float32)
|
|
out = session.run(None, {"input": inp})[0]
|
|
assert out.shape == (1, 10, 2)
|
|
|
|
with torch.no_grad():
|
|
torch_out = model(torch.from_numpy(inp)).numpy()
|
|
np.testing.assert_allclose(out, torch_out, atol=1e-4)
|