File size: 2,435 Bytes
32f5a65 | 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 | from __future__ import annotations
import torch
from sepsis_mcp.grud_model import GRUDModel
def test_grud_forward_returns_batch_probabilities() -> None:
model = GRUDModel(
input_size=2,
static_size=5,
hidden_size=8,
)
values = torch.tensor(
[
[[80.0, 95.0], [82.0, 0.0], [84.0, 93.0]],
[[70.0, 97.0], [71.0, 96.0], [72.0, 95.0]],
],
dtype=torch.float32,
)
masks = torch.tensor(
[
[[1.0, 1.0], [1.0, 0.0], [1.0, 1.0]],
[[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]],
],
dtype=torch.float32,
)
deltas = torch.tensor(
[
[[0.0, 0.0], [0.0, 1.0], [0.0, 0.0]],
[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],
],
dtype=torch.float32,
)
static = torch.ones((2, 5), dtype=torch.float32)
probabilities = model(values, masks, deltas, static)
assert probabilities.shape == (2,)
assert torch.all(probabilities >= 0.0)
assert torch.all(probabilities <= 1.0)
def test_grud_forward_logits_matches_probability_output() -> None:
model = GRUDModel(
input_size=2,
static_size=5,
hidden_size=8,
)
values = torch.rand((2, 3, 2), dtype=torch.float32)
masks = torch.ones((2, 3, 2), dtype=torch.float32)
deltas = torch.zeros((2, 3, 2), dtype=torch.float32)
static = torch.rand((2, 5), dtype=torch.float32)
logits = model.forward_logits(values, masks, deltas, static)
probabilities = model(values, masks, deltas, static)
assert logits.shape == (2,)
assert torch.allclose(torch.sigmoid(logits), probabilities)
def test_grud_supports_single_optimizer_step() -> None:
torch.manual_seed(0)
model = GRUDModel(
input_size=2,
static_size=5,
hidden_size=8,
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.BCELoss()
values = torch.rand((4, 3, 2), dtype=torch.float32)
masks = torch.randint(0, 2, (4, 3, 2), dtype=torch.float32)
deltas = torch.rand((4, 3, 2), dtype=torch.float32)
static = torch.rand((4, 5), dtype=torch.float32)
labels = torch.tensor([0.0, 1.0, 0.0, 1.0], dtype=torch.float32)
optimizer.zero_grad()
probabilities = model(values, masks, deltas, static)
loss = criterion(probabilities, labels)
loss.backward()
optimizer.step()
assert torch.isfinite(loss)
|