From f3558905b051719c374d15d1cc3a55d7174e19bf Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 19:03:27 +0200 Subject: [PATCH 01/64] docs: add MemPalace setup instructions for Kimi Code --- .agents/mempalace-setup.md | 159 +++++++++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 .agents/mempalace-setup.md diff --git a/.agents/mempalace-setup.md b/.agents/mempalace-setup.md new file mode 100644 index 0000000..8628a08 --- /dev/null +++ b/.agents/mempalace-setup.md @@ -0,0 +1,159 @@ +# MemPalace Setup for Kimi Code + +MemPalace is configured for this project. To recreate the setup on a new machine, follow the steps below. + +## 1. Install MemPalace + +```bash +uv tool install mempalace +``` + +`pipx install mempalace` works too. Plain `pip` is fine inside an activated virtualenv. + +## 2. Register the MemPalace MCP server with Kimi Code + +Add to `~/.kimi-code/mcp.json`: + +```json +{ + "mcpServers": { + "mempalace": { + "command": "mempalace-mcp", + "args": [], + "env": {} + } + } +} +``` + +Use a custom palace path if needed: + +```json +{ + "mcpServers": { + "mempalace": { + "command": "mempalace-mcp", + "args": ["--palace", "/path/to/palace"], + "env": {} + } + } +} +``` + +## 3. Initialize and mine the project + +From the repository root: + +```bash +mempalace init . --yes +mempalace mine . +``` + +The project already includes a `mempalace.yaml` that maps the `jaeger` wing to the repository folders. + +## 4. Verify the integration + +Test the MCP server directly: + +```bash +echo '{"jsonrpc":"2.0","method":"tools/list","id":1}' | mempalace-mcp | python -m json.tool +``` + +In a fresh Kimi Code session: + +```bash +kimi -p "List the tools from the mempalace MCP server." +``` + +You can then ask Kimi to search the palace, e.g.: + +```bash +kimi -p "Search the palace for training speed comparisons." +``` + +## 5. Optional — pre-approve safe MemPalace tools + +Add to `~/.kimi-code/config.toml`: + +```toml +[[permission.rules]] +decision = "allow" +pattern = "mcp__mempalace__status" + +[[permission.rules]] +decision = "allow" +pattern = "mcp__mempalace__list_*" + +[[permission.rules]] +decision = "allow" +pattern = "mcp__mempalace__search" + +[[permission.rules]] +decision = "allow" +pattern = "mcp__mempalace__get_*" + +[[permission.rules]] +decision = "deny" +pattern = "mcp__mempalace__delete_*" + +[[permission.rules]] +decision = "deny" +pattern = "mcp__mempalace__kg_invalidate" +``` + +This allows read-only/search operations and keeps destructive operations behind manual approval. + +## 6. Autosave hooks + +MemPalace's autosave hooks are adapted to Kimi Code via a small shim because Kimi Code's hook JSON and transcript layout differ from Claude Code/Codex. + +### 6.1 Install the shim + +Create `~/.kimi-code/hooks/mempalace-autosave.py` with the contents from the project plan or from the canonical version maintained in this repository (if committed). The shim: + +- Receives Kimi Code `Stop` and `PreCompact` hook events. +- Looks up the session directory in `~/.kimi-code/session_index.jsonl`. +- Converts `agents/main/wire.jsonl` to a Claude-Code-style transcript. +- Writes the transcript under `~/.mempalace/kimi_transcripts//` so MemPalace only mines that directory. +- Invokes `mempalace hook run --hook stop|precompact --harness claude-code`. +- Does **not** re-mine the project automatically (to avoid lock contention); run `mempalace mine .` manually after significant code changes. + +Make it executable: + +```bash +chmod +x ~/.kimi-code/hooks/mempalace-autosave.py +``` + +### 6.2 Register the hooks + +Append to `~/.kimi-code/config.toml`: + +```toml +[[hooks]] +event = "Stop" +command = "/home/yasas-wijesekara/.kimi-code/hooks/mempalace-autosave.py" +timeout = 60 + +[[hooks]] +event = "PreCompact" +command = "/home/yasas-wijesekara/.kimi-code/hooks/mempalace-autosave.py" +timeout = 60 +``` + +Adjust the absolute path to match your home directory. + +### 6.3 Verify hooks + +After a fresh Kimi Code session reaches ~15 user messages, check: + +```bash +tail -20 ~/.mempalace/hook_state/hook.log +``` + +You should see `TRIGGERING SAVE at exchange 15` and a diary checkpoint entry. Trigger `/compact` and check the log again for a `PRE-COMPACT triggered` line. + +## References + +- MemPalace: +- MemPalace hooks guide: +- Kimi Code hooks docs: From b4e2c5669a4296f4e2a4b02880b1a04b5013120b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 19:15:23 +0200 Subject: [PATCH 02/64] docs: add AGENTS.md memory protocol and update MemPalace setup instructions --- .agents/mempalace-setup.md | 19 +++++++++++++++---- AGENTS.md | 30 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) create mode 100644 AGENTS.md diff --git a/.agents/mempalace-setup.md b/.agents/mempalace-setup.md index 8628a08..be3d11e 100644 --- a/.agents/mempalace-setup.md +++ b/.agents/mempalace-setup.md @@ -103,11 +103,22 @@ pattern = "mcp__mempalace__kg_invalidate" This allows read-only/search operations and keeps destructive operations behind manual approval. -## 6. Autosave hooks +## 6. Always save plans and decisions + +Create `AGENTS.md` in the project root with explicit MemPalace memory instructions. Kimi Code loads it as a project-level system prompt, so every session is reminded to: + +- Search the palace before answering questions about past plans or decisions. +- Save every new or updated plan to the `jaeger/plans` room via `mcp__mempalace__add_drawer`. +- Record decisions as knowledge-graph facts via `mcp__mempalace__kg_add`. +- Write a session diary entry at the end of each session via `mcp__mempalace__diary_write`. + +See the current `AGENTS.md` in this repo for the exact wording. + +## 7. Autosave hooks MemPalace's autosave hooks are adapted to Kimi Code via a small shim because Kimi Code's hook JSON and transcript layout differ from Claude Code/Codex. -### 6.1 Install the shim +### 7.1 Install the shim Create `~/.kimi-code/hooks/mempalace-autosave.py` with the contents from the project plan or from the canonical version maintained in this repository (if committed). The shim: @@ -124,7 +135,7 @@ Make it executable: chmod +x ~/.kimi-code/hooks/mempalace-autosave.py ``` -### 6.2 Register the hooks +### 7.2 Register the hooks Append to `~/.kimi-code/config.toml`: @@ -142,7 +153,7 @@ timeout = 60 Adjust the absolute path to match your home directory. -### 6.3 Verify hooks +### 7.3 Verify hooks After a fresh Kimi Code session reaches ~15 user messages, check: diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..b6bc234 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,30 @@ +# Jaeger — Agent Instructions + +## Project background + +Jaeger (`jaeger-bio`) is a homology-free deep-learning command-line tool for identifying bacteriophage genome sequences. The repository lives at `/home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/Jaeger` and uses Python, TensorFlow, NumPy, and YAML training configs. + +## Memory protocol — MemPalace + +This project uses [MemPalace](https://mempalaceofficial.com/) for durable, searchable memory across Kimi Code sessions. You MUST follow this protocol on every turn: + +1. **On wake-up / session start**: call `mcp__mempalace__mempalace_status` to load the palace overview and orient yourself. +2. **Before answering** any question about past plans, decisions, people, or project facts: call `mcp__mempalace__mempalace_search` or `mcp__mempalace__mempalace_kg_query`. Do not guess. +3. **After creating or updating a plan**: save it immediately with `mcp__mempalace__mempalace_add_drawer`: + - `wing`: `jaeger` + - `room`: `plans` + - `topic`: `plan` + - `content`: the full plan text, including goals, tasks, exact file paths, and decisions. +4. **After making a decision**: record it with `mcp__mempalace__mempalace_kg_add`: + - Example: `subject="Jaeger project"`, `predicate="decided_to"`, `object="adopt AGENTS.md memory protocol"`, `valid_from="YYYY-MM-DD"`. +5. **After each session**: write a diary entry with `mcp__mempalace__mempalace_diary_write`: + - `agent_name`: `kimi-code` + - `wing`: `jaeger` + - `topic`: `session-summary` + - `entry`: a concise, factual summary of what was discussed, decided, implemented, and left unfinished. +6. **When facts change**: invalidate the old fact with `mcp__mempalace__mempalace_kg_invalidate` and add the corrected fact with `mcp__mempalace__mempalace_kg_add`. +7. **Do not bulk-delete drawers or knowledge-graph facts** without explicit user approval. + +## Tool permission note + +Read/search MemPalace tools (`mempalace_status`, `mempalace_search`, `mempalace_get_*`, `mempalace_kg_query`) are pre-approved in `~/.kimi-code/config.toml`. Write tools (`add_drawer`, `kg_add`, `diary_write`) will trigger approval prompts unless the user enables YOLO mode. From 325d0354930b21ea740966860416c4dee4cc9dc7 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 22:54:07 +0200 Subject: [PATCH 03/64] chore(pytorch): replace TensorFlow optional deps with PyTorch, keep legacy extra --- pyproject.toml | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index faa6fb7..2721963 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,26 +61,28 @@ jaeger = "jaeger.cli:main" Homepage = "https://github.com/Yasas1994/Jaeger" [project.optional-dependencies] -darwin-arm = [ - "tensorflow >=2.21, <2.22", - "tensorflow-metal", -] gpu = [ - "tensorflow[and-cuda] >=2.21, <2.22", + "torch >=2.0, <3.0", ] cpu = [ + "torch >=2.0, <3.0", +] +darwin-arm = [ + "torch >=2.0, <3.0", +] +legacy = [ "tensorflow >=2.21, <2.22", ] +test = [ + "pytest >=8.0", + "pytest-mock >=3.14", +] onnx = [ "onnxruntime-gpu >=1.26", "tf2onnx >=1.17", "sympy >=1.13", "onnx >=1.16", ] -test = [ - "pytest >=8.0", - "pytest-mock >=3.14", -] taxonomy = [ "taxopy", "faiss-cpu", From fa7de0fe5b656cf416fd0509da28dc1c26e02635 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:01:21 +0200 Subject: [PATCH 04/64] feat(pytorch): add GeLU activation layer --- src/jaeger/nnlib/pytorch/__init__.py | 1 + src/jaeger/nnlib/pytorch/layers.py | 16 ++++++++++++++++ tests/unit/nnlib/pytorch/test_layers.py | 11 +++++++++++ 3 files changed, 28 insertions(+) create mode 100644 src/jaeger/nnlib/pytorch/__init__.py create mode 100644 src/jaeger/nnlib/pytorch/layers.py create mode 100644 tests/unit/nnlib/pytorch/test_layers.py diff --git a/src/jaeger/nnlib/pytorch/__init__.py b/src/jaeger/nnlib/pytorch/__init__.py new file mode 100644 index 0000000..7440390 --- /dev/null +++ b/src/jaeger/nnlib/pytorch/__init__.py @@ -0,0 +1 @@ +"""PyTorch neural network layers and utilities for Jaeger.""" diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py new file mode 100644 index 0000000..6178d5a --- /dev/null +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -0,0 +1,16 @@ +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class GeLU(nn.Module): + """Tanh-approximated GELU for TFLite-compatible graph export.""" + + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="tanh") diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py new file mode 100644 index 0000000..f1376de --- /dev/null +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -0,0 +1,11 @@ +import torch +import pytest +from jaeger.nnlib.pytorch.layers import GeLU + + +def test_gelu_matches_torch_nn_gelu(): + x = torch.tensor([-1.0, 0.0, 1.0, 2.0]) + layer = GeLU() + out = layer(x) + expected = torch.nn.functional.gelu(x, approximate="tanh") + assert torch.allclose(out, expected, atol=1e-6) From e5fd7a52e7dfa9207c9a0122887c4900bff680c1 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:04:04 +0200 Subject: [PATCH 05/64] feat(pytorch): add MaskedConv1D with mask propagation --- src/jaeger/nnlib/pytorch/layers.py | 101 ++++++++++++++++++++++++ tests/unit/nnlib/pytorch/test_layers.py | 14 +++- 2 files changed, 114 insertions(+), 1 deletion(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 6178d5a..74b37c0 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -14,3 +14,104 @@ def __init__(self): def forward(self, x: torch.Tensor) -> torch.Tensor: return F.gelu(x, approximate="tanh") + + +class MaskedConv1D(nn.Module): + """1D convolution over the sequence axis of (B, F, L, C) inputs. + + Masked positions are zeroed before convolution and the output mask is + propagated based on whether the full kernel window contained valid values. + """ + + def __init__( + self, + filters: int, + kernel_size: int, + strides: int = 1, + padding: str = "valid", + dilation_rate: int = 1, + activation: Optional[str] = None, + use_bias: bool = True, + kernel_initializer: str = "glorot_uniform", + bias_initializer: str = "zeros", + kernel_regularizer: Optional[float] = None, + ): + super().__init__() + self.filters = filters + self.kernel_size = kernel_size + self.strides = strides + self.padding = padding.lower() + self.dilation_rate = dilation_rate + self.use_bias = use_bias + self.activation = activation + + self.conv = nn.Conv1d( + in_channels=filters, # placeholder; set in build + out_channels=filters, + kernel_size=kernel_size, + stride=strides, + padding=0, # handled manually + dilation=dilation_rate, + bias=use_bias, + ) + # Actual in_channels will be determined at first forward; override then. + + def _resolve_padding(self, length: int) -> int: + if self.padding == "same": + dilated = self.dilation_rate * (self.kernel_size - 1) + return (length + self.strides - 1) // self.strides * self.strides - length + dilated + return 0 + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + b, f, l, c = x.shape + # Lazy set conv in_channels on first call + if self.conv.in_channels != c: + self.conv = nn.Conv1d( + in_channels=c, + out_channels=self.filters, + kernel_size=self.kernel_size, + stride=self.strides, + padding=0, + dilation=self.dilation_rate, + bias=self.use_bias, + ).to(x.device) + + if mask is not None: + x = x * mask.unsqueeze(-1).to(x.dtype) + + # Merge batch and frame dims: (B*F, C, L) + x_2d = x.reshape(b * f, c, l) + if self.padding == "same": + pad_total = self._resolve_padding(l) + pad_left = pad_total // 2 + pad_right = pad_total - pad_left + x_2d = F.pad(x_2d, (pad_left, pad_right)) + + out = self.conv(x_2d) + _, _, l_out = out.shape + out = out.reshape(b, f, l_out, self.filters) + + if self.activation: + out = getattr(F, self.activation)(out) + + out_mask = None + if mask is not None: + mask_f = mask.reshape(b * f, 1, l).to(x.dtype) + if self.padding == "same": + mask_f = F.pad(mask_f, (pad_left, pad_right)) + with torch.no_grad(): + kernel = torch.ones( + (1, 1, self.kernel_size), dtype=mask_f.dtype, device=mask_f.device + ) + out_mask = F.conv1d( + mask_f, + kernel, + stride=self.strides, + dilation=self.dilation_rate, + ) + out_mask = (out_mask >= self.kernel_size).squeeze(1) + out_mask = out_mask.reshape(b, f, l_out) + + return out, out_mask diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index f1376de..5aa8b9b 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -1,6 +1,6 @@ import torch import pytest -from jaeger.nnlib.pytorch.layers import GeLU +from jaeger.nnlib.pytorch.layers import GeLU, MaskedConv1D def test_gelu_matches_torch_nn_gelu(): @@ -9,3 +9,15 @@ def test_gelu_matches_torch_nn_gelu(): out = layer(x) expected = torch.nn.functional.gelu(x, approximate="tanh") assert torch.allclose(out, expected, atol=1e-6) + + +def test_masked_conv1d_shape_and_mask(): + # (B=2, F=6, L=10, C=4) + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[0, :, 5:] = False + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") + out, out_mask = layer(x, mask) + assert out.shape == (2, 6, 10, 8) + assert out_mask.shape == (2, 6, 10) + assert not out_mask[0, :, 5:].any() From 4b91aa5d2f32128cfec927f0b38c2a05eb1fdfaf Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:09:44 +0200 Subject: [PATCH 06/64] feat(pytorch): improve MaskedConv1D init and tests --- src/jaeger/nnlib/pytorch/layers.py | 88 +++++++++++++++++++++---- tests/unit/nnlib/pytorch/test_layers.py | 65 ++++++++++++++++++ 2 files changed, 142 insertions(+), 11 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 74b37c0..9acd32c 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -1,4 +1,3 @@ -import math from typing import Optional, Tuple import torch @@ -21,8 +20,22 @@ class MaskedConv1D(nn.Module): Masked positions are zeroed before convolution and the output mask is propagated based on whether the full kernel window contained valid values. + Bias is added at all positions, so downstream code should apply the + returned mask if masked positions must remain suppressed. + + For API compatibility with Keras-style configs, ``kernel_initializer``, + ``bias_initializer`` and ``kernel_regularizer`` are accepted in the + signature but only their default values are supported; passing a non-default + value raises ``NotImplementedError``. """ + _SUPPORTED_ACTIVATIONS = { + "relu": F.relu, + "gelu": F.gelu, + "sigmoid": F.sigmoid, + "tanh": F.tanh, + } + def __init__( self, filters: int, @@ -37,16 +50,51 @@ def __init__( kernel_regularizer: Optional[float] = None, ): super().__init__() + padding_norm = padding.lower() + if padding_norm not in ("same", "valid"): + raise ValueError( + f"padding must be 'same' or 'valid', got {padding!r}" + ) + + if kernel_initializer != "glorot_uniform": + raise NotImplementedError( + f"kernel_initializer {kernel_initializer!r} is not supported; " + "only 'glorot_uniform' is implemented." + ) + if bias_initializer != "zeros": + raise NotImplementedError( + f"bias_initializer {bias_initializer!r} is not supported; " + "only 'zeros' is implemented." + ) + if kernel_regularizer is not None: + raise NotImplementedError( + f"kernel_regularizer {kernel_regularizer!r} is not supported." + ) + self.filters = filters self.kernel_size = kernel_size self.strides = strides - self.padding = padding.lower() + self.padding = padding_norm self.dilation_rate = dilation_rate self.use_bias = use_bias - self.activation = activation + if activation is None: + self.activation: Optional[callable] = None + else: + act_norm = activation.lower() + if act_norm not in self._SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of " + f"{list(self._SUPPORTED_ACTIVATIONS.keys())} or None, " + f"got {activation!r}" + ) + self.activation = self._SUPPORTED_ACTIVATIONS[act_norm] + + # Placeholder conv so that ``parameters()`` and ``state_dict()`` are + # available before the first forward. ``in_channels`` is set to the + # real value on the first forward call. self.conv = nn.Conv1d( - in_channels=filters, # placeholder; set in build + in_channels=1, out_channels=filters, kernel_size=kernel_size, stride=strides, @@ -54,20 +102,30 @@ def __init__( dilation=dilation_rate, bias=use_bias, ) - # Actual in_channels will be determined at first forward; override then. + self._in_channels: Optional[int] = None def _resolve_padding(self, length: int) -> int: if self.padding == "same": dilated = self.dilation_rate * (self.kernel_size - 1) - return (length + self.strides - 1) // self.strides * self.strides - length + dilated + return ( + (length + self.strides - 1) // self.strides * self.strides + - length + + dilated + ) return 0 def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if x.ndim != 4: + raise ValueError( + f"Expected 4D input of shape (B, F, L, C), got {x.ndim}D" + ) + b, f, l, c = x.shape - # Lazy set conv in_channels on first call - if self.conv.in_channels != c: + + if self._in_channels is None: + # First forward: create the real conv with the correct in_channels. self.conv = nn.Conv1d( in_channels=c, out_channels=self.filters, @@ -77,6 +135,12 @@ def forward( dilation=self.dilation_rate, bias=self.use_bias, ).to(x.device) + self._in_channels = c + elif c != self._in_channels: + raise ValueError( + f"Expected input channels {self._in_channels}, got {c}. " + "MaskedConv1D does not support variable channel dimensions." + ) if mask is not None: x = x * mask.unsqueeze(-1).to(x.dtype) @@ -93,8 +157,8 @@ def forward( _, _, l_out = out.shape out = out.reshape(b, f, l_out, self.filters) - if self.activation: - out = getattr(F, self.activation)(out) + if self.activation is not None: + out = self.activation(out) out_mask = None if mask is not None: @@ -103,7 +167,9 @@ def forward( mask_f = F.pad(mask_f, (pad_left, pad_right)) with torch.no_grad(): kernel = torch.ones( - (1, 1, self.kernel_size), dtype=mask_f.dtype, device=mask_f.device + (1, 1, self.kernel_size), + dtype=mask_f.dtype, + device=mask_f.device, ) out_mask = F.conv1d( mask_f, diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 5aa8b9b..25dcfe2 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -11,6 +11,13 @@ def test_gelu_matches_torch_nn_gelu(): assert torch.allclose(out, expected, atol=1e-6) +def test_gelu_preserves_shape(): + x = torch.randn(2, 4, 8) + layer = GeLU() + out = layer(x) + assert out.shape == x.shape + + def test_masked_conv1d_shape_and_mask(): # (B=2, F=6, L=10, C=4) x = torch.randn(2, 6, 10, 4) @@ -21,3 +28,61 @@ def test_masked_conv1d_shape_and_mask(): assert out.shape == (2, 6, 10, 8) assert out_mask.shape == (2, 6, 10) assert not out_mask[0, :, 5:].any() + + +def test_masked_conv1d_valid_padding(): + x = torch.randn(2, 3, 10, 4) + layer = MaskedConv1D(filters=8, kernel_size=3, padding="valid") + out, out_mask = layer(x, mask=None) + assert out.shape == (2, 3, 8, 8) + + +def test_masked_conv1d_stride_and_dilation(): + x = torch.randn(2, 3, 10, 4) + layer = MaskedConv1D( + filters=8, kernel_size=3, padding="same", strides=2, dilation_rate=2 + ) + out, out_mask = layer(x, mask=None) + assert out.shape == (2, 3, 5, 8) + assert out_mask is None + + +def test_masked_conv1d_no_mask(): + x = torch.randn(2, 3, 10, 4) + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") + out, out_mask = layer(x, mask=None) + assert out.shape == (2, 3, 10, 8) + assert out_mask is None + + +def test_masked_conv1d_activation(): + x = torch.abs(torch.randn(2, 3, 10, 4)) + 1.0 + layer = MaskedConv1D( + filters=8, kernel_size=3, padding="same", activation="relu" + ) + out, _ = layer(x, mask=None) + assert (out >= 0).all() + + +def test_masked_conv1d_interior_mask_hole(): + x = torch.randn(1, 1, 10, 4) + mask = torch.ones(1, 1, 10, dtype=torch.bool) + mask[0, 0, 5] = False + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") + _, out_mask = layer(x, mask) + assert not out_mask[0, 0, 5] + assert out_mask[0, 0, 2] + assert out_mask[0, 0, 7] + + +def test_masked_conv1d_parameters_before_forward(): + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") + params = list(layer.parameters()) + assert len(params) > 0 + + +def test_masked_conv1d_state_dict_before_forward(): + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") + state = layer.state_dict() + assert "conv.weight" in state + assert "conv.bias" in state From 33970f63bcdb60b9537659d3d4d9d8e1f7f52393 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:12:49 +0200 Subject: [PATCH 07/64] feat(pytorch): add MaskedBatchNorm and MaskedLayerNorm --- src/jaeger/nnlib/pytorch/layers.py | 95 +++++++++++++++++++++++++ tests/unit/nnlib/pytorch/test_layers.py | 21 +++++- 2 files changed, 115 insertions(+), 1 deletion(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 9acd32c..ad5179b 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -181,3 +181,98 @@ def forward( out_mask = out_mask.reshape(b, f, l_out) return out, out_mask + + +class MaskedBatchNorm(nn.Module): + """Batch normalization that excludes masked positions from statistics. + + Can optionally return normalized mean difference (nmd) vectors. + """ + + def __init__( + self, + num_features: int, + eps: float = 1e-5, + momentum: float = 0.9, + return_nmd: bool = False, + ): + super().__init__() + self.num_features = num_features + self.eps = eps + self.momentum = momentum + self.return_nmd = return_nmd + + self.gamma = nn.Parameter(torch.ones(num_features)) + self.beta = nn.Parameter(torch.zeros(num_features)) + self.register_buffer("running_mean", torch.zeros(num_features)) + self.register_buffer("running_var", torch.ones(num_features)) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + xf = x.to(torch.float32) + b, f, l, c = xf.shape + + if mask is not None: + mask_f = mask.unsqueeze(-1).to(torch.float32) + masked_x = xf * mask_f + valid_count = mask_f.sum(dim=(0, 1, 2)) + self.eps + mean = masked_x.sum(dim=(0, 1, 2)) / valid_count + var = ((masked_x - mean) * mask_f).pow(2).sum(dim=(0, 1, 2)) / valid_count + else: + mean = xf.mean(dim=(0, 1, 2)) + var = xf.var(dim=(0, 1, 2), unbiased=False) + + if self.training: + self.running_mean = self.momentum * self.running_mean + (1 - self.momentum) * mean.detach() + self.running_var = self.momentum * self.running_var + (1 - self.momentum) * var.detach() + mean_use, var_use = mean, var + else: + mean_use, var_use = self.running_mean, self.running_var + + mean_use = mean_use.view(1, 1, 1, -1) + var_use = var_use.view(1, 1, 1, -1) + normalized = (xf - mean_use) / torch.sqrt(var_use + self.eps) + out = normalized * self.gamma.view(1, 1, 1, -1) + self.beta.view(1, 1, 1, -1) + out = out.to(x.dtype) + + if self.return_nmd: + if mask is not None: + per_ex_sum = masked_x.sum(dim=(1, 2)) + per_ex_count = mask_f.sum(dim=(1, 2)) + self.eps + mean_channel = per_ex_sum / per_ex_count + else: + mean_channel = xf.mean(dim=(1, 2)) + nmd = (mean_channel - mean).to(x.dtype) + return out, nmd + + return out, None + + +class MaskedLayerNorm(nn.Module): + """Layer normalization that excludes masked positions.""" + + def __init__(self, num_features: int, eps: float = 1e-3): + super().__init__() + self.num_features = num_features + self.eps = eps + self.gamma = nn.Parameter(torch.ones(num_features)) + self.beta = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + xf = x.to(torch.float32) + if mask is not None: + mask_f = mask.unsqueeze(-1).to(torch.float32) + masked_x = xf * mask_f + count = mask_f.sum(dim=-1, keepdim=True) + self.eps + mean = masked_x.sum(dim=-1, keepdim=True) / count + var = ((masked_x - mean) * mask_f).pow(2).sum(dim=-1, keepdim=True) / count + else: + mean = xf.mean(dim=-1, keepdim=True) + var = xf.var(dim=-1, keepdim=True, unbiased=False) + + normalized = (xf - mean) / torch.sqrt(var + self.eps) + out = normalized * self.gamma.view(1, 1, 1, -1) + self.beta.view(1, 1, 1, -1) + if mask is not None: + out = out * mask_f + return out.to(x.dtype) diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 25dcfe2..5976135 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -1,6 +1,6 @@ import torch import pytest -from jaeger.nnlib.pytorch.layers import GeLU, MaskedConv1D +from jaeger.nnlib.pytorch.layers import GeLU, MaskedBatchNorm, MaskedConv1D, MaskedLayerNorm def test_gelu_matches_torch_nn_gelu(): @@ -86,3 +86,22 @@ def test_masked_conv1d_state_dict_before_forward(): state = layer.state_dict() assert "conv.weight" in state assert "conv.bias" in state + + +def test_masked_batchnorm_output_shape_and_nmd(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[0, :, 5:] = False + layer = MaskedBatchNorm(num_features=4, return_nmd=True) + out, nmd = layer(x, mask) + assert out.shape == (2, 6, 10, 4) + assert nmd.shape == (2, 4) + + +def test_masked_layer_norm_shape(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[0, :, 5:] = False + layer = MaskedLayerNorm(num_features=4) + out = layer(x, mask) + assert out.shape == (2, 6, 10, 4) From b81081f300eadcc31c09a1a0ad8dfe7c53b7f6f4 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:19:03 +0200 Subject: [PATCH 08/64] feat(pytorch): fix MaskedLayerNorm stats and lint issues --- src/jaeger/nnlib/pytorch/layers.py | 42 +++++++++++-------- tests/unit/nnlib/pytorch/test_layers.py | 54 ++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 17 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index ad5179b..f4d75e8 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple +from typing import Callable, Optional, Tuple import torch import torch.nn as nn @@ -32,8 +32,8 @@ class MaskedConv1D(nn.Module): _SUPPORTED_ACTIVATIONS = { "relu": F.relu, "gelu": F.gelu, - "sigmoid": F.sigmoid, - "tanh": F.tanh, + "sigmoid": torch.sigmoid, + "tanh": torch.tanh, } def __init__( @@ -79,7 +79,7 @@ def __init__( self.use_bias = use_bias if activation is None: - self.activation: Optional[callable] = None + self.activation: Optional[Callable[[torch.Tensor], torch.Tensor]] = None else: act_norm = activation.lower() if act_norm not in self._SUPPORTED_ACTIVATIONS: @@ -122,7 +122,7 @@ def forward( f"Expected 4D input of shape (B, F, L, C), got {x.ndim}D" ) - b, f, l, c = x.shape + b, f, length, c = x.shape if self._in_channels is None: # First forward: create the real conv with the correct in_channels. @@ -146,23 +146,23 @@ def forward( x = x * mask.unsqueeze(-1).to(x.dtype) # Merge batch and frame dims: (B*F, C, L) - x_2d = x.reshape(b * f, c, l) + x_2d = x.reshape(b * f, c, length) if self.padding == "same": - pad_total = self._resolve_padding(l) + pad_total = self._resolve_padding(length) pad_left = pad_total // 2 pad_right = pad_total - pad_left x_2d = F.pad(x_2d, (pad_left, pad_right)) out = self.conv(x_2d) - _, _, l_out = out.shape - out = out.reshape(b, f, l_out, self.filters) + _, _, length_out = out.shape + out = out.reshape(b, f, length_out, self.filters) if self.activation is not None: out = self.activation(out) out_mask = None if mask is not None: - mask_f = mask.reshape(b * f, 1, l).to(x.dtype) + mask_f = mask.reshape(b * f, 1, length).to(x.dtype) if self.padding == "same": mask_f = F.pad(mask_f, (pad_left, pad_right)) with torch.no_grad(): @@ -178,7 +178,7 @@ def forward( dilation=self.dilation_rate, ) out_mask = (out_mask >= self.kernel_size).squeeze(1) - out_mask = out_mask.reshape(b, f, l_out) + out_mask = out_mask.reshape(b, f, length_out) return out, out_mask @@ -187,6 +187,9 @@ class MaskedBatchNorm(nn.Module): """Batch normalization that excludes masked positions from statistics. Can optionally return normalized mean difference (nmd) vectors. + + The running statistics are updated with Keras-style momentum: + ``running = momentum * running + (1 - momentum) * batch``. """ def __init__( @@ -211,7 +214,7 @@ def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: xf = x.to(torch.float32) - b, f, l, c = xf.shape + b, f, length, c = xf.shape if mask is not None: mask_f = mask.unsqueeze(-1).to(torch.float32) @@ -224,8 +227,8 @@ def forward( var = xf.var(dim=(0, 1, 2), unbiased=False) if self.training: - self.running_mean = self.momentum * self.running_mean + (1 - self.momentum) * mean.detach() - self.running_var = self.momentum * self.running_var + (1 - self.momentum) * var.detach() + self.running_mean.lerp_(mean.detach(), 1 - self.momentum) + self.running_var.lerp_(var.detach(), 1 - self.momentum) mean_use, var_use = mean, var else: mean_use, var_use = self.running_mean, self.running_var @@ -250,7 +253,13 @@ def forward( class MaskedLayerNorm(nn.Module): - """Layer normalization that excludes masked positions.""" + """Layer normalization that excludes masked positions. + + Normalization is performed over the channel dimension for each + (batch, frame, position) tuple. Positions excluded by ``mask`` are + omitted from the mean and variance, and the corresponding output is + zeroed. + """ def __init__(self, num_features: int, eps: float = 1e-3): super().__init__() @@ -264,7 +273,8 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch if mask is not None: mask_f = mask.unsqueeze(-1).to(torch.float32) masked_x = xf * mask_f - count = mask_f.sum(dim=-1, keepdim=True) + self.eps + count = mask_f.sum(dim=-1, keepdim=True) * x.shape[-1] + count = count.masked_fill(count == 0, self.eps) mean = masked_x.sum(dim=-1, keepdim=True) / count var = ((masked_x - mean) * mask_f).pow(2).sum(dim=-1, keepdim=True) / count else: diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 5976135..305a8b8 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -1,5 +1,4 @@ import torch -import pytest from jaeger.nnlib.pytorch.layers import GeLU, MaskedBatchNorm, MaskedConv1D, MaskedLayerNorm @@ -98,6 +97,34 @@ def test_masked_batchnorm_output_shape_and_nmd(): assert nmd.shape == (2, 4) +def test_masked_batchnorm_no_mask(): + x = torch.randn(2, 6, 10, 4) + layer = MaskedBatchNorm(num_features=4) + out, _ = layer(x, mask=None) + assert out.shape == (2, 6, 10, 4) + # running stats should have updated in train mode + assert not torch.allclose(layer.running_mean, torch.zeros(4)) + assert not torch.allclose(layer.running_var, torch.ones(4)) + + +def test_masked_batchnorm_eval_uses_running_stats(): + x = torch.randn(2, 6, 10, 4) + layer = MaskedBatchNorm(num_features=4) + layer.train() + _ = layer(x) + running_mean = layer.running_mean.clone() + running_var = layer.running_var.clone() + layer.eval() + out, _ = layer(x, mask=None) + assert out.shape == (2, 6, 10, 4) + # manual normalization using running stats should match + expected = (x - running_mean.view(1, 1, 1, -1)) / torch.sqrt( + running_var.view(1, 1, 1, -1) + layer.eps + ) + expected = expected * layer.gamma.view(1, 1, 1, -1) + layer.beta.view(1, 1, 1, -1) + assert torch.allclose(out, expected, atol=1e-5) + + def test_masked_layer_norm_shape(): x = torch.randn(2, 6, 10, 4) mask = torch.ones(2, 6, 10, dtype=torch.bool) @@ -105,3 +132,28 @@ def test_masked_layer_norm_shape(): layer = MaskedLayerNorm(num_features=4) out = layer(x, mask) assert out.shape == (2, 6, 10, 4) + # masked positions should be zeroed + assert torch.allclose(out[0, :, 5:, :], torch.zeros_like(out[0, :, 5:, :])) + + +def test_masked_layer_norm_no_mask(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + layer = MaskedLayerNorm(num_features=4) + out_masked = layer(x, mask) + out_unmasked = layer(x, mask=None) + assert torch.allclose(out_masked, out_unmasked, atol=1e-5) + + +def test_masked_layer_norm_numerical(): + # Simple 1 position, 4 channels, all valid. + x = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]]]]) + layer = MaskedLayerNorm(num_features=4) + with torch.no_grad(): + layer.gamma.fill_(1.0) + layer.beta.fill_(0.0) + out = layer(x) + mean = x.mean(dim=-1) + var = x.var(dim=-1, unbiased=False) + expected = (x - mean) / torch.sqrt(var + layer.eps) + assert torch.allclose(out, expected, atol=1e-5) From c4d0fa82fcbdb2c9e41f30f8a320653f78691dd1 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:23:25 +0200 Subject: [PATCH 09/64] feat(pytorch): add attention and pooling layers --- src/jaeger/nnlib/pytorch/layers.py | 63 +++++++++++++++++++++++++ tests/unit/nnlib/pytorch/test_layers.py | 23 ++++++++- 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index f4d75e8..7710235 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -286,3 +286,66 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch if mask is not None: out = out * mask_f return out.to(x.dtype) + + +class GatedFrameGlobalMaxPooling(nn.Module): + """Frame-aware global max pooling. Input (B,F,L,D) -> output (B,D).""" + + def __init__(self, return_gate: bool = False): + super().__init__() + self.return_gate = return_gate + # Placeholder; real in_features is set on first forward. + self.score_dense = nn.Linear(1, 1) + self._in_features: Optional[int] = None + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]: + # x: (B, F, L, D) + b, f, length, d = x.shape + if self._in_features is None: + self.score_dense = nn.Linear(d, 1).to(x.device) + self._in_features = d + elif d != self._in_features: + raise ValueError( + f"Expected input channels {self._in_features}, got {d}. " + "GatedFrameGlobalMaxPooling does not support variable channel dimensions." + ) + + per_frame = x.max(dim=2)[0] # (B, F, D) + logits = self.score_dense(per_frame).squeeze(-1) # (B, F) + gates = torch.softmax(logits, dim=1) + pooled = (per_frame * gates.unsqueeze(-1)).sum(dim=1) # (B, D) + if self.return_gate: + return pooled, gates + return pooled + + +class AxialAttention(nn.Module): + """Axial attention over the sequence axis.""" + + def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0): + super().__init__() + self.attn = nn.MultiheadAttention( + embed_dim=embed_dim, + num_heads=num_heads, + dropout=dropout, + batch_first=True, + ) + self.norm = nn.LayerNorm(embed_dim) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + b, f, l, d = x.shape + # Reshape to (B*F, L, D) + x_2d = x.reshape(b * f, l, d) + key_padding_mask = None + if mask is not None: + key_padding_mask = ~mask.reshape(b * f, l) + attn_out, _ = self.attn( + x_2d, x_2d, x_2d, key_padding_mask=key_padding_mask, need_weights=False + ) + out = self.norm(x_2d + attn_out).reshape(b, f, l, d) + out_mask = mask + return out, out_mask diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 305a8b8..2fb880f 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -1,5 +1,12 @@ import torch -from jaeger.nnlib.pytorch.layers import GeLU, MaskedBatchNorm, MaskedConv1D, MaskedLayerNorm +from jaeger.nnlib.pytorch.layers import ( + AxialAttention, + GeLU, + GatedFrameGlobalMaxPooling, + MaskedBatchNorm, + MaskedConv1D, + MaskedLayerNorm, +) def test_gelu_matches_torch_nn_gelu(): @@ -157,3 +164,17 @@ def test_masked_layer_norm_numerical(): var = x.var(dim=-1, unbiased=False) expected = (x - mean) / torch.sqrt(var + layer.eps) assert torch.allclose(out, expected, atol=1e-5) + + +def test_gated_frame_pooling_shape(): + x = torch.randn(2, 6, 10, 4) + layer = GatedFrameGlobalMaxPooling(return_gate=False) + out = layer(x) + assert out.shape == (2, 4) + + +def test_axial_attention_shape(): + x = torch.randn(2, 6, 10, 4) + layer = AxialAttention(embed_dim=4, num_heads=2) + out, mask = layer(x) + assert out.shape == (2, 6, 10, 4) From 81e176e0bb1f4484ae6715091f19c07fec6c04ac Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:26:44 +0200 Subject: [PATCH 10/64] feat(pytorch): refine pooling/attention layers and tests --- src/jaeger/nnlib/pytorch/layers.py | 26 ++++++------------- tests/unit/nnlib/pytorch/test_layers.py | 34 +++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 18 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 7710235..a963749 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -294,27 +294,17 @@ class GatedFrameGlobalMaxPooling(nn.Module): def __init__(self, return_gate: bool = False): super().__init__() self.return_gate = return_gate - # Placeholder; real in_features is set on first forward. - self.score_dense = nn.Linear(1, 1) - self._in_features: Optional[int] = None + self.score_dense = nn.LazyLinear(1) def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]: # x: (B, F, L, D) - b, f, length, d = x.shape - if self._in_features is None: - self.score_dense = nn.Linear(d, 1).to(x.device) - self._in_features = d - elif d != self._in_features: - raise ValueError( - f"Expected input channels {self._in_features}, got {d}. " - "GatedFrameGlobalMaxPooling does not support variable channel dimensions." - ) - + if mask is not None: + x = x.masked_fill(~mask.unsqueeze(-1), -1e9) per_frame = x.max(dim=2)[0] # (B, F, D) logits = self.score_dense(per_frame).squeeze(-1) # (B, F) - gates = torch.softmax(logits, dim=1) + gates = F.softmax(logits, dim=1) pooled = (per_frame * gates.unsqueeze(-1)).sum(dim=1) # (B, D) if self.return_gate: return pooled, gates @@ -337,15 +327,15 @@ def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0): def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - b, f, l, d = x.shape + b, f, seq_len, d = x.shape # Reshape to (B*F, L, D) - x_2d = x.reshape(b * f, l, d) + x_2d = x.reshape(b * f, seq_len, d) key_padding_mask = None if mask is not None: - key_padding_mask = ~mask.reshape(b * f, l) + key_padding_mask = ~mask.reshape(b * f, seq_len) attn_out, _ = self.attn( x_2d, x_2d, x_2d, key_padding_mask=key_padding_mask, need_weights=False ) - out = self.norm(x_2d + attn_out).reshape(b, f, l, d) + out = self.norm(x_2d + attn_out).reshape(b, f, seq_len, d) out_mask = mask return out, out_mask diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 2fb880f..850312a 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -178,3 +178,37 @@ def test_axial_attention_shape(): layer = AxialAttention(embed_dim=4, num_heads=2) out, mask = layer(x) assert out.shape == (2, 6, 10, 4) + + +def test_gated_frame_pooling_return_gate(): + x = torch.randn(2, 6, 10, 4) + layer = GatedFrameGlobalMaxPooling(return_gate=True) + pooled, gates = layer(x) + assert pooled.shape == (2, 4) + assert gates.shape == (2, 6) + assert torch.allclose(gates.sum(dim=1), torch.ones(2), atol=1e-5) + + +def test_gated_frame_pooling_with_mask(): + x = torch.zeros(1, 2, 4, 4) + x[0, 0, 0, :] = 10.0 # unmasked high value in frame 0 + x[0, 1, 0, :] = 10.0 # high value in frame 1, will be masked out + mask = torch.ones(1, 2, 4, dtype=torch.bool) + mask[0, 1, 0] = False # mask out the high position in frame 1 + layer = GatedFrameGlobalMaxPooling(return_gate=False) + out = layer(x, mask) + # Masked position should not affect per-frame max; frame 1 max becomes 0.0, + # so the gated pooled output must be strictly less than 10.0. + assert out.shape == (1, 4) + assert (out < 10.0).all() + + +def test_axial_attention_with_mask(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[:, :, 5:] = False + layer = AxialAttention(embed_dim=4, num_heads=2) + out, out_mask = layer(x, mask) + assert out.shape == (2, 6, 10, 4) + assert torch.isfinite(out).all() + assert out_mask is mask From cb55bcb39c26dd864339b38d4d1901b87ad5c860 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:31:33 +0200 Subject: [PATCH 11/64] feat(pytorch): add remaining custom layers and RepresentationModel --- src/jaeger/nnlib/pytorch/layers.py | 116 ++++++++++++++++++++++++ src/jaeger/nnlib/pytorch/models.py | 74 +++++++++++++++ tests/unit/nnlib/pytorch/test_layers.py | 61 +++++++++++++ 3 files changed, 251 insertions(+) create mode 100644 src/jaeger/nnlib/pytorch/models.py diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index a963749..03d8e08 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -1,3 +1,4 @@ +import math from typing import Callable, Optional, Tuple import torch @@ -339,3 +340,118 @@ def forward( out = self.norm(x_2d + attn_out).reshape(b, f, seq_len, d) out_mask = mask return out, out_mask + + +class MaskedGlobalAvgPooling(nn.Module): + """Global average pooling over the sequence axis, respecting masks.""" + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + # x: (B, F, L, D) + if mask is not None: + mask_f = mask.unsqueeze(-1).to(x.dtype) + masked_x = x * mask_f + count = mask_f.sum(dim=2, keepdim=True).clamp(min=1.0) + pooled = masked_x.sum(dim=2, keepdim=True) / count + return pooled.mean(dim=(1, 2)) + return x.mean(dim=(1, 2)) + + +class SinusoidalPositionEmbedding(nn.Module): + def __init__(self, max_wavelength: int = 10000): + super().__init__() + self.max_wavelength = max_wavelength + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, F, L, D) + b, f, length, d = x.shape + position = torch.arange(length, device=x.device).unsqueeze(1).float() + div_term = torch.exp( + torch.arange(0, d, 2, device=x.device).float() + * (-math.log(self.max_wavelength) / d) + ) + pe = torch.zeros(length, d, device=x.device) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + return x + pe.view(1, 1, length, d) + + +class ResidualBlock(nn.Module): + def __init__(self, layer: nn.Module): + super().__init__() + self.layer = layer + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + if hasattr(self.layer, "forward") and "mask" in self.layer.forward.__code__.co_varnames: + out = self.layer(x, mask) + if isinstance(out, tuple): + out, out_mask = out + return out + x, out_mask + return out + x, mask + out = self.layer(x) + return out + x, mask + + +class TransformerEncoder(nn.Module): + def __init__(self, embed_dim: int, num_heads: int, num_layers: int = 1, dropout: float = 0.0): + super().__init__() + layer = nn.TransformerEncoderLayer( + d_model=embed_dim, + nhead=num_heads, + dim_feedforward=embed_dim * 4, + dropout=dropout, + batch_first=True, + ) + self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + b, f, length, d = x.shape + x_2d = x.reshape(b * f, length, d) + key_mask = None + if mask is not None: + key_mask = ~mask.reshape(b * f, length) + out = self.encoder(x_2d, src_key_padding_mask=key_mask) + return out.reshape(b, f, length, d), mask + + +class CrossFrameAttention(nn.Module): + def __init__(self, embed_dim: int, num_heads: int): + super().__init__() + self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) + self.norm = nn.LayerNorm(embed_dim) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + b, f, length, d = x.shape + # Treat frames as sequence: (B*L, F, D) + x_t = x.permute(0, 2, 1, 3).reshape(b * length, f, d) + key_mask = None + if mask is not None: + key_mask = ~mask.permute(0, 2, 1).reshape(b * length, f) + out, _ = self.attn(x_t, x_t, x_t, key_padding_mask=key_mask, need_weights=False) + out = self.norm(x_t + out).reshape(b, length, f, d).permute(0, 2, 1, 3) + return out, mask + + +class Embedding(nn.Module): + """Minimal embedding stub for compatibility until Task 6 expands it.""" + + def __init__( + self, + input_type: str, + vocab_size: int, + embedding_size: int, + use_embedding_layer: bool = True, + ): + super().__init__() + self.input_type = input_type + self.vocab_size = vocab_size + self.embedding_size = embedding_size + self.use_embedding_layer = use_embedding_layer + if use_embedding_layer: + self.embed = nn.Embedding(vocab_size, embedding_size) + else: + self.project = nn.LazyLinear(embedding_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_embedding_layer: + return self.embed(x) + return self.project(x) diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py new file mode 100644 index 0000000..9d51947 --- /dev/null +++ b/src/jaeger/nnlib/pytorch/models.py @@ -0,0 +1,74 @@ +from typing import List, Optional + +import torch +import torch.nn as nn + +from jaeger.nnlib.pytorch.layers import ( + AxialAttention, + CrossFrameAttention, + GatedFrameGlobalMaxPooling, + GeLU, + MaskedBatchNorm, + MaskedConv1D, + MaskedGlobalAvgPooling, + MaskedLayerNorm, + ResidualBlock, + TransformerEncoder, +) + + +class RepresentationModel(nn.Module): + def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str = "average"): + super().__init__() + self.embedding = embedding + self.blocks = nn.ModuleList() + self.return_nmd = False + self.output_dim = None + self.nmd_dim = None + for cfg in hidden_layers: + name = cfg["name"].lower() + config = cfg.get("config", {}) + if name == "masked_conv1d": + self.blocks.append(MaskedConv1D(**config)) + elif name == "masked_batchnorm": + self.blocks.append(MaskedBatchNorm(**config)) + elif name == "masked_layer_norm": + self.blocks.append(MaskedLayerNorm(**config)) + elif name == "axial_attention": + self.blocks.append(AxialAttention(**config)) + elif name == "cross_frame_attention": + self.blocks.append(CrossFrameAttention(**config)) + elif name == "transformer_encoder": + self.blocks.append(TransformerEncoder(**config)) + elif name == "residual_block": + inner = self.blocks.pop() if self.blocks else None + self.blocks.append(ResidualBlock(inner)) + elif name == "dense": + self.blocks.append(nn.Linear(**config)) + elif name == "activation": + self.blocks.append(GeLU() if config.get("activation") == "gelu" else nn.ReLU()) + elif name == "dropout": + self.blocks.append(nn.Dropout(config.get("rate", 0.0))) + else: + raise ValueError(f"Unknown layer type: {name}") + self.pooler = self._build_pooler(pooling) + + def _build_pooler(self, pooling: str): + if pooling == "average": + return MaskedGlobalAvgPooling() + elif pooling == "gatedframe": + return GatedFrameGlobalMaxPooling(return_gate=False) + raise ValueError(f"Unknown pooling: {pooling}") + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + x = self.embedding(x) + nmd = None + for block in self.blocks: + if hasattr(block, "return_nmd") and block.return_nmd: + x, nmd = block(x, mask) + else: + x, mask = block(x, mask) + pooled = self.pooler(x, mask) + if nmd is not None: + return pooled, nmd + return pooled diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 850312a..919bb46 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -1,11 +1,15 @@ import torch from jaeger.nnlib.pytorch.layers import ( AxialAttention, + CrossFrameAttention, GeLU, GatedFrameGlobalMaxPooling, MaskedBatchNorm, MaskedConv1D, MaskedLayerNorm, + ResidualBlock, + SinusoidalPositionEmbedding, + TransformerEncoder, ) @@ -212,3 +216,60 @@ def test_axial_attention_with_mask(): assert out.shape == (2, 6, 10, 4) assert torch.isfinite(out).all() assert out_mask is mask + + +def test_sinusoidal_position_embedding_shape(): + x = torch.randn(2, 6, 10, 4) + layer = SinusoidalPositionEmbedding() + out = layer(x) + assert out.shape == (2, 6, 10, 4) + + +def test_residual_block_with_masked_layer(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + layer = ResidualBlock(MaskedLayerNorm(num_features=4)) + out, out_mask = layer(x, mask) + assert out.shape == (2, 6, 10, 4) + assert out_mask.shape == (2, 6, 10) + + +def test_transformer_encoder_shape(): + x = torch.randn(2, 6, 10, 4) + layer = TransformerEncoder(embed_dim=4, num_heads=2) + out, out_mask = layer(x) + assert out.shape == (2, 6, 10, 4) + + +def test_cross_frame_attention_shape(): + x = torch.randn(2, 6, 10, 4) + layer = CrossFrameAttention(embed_dim=4, num_heads=2) + out, out_mask = layer(x) + assert out.shape == (2, 6, 10, 4) + + +def test_representation_model_forward(): + from jaeger.nnlib.pytorch.models import RepresentationModel + from jaeger.nnlib.pytorch.layers import Embedding + + embedding = Embedding( + input_type="translated", + vocab_size=65, + embedding_size=32, + use_embedding_layer=True, + ) + model = RepresentationModel( + embedding=embedding, + hidden_layers=[ + {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}}, + {"name": "masked_batchnorm", "config": {"num_features": 16, "return_nmd": True}}, + ], + pooling="average", + ) + x = torch.randint(0, 65, (2, 6, 50)) + mask = torch.ones(2, 6, 50, dtype=torch.bool) + out = model(x, mask) + assert isinstance(out, tuple) + pooled, nmd = out + assert pooled.shape == (2, 16) + assert nmd.shape == (2, 16) From df9123ed193e9543cdb7e07262ac2c2c0ae40bf4 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:35:15 +0200 Subject: [PATCH 12/64] feat(pytorch): fix RepresentationModel and remaining layer issues --- src/jaeger/nnlib/pytorch/layers.py | 4 +-- src/jaeger/nnlib/pytorch/models.py | 17 +++++++--- tests/unit/nnlib/pytorch/test_layers.py | 42 +++++++++++++++++++++++-- 3 files changed, 54 insertions(+), 9 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 03d8e08..1625f12 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -371,7 +371,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) pe = torch.zeros(length, d, device=x.device) pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]]) return x + pe.view(1, 1, length, d) @@ -431,7 +431,7 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): return out, mask -class Embedding(nn.Module): +class InputEmbedding(nn.Module): """Minimal embedding stub for compatibility until Task 6 expands it.""" def __init__( diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index 9d51947..96a704d 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -23,8 +23,6 @@ def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str self.embedding = embedding self.blocks = nn.ModuleList() self.return_nmd = False - self.output_dim = None - self.nmd_dim = None for cfg in hidden_layers: name = cfg["name"].lower() config = cfg.get("config", {}) @@ -44,6 +42,9 @@ def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str inner = self.blocks.pop() if self.blocks else None self.blocks.append(ResidualBlock(inner)) elif name == "dense": + if "units" in config: + config = dict(config) + config["out_features"] = config.pop("units") self.blocks.append(nn.Linear(**config)) elif name == "activation": self.blocks.append(GeLU() if config.get("activation") == "gelu" else nn.ReLU()) @@ -60,14 +61,20 @@ def _build_pooler(self, pooling: str): return GatedFrameGlobalMaxPooling(return_gate=False) raise ValueError(f"Unknown pooling: {pooling}") + def _accepts_mask(self, block): + return "mask" in block.forward.__code__.co_varnames + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): x = self.embedding(x) nmd = None for block in self.blocks: - if hasattr(block, "return_nmd") and block.return_nmd: - x, nmd = block(x, mask) + if self._accepts_mask(block): + if hasattr(block, "return_nmd") and block.return_nmd: + x, nmd = block(x, mask) + else: + x, mask = block(x, mask) else: - x, mask = block(x, mask) + x = block(x) pooled = self.pooler(x, mask) if nmd is not None: return pooled, nmd diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 919bb46..7888da2 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -4,8 +4,10 @@ CrossFrameAttention, GeLU, GatedFrameGlobalMaxPooling, + InputEmbedding, MaskedBatchNorm, MaskedConv1D, + MaskedGlobalAvgPooling, MaskedLayerNorm, ResidualBlock, SinusoidalPositionEmbedding, @@ -250,9 +252,8 @@ def test_cross_frame_attention_shape(): def test_representation_model_forward(): from jaeger.nnlib.pytorch.models import RepresentationModel - from jaeger.nnlib.pytorch.layers import Embedding - embedding = Embedding( + embedding = InputEmbedding( input_type="translated", vocab_size=65, embedding_size=32, @@ -273,3 +274,40 @@ def test_representation_model_forward(): pooled, nmd = out assert pooled.shape == (2, 16) assert nmd.shape == (2, 16) + + +def test_sinusoidal_position_embedding_odd_dim(): + x = torch.randn(2, 6, 10, 5) + layer = SinusoidalPositionEmbedding() + out = layer(x) + assert out.shape == (2, 6, 10, 5) + + +def test_masked_global_avg_pooling_shape(): + x = torch.randn(2, 6, 10, 4) + layer = MaskedGlobalAvgPooling() + out = layer(x) + assert out.shape == (2, 4) + + +def test_masked_global_avg_pooling_ignores_masked(): + x = torch.zeros(1, 1, 4, 1) + x[0, 0, 0, 0] = 100.0 + x[0, 0, 1, 0] = 0.0 + x[0, 0, 2, 0] = 0.0 + x[0, 0, 3, 0] = 0.0 + mask = torch.ones(1, 1, 4, dtype=torch.bool) + mask[0, 0, 0] = False + layer = MaskedGlobalAvgPooling() + out = layer(x, mask) + assert torch.allclose(out, torch.zeros_like(out)) + + +def test_masked_global_avg_pooling_zero_frame(): + x = torch.randn(2, 3, 5, 4) + mask = torch.ones(2, 3, 5, dtype=torch.bool) + mask[0, 0, :] = False + layer = MaskedGlobalAvgPooling() + out = layer(x, mask) + assert out.shape == (2, 4) + assert torch.isfinite(out).all() From 6cf8903163a575139ca57e76e1a7dbc5cdaff2e4 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:39:56 +0200 Subject: [PATCH 13/64] feat(pytorch): add model modules (embedding, heads, JaegerModel) --- src/jaeger/nnlib/pytorch/layers.py | 24 ---- src/jaeger/nnlib/pytorch/models.py | 139 +++++++++++++++++++++++- tests/unit/nnlib/pytorch/test_layers.py | 5 +- tests/unit/nnlib/pytorch/test_models.py | 53 +++++++++ 4 files changed, 193 insertions(+), 28 deletions(-) create mode 100644 tests/unit/nnlib/pytorch/test_models.py diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 1625f12..f5b9f09 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -431,27 +431,3 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): return out, mask -class InputEmbedding(nn.Module): - """Minimal embedding stub for compatibility until Task 6 expands it.""" - - def __init__( - self, - input_type: str, - vocab_size: int, - embedding_size: int, - use_embedding_layer: bool = True, - ): - super().__init__() - self.input_type = input_type - self.vocab_size = vocab_size - self.embedding_size = embedding_size - self.use_embedding_layer = use_embedding_layer - if use_embedding_layer: - self.embed = nn.Embedding(vocab_size, embedding_size) - else: - self.project = nn.LazyLinear(embedding_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.use_embedding_layer: - return self.embed(x) - return self.project(x) diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index 96a704d..dd04adb 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import Dict, List, Optional import torch import torch.nn as nn @@ -17,6 +17,107 @@ ) +class Embedding(nn.Module): + """Translates DNA into 6-frame codon embeddings or nucleotide one-hot.""" + + def __init__( + self, + input_type: str, + vocab_size: Optional[int], + embedding_size: int, + use_embedding_layer: bool, + use_positional_embeddings: bool = False, + positional_embedding_length: Optional[int] = None, + ): + super().__init__() + self.input_type = input_type + self.use_embedding_layer = use_embedding_layer + self.use_positional_embeddings = use_positional_embeddings + + if input_type == "translated": + if use_embedding_layer: + self.embed = nn.Embedding(vocab_size, embedding_size, padding_idx=0) + nn.init.orthogonal_(self.embed.weight) + else: + self.embed = nn.Linear(vocab_size, embedding_size, bias=False) + nn.init.orthogonal_(self.embed.weight) + elif input_type == "nucleotide": + self.embed = None + else: + raise ValueError(f"Invalid input_type: {input_type}") + + if use_positional_embeddings: + from jaeger.nnlib.pytorch.layers import SinusoidalPositionEmbedding + + self.positional = SinusoidalPositionEmbedding(positional_embedding_length) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x shape depends on input_type; caller must pass correct shape + if self.input_type == "translated": + if self.use_embedding_layer: + return self.embed(x) + else: + return self.embed(x) + return x + + +class ClassificationHead(nn.Module): + """Dense head for class prediction.""" + + def __init__(self, input_dim: int, num_classes: int, hidden_units: List[int], dropout: float = 0.0): + super().__init__() + layers = [] + prev = input_dim + for units in hidden_units: + layers.append(nn.Linear(prev, units)) + layers.append(nn.LayerNorm(units)) + layers.append(GeLU()) + if dropout > 0: + layers.append(nn.Dropout(dropout)) + prev = units + layers.append(nn.Linear(prev, num_classes)) + self.net = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class ReliabilityHead(nn.Module): + """Dense head for confidence estimation.""" + + def __init__(self, input_dim: int, num_classes: int, hidden_units: List[int], dropout: float = 0.0): + super().__init__() + layers = [] + prev = input_dim + for units in hidden_units: + layers.append(nn.Linear(prev, units)) + layers.append(nn.LayerNorm(units)) + layers.append(GeLU()) + if dropout > 0: + layers.append(nn.Dropout(dropout)) + prev = units + layers.append(nn.Linear(prev, num_classes)) + self.net = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class ProjectionHead(nn.Module): + """Projection head for ArcFace self-supervised pretraining.""" + + def __init__(self, input_dim: int, projection_dim: int): + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, projection_dim), + nn.LayerNorm(projection_dim), + GeLU(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + class RepresentationModel(nn.Module): def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str = "average"): super().__init__() @@ -79,3 +180,39 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): if nmd is not None: return pooled, nmd return pooled + + +class JaegerModel(nn.Module): + """Combined model exposing all outputs for inference.""" + + def __init__( + self, + rep_model: nn.Module, + classification_head: nn.Module, + reliability_head: Optional[nn.Module] = None, + ): + super().__init__() + self.rep_model = rep_model + self.classification_head = classification_head + self.reliability_head = reliability_head + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + outputs = self.rep_model(x, mask) + if isinstance(outputs, tuple): + embedding = outputs[0] + nmd = outputs[1] if len(outputs) > 1 else None + gate = outputs[2] if len(outputs) > 2 else None + else: + embedding = outputs + nmd = None + gate = None + + prediction = self.classification_head(embedding) + result: Dict[str, torch.Tensor] = {"prediction": prediction, "embedding": embedding} + if nmd is not None: + result["nmd"] = nmd + if gate is not None: + result["gate"] = gate + if self.reliability_head is not None and nmd is not None: + result["reliability"] = self.reliability_head(nmd) + return result diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index 7888da2..cef8287 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -4,7 +4,6 @@ CrossFrameAttention, GeLU, GatedFrameGlobalMaxPooling, - InputEmbedding, MaskedBatchNorm, MaskedConv1D, MaskedGlobalAvgPooling, @@ -251,9 +250,9 @@ def test_cross_frame_attention_shape(): def test_representation_model_forward(): - from jaeger.nnlib.pytorch.models import RepresentationModel + from jaeger.nnlib.pytorch.models import Embedding, RepresentationModel - embedding = InputEmbedding( + embedding = Embedding( input_type="translated", vocab_size=65, embedding_size=32, diff --git a/tests/unit/nnlib/pytorch/test_models.py b/tests/unit/nnlib/pytorch/test_models.py new file mode 100644 index 0000000..c32d60c --- /dev/null +++ b/tests/unit/nnlib/pytorch/test_models.py @@ -0,0 +1,53 @@ +import torch +from jaeger.nnlib.pytorch.models import ClassificationHead, Embedding, JaegerModel, ReliabilityHead, RepresentationModel + + +def test_embedding_translated_with_embedding_layer(): + emb = Embedding(input_type="translated", vocab_size=65, embedding_size=32, use_embedding_layer=True) + x = torch.randint(0, 65, (2, 6, 50)) + out = emb(x) + assert out.shape == (2, 6, 50, 32) + + +def test_embedding_translated_without_embedding_layer(): + emb = Embedding(input_type="translated", vocab_size=65, embedding_size=32, use_embedding_layer=False) + x = torch.randn(2, 6, 50, 65) + out = emb(x) + assert out.shape == (2, 6, 50, 32) + + +def test_classification_head_output_shape(): + head = ClassificationHead(input_dim=64, num_classes=3, hidden_units=[128, 64]) + x = torch.randn(2, 64) + out = head(x) + assert out.shape == (2, 3) + + +def test_reliability_head_output_shape(): + head = ReliabilityHead(input_dim=64, num_classes=1, hidden_units=[32]) + x = torch.randn(2, 64) + out = head(x) + assert out.shape == (2, 1) + + +def test_jaeger_model_forward(): + embedding = Embedding(input_type="translated", vocab_size=65, embedding_size=32, use_embedding_layer=True) + rep_model = RepresentationModel( + embedding=embedding, + hidden_layers=[ + {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}}, + {"name": "masked_batchnorm", "config": {"num_features": 16, "return_nmd": True}}, + ], + pooling="average", + ) + classification_head = ClassificationHead(input_dim=16, num_classes=3, hidden_units=[8]) + reliability_head = ReliabilityHead(input_dim=16, num_classes=1, hidden_units=[8]) + model = JaegerModel(rep_model=rep_model, classification_head=classification_head, reliability_head=reliability_head) + + x = torch.randint(0, 65, (2, 6, 50)) + mask = torch.ones(2, 6, 50, dtype=torch.bool) + out = model(x, mask) + assert out["prediction"].shape == (2, 3) + assert out["embedding"].shape == (2, 16) + assert out["nmd"].shape == (2, 16) + assert out["reliability"].shape == (2, 1) From 707e79e7e0f4dd28a79a5346d80d54735d8cadd7 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:42:00 +0200 Subject: [PATCH 14/64] feat(pytorch): apply positional embeddings and add ProjectionHead test --- src/jaeger/nnlib/pytorch/models.py | 7 +++--- tests/unit/nnlib/pytorch/test_models.py | 33 ++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index dd04adb..16221ef 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -54,10 +54,9 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: # x shape depends on input_type; caller must pass correct shape if self.input_type == "translated": - if self.use_embedding_layer: - return self.embed(x) - else: - return self.embed(x) + x = self.embed(x) + if self.use_positional_embeddings: + x = self.positional(x) return x diff --git a/tests/unit/nnlib/pytorch/test_models.py b/tests/unit/nnlib/pytorch/test_models.py index c32d60c..d917ec5 100644 --- a/tests/unit/nnlib/pytorch/test_models.py +++ b/tests/unit/nnlib/pytorch/test_models.py @@ -1,5 +1,5 @@ import torch -from jaeger.nnlib.pytorch.models import ClassificationHead, Embedding, JaegerModel, ReliabilityHead, RepresentationModel +from jaeger.nnlib.pytorch.models import ClassificationHead, Embedding, JaegerModel, ProjectionHead, ReliabilityHead, RepresentationModel def test_embedding_translated_with_embedding_layer(): @@ -16,6 +16,37 @@ def test_embedding_translated_without_embedding_layer(): assert out.shape == (2, 6, 50, 32) +def test_embedding_with_positional_embeddings(): + emb_pos = Embedding( + input_type="translated", + vocab_size=65, + embedding_size=32, + use_embedding_layer=True, + use_positional_embeddings=True, + positional_embedding_length=10000, + ) + emb_no_pos = Embedding( + input_type="translated", + vocab_size=65, + embedding_size=32, + use_embedding_layer=True, + use_positional_embeddings=False, + ) + x = torch.randint(0, 65, (2, 6, 50)) + out_pos = emb_pos(x) + out_no_pos = emb_no_pos(x) + assert out_pos.shape == (2, 6, 50, 32) + assert out_no_pos.shape == (2, 6, 50, 32) + assert not torch.allclose(out_pos, out_no_pos) + + +def test_projection_head_output_shape(): + head = ProjectionHead(input_dim=64, projection_dim=128) + x = torch.randn(2, 64) + out = head(x) + assert out.shape == (2, 128) + + def test_classification_head_output_shape(): head = ClassificationHead(input_dim=64, num_classes=3, hidden_units=[128, 64]) x = torch.randn(2, 64) From 6c61c57b9992cfa714b403ec356d6108894fc368 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:45:25 +0200 Subject: [PATCH 15/64] feat(pytorch): add ArcFace loss --- src/jaeger/nnlib/pytorch/losses.py | 47 +++++++++++++++++++++++++ tests/unit/nnlib/pytorch/test_losses.py | 31 ++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 src/jaeger/nnlib/pytorch/losses.py create mode 100644 tests/unit/nnlib/pytorch/test_losses.py diff --git a/src/jaeger/nnlib/pytorch/losses.py b/src/jaeger/nnlib/pytorch/losses.py new file mode 100644 index 0000000..f7717cf --- /dev/null +++ b/src/jaeger/nnlib/pytorch/losses.py @@ -0,0 +1,47 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ArcFaceLoss(nn.Module): + """ArcFace additive angular margin loss.""" + + def __init__( + self, + num_classes: int, + embedding_dim: int, + margin: float = 0.5, + scale: float = 30.0, + onehot: bool = True, + ): + super().__init__() + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.margin = margin + self.scale = scale + self.onehot = onehot + self.weight = nn.Parameter(torch.Tensor(num_classes, embedding_dim)) + nn.init.xavier_uniform_(self.weight) + + def forward(self, labels: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor: + # embeddings: (B, D), labels: (B, C) one-hot or (B,) long + embeddings_norm = F.normalize(embeddings, p=2, dim=1) + weight_norm = F.normalize(self.weight, p=2, dim=1) + cos_t = torch.matmul(embeddings_norm, weight_norm.t()) + + if self.onehot: + target = labels.argmax(dim=1) + else: + target = labels + + # Additive angular margin + cos_m = math.cos(self.margin) + sin_m = math.sin(self.margin) + sin_t = torch.sqrt(1.0 - cos_t.pow(2) + 1e-6) + cos_t_plus_m = cos_t * cos_m - sin_t * sin_m + one_hot = F.one_hot(target, num_classes=self.num_classes).to(cos_t.dtype) + logits = one_hot * cos_t_plus_m + (1.0 - one_hot) * cos_t + logits = logits * self.scale + return F.cross_entropy(logits, target) diff --git a/tests/unit/nnlib/pytorch/test_losses.py b/tests/unit/nnlib/pytorch/test_losses.py new file mode 100644 index 0000000..f0fd81f --- /dev/null +++ b/tests/unit/nnlib/pytorch/test_losses.py @@ -0,0 +1,31 @@ +import torch + +from jaeger.nnlib.pytorch.losses import ArcFaceLoss + + +def test_arcface_loss_shape(): + loss = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.5, scale=30.0) + labels = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float32) + embeddings = torch.randn(3, 64) + out = loss(labels, embeddings) + assert out.ndim == 0 + + +def test_arcface_loss_class_index_labels(): + loss = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.5, scale=30.0, onehot=False) + labels = torch.tensor([0, 1, 2], dtype=torch.long) + embeddings = torch.randn(3, 64) + out = loss(labels, embeddings) + assert out.ndim == 0 + + +def test_arcface_loss_changes_with_margin(): + labels = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float32) + embeddings = torch.randn(3, 64) + + loss_a = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.1, scale=30.0) + loss_b = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.9, scale=30.0) + + out_a = loss_a(labels, embeddings) + out_b = loss_b(labels, embeddings) + assert out_a.item() != out_b.item() From 895d454a3c41bb5da56367b587011332836da31a Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:48:35 +0200 Subject: [PATCH 16/64] feat(pytorch): add per-class precision/recall metrics --- src/jaeger/nnlib/pytorch/metrics.py | 38 ++++++++++++++++++++++ tests/unit/nnlib/pytorch/test_metrics.py | 41 ++++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 src/jaeger/nnlib/pytorch/metrics.py create mode 100644 tests/unit/nnlib/pytorch/test_metrics.py diff --git a/src/jaeger/nnlib/pytorch/metrics.py b/src/jaeger/nnlib/pytorch/metrics.py new file mode 100644 index 0000000..6e20e09 --- /dev/null +++ b/src/jaeger/nnlib/pytorch/metrics.py @@ -0,0 +1,38 @@ +import torch + + +class PrecisionForClass: + def __init__(self, class_id: int): + self.class_id = class_id + self.tp = 0 + self.fp = 0 + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + pred_labels = y_pred.argmax(dim=-1) + true_labels = y_true.argmax(dim=-1) + cls = self.class_id + self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) + self.fp += int(((pred_labels == cls) & (true_labels != cls)).sum().item()) + + def compute(self) -> float: + if self.tp + self.fp == 0: + return 0.0 + return self.tp / (self.tp + self.fp) + + +class RecallForClass(PrecisionForClass): + def __init__(self, class_id: int): + super().__init__(class_id) + self.fn = 0 + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + pred_labels = y_pred.argmax(dim=-1) + true_labels = y_true.argmax(dim=-1) + cls = self.class_id + self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) + self.fn += int(((pred_labels != cls) & (true_labels == cls)).sum().item()) + + def compute(self) -> float: + if self.tp + self.fn == 0: + return 0.0 + return self.tp / (self.tp + self.fn) diff --git a/tests/unit/nnlib/pytorch/test_metrics.py b/tests/unit/nnlib/pytorch/test_metrics.py new file mode 100644 index 0000000..5573dea --- /dev/null +++ b/tests/unit/nnlib/pytorch/test_metrics.py @@ -0,0 +1,41 @@ +import torch +from jaeger.nnlib.pytorch.metrics import PrecisionForClass, RecallForClass + + +def test_precision_for_class(): + metric = PrecisionForClass(class_id=1) + preds = torch.tensor([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]) + labels = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.float32) + metric.update(preds, labels) + result = metric.compute() + assert 0.0 <= result <= 1.0 + + +def test_recall_for_class(): + metric = RecallForClass(class_id=1) + preds = torch.tensor([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]) + labels = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.float32) + metric.update(preds, labels) + result = metric.compute() + assert 0.0 <= result <= 1.0 + + +def test_precision_for_class_zero_division(): + metric = PrecisionForClass(class_id=0) + preds = torch.tensor([[0.1, 0.9], [0.3, 0.7]]) + labels = torch.tensor([[0, 1], [0, 1]], dtype=torch.float32) + metric.update(preds, labels) + result = metric.compute() + assert result == 0.0 + + +def test_metrics_multiple_updates(): + metric = PrecisionForClass(class_id=1) + preds1 = torch.tensor([[0.1, 0.9], [0.8, 0.2]]) + labels1 = torch.tensor([[0, 1], [1, 0]], dtype=torch.float32) + preds2 = torch.tensor([[0.3, 0.7], [0.9, 0.1]]) + labels2 = torch.tensor([[0, 1], [1, 0]], dtype=torch.float32) + metric.update(preds1, labels1) + metric.update(preds2, labels2) + result = metric.compute() + assert 0.0 <= result <= 1.0 From e3486e05806b5bf90df3f8d76e0db8c92d742c94 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Sun, 14 Jun 2026 23:54:23 +0200 Subject: [PATCH 17/64] feat(pytorch): add ModelBuilder --- src/jaeger/nnlib/pytorch/builder.py | 122 +++++++++++++++++++++++ src/jaeger/nnlib/pytorch/models.py | 51 +++++++++- tests/unit/nnlib/pytorch/test_builder.py | 54 ++++++++++ 3 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 src/jaeger/nnlib/pytorch/builder.py create mode 100644 tests/unit/nnlib/pytorch/test_builder.py diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py new file mode 100644 index 0000000..b4ced70 --- /dev/null +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -0,0 +1,122 @@ +from typing import Any, Dict + +import torch +import torch.nn as nn + +from jaeger.nnlib.pytorch.models import ( + ClassificationHead, + Embedding, + JaegerModel, + ReliabilityHead, + RepresentationModel, +) + + +class ModelBuilder: + """Build PyTorch Jaeger models from YAML config.""" + + def __init__(self, config: Dict[str, Any]): + self.cfg = config + self.model_cfg = config.get("model", {}) + self.train_cfg = config.get("training", {}) + self.classifier_out_dim = int(self.model_cfg.get("classifier_out_dim", 0)) + self.reliability_out_dim = int(self.model_cfg.get("reliability_out_dim", 0)) + + def build_fragment_classifier(self) -> Dict[str, nn.Module]: + models: Dict[str, nn.Module] = {} + + embedding_cfg = self.model_cfg.get("embedding", {}) + embedding = Embedding( + input_type=embedding_cfg.get("input_type"), + vocab_size=embedding_cfg.get("vocab_size"), + embedding_size=embedding_cfg.get("embedding_size", 4), + use_embedding_layer=embedding_cfg.get("use_embedding_layer", False), + ) + + rep_cfg = self.model_cfg.get("representation_learner", {}) + rep_model = RepresentationModel( + embedding=embedding, + hidden_layers=rep_cfg.get("hidden_layers", []), + pooling=rep_cfg.get("pooling", "average"), + ) + models["rep_model"] = rep_model + + if "classifier" in self.model_cfg: + cls_cfg = self.model_cfg["classifier"] + head = ClassificationHead( + input_dim=cls_cfg.get("input_shape", rep_model.output_dim), + num_classes=self.classifier_out_dim, + hidden_units=[ + layer["config"]["units"] + for layer in cls_cfg.get("hidden_layers", [])[:-1] + ], + ) + models["classification_head"] = head + models["jaeger_classifier"] = nn.Sequential(rep_model, head) + + if "reliability_model" in self.model_cfg: + rel_cfg = self.model_cfg["reliability_model"] + rel_head = ReliabilityHead( + input_dim=rel_cfg.get("input_shape", rep_model.nmd_dim), + num_classes=self.reliability_out_dim, + hidden_units=[ + layer["config"]["units"] + for layer in rel_cfg.get("hidden_layers", [])[:-1] + ], + ) + models["reliability_head"] = rel_head + + models["jaeger_model"] = JaegerModel( + rep_model=rep_model, + classification_head=models["classification_head"], + reliability_head=models.get("reliability_head"), + ) + return models + + def compile_model( + self, models: Dict[str, nn.Module], train_branch: str = "classifier" + ) -> tuple: + opt_name = self.train_cfg.get("optimizer", "adam").lower() + opt_params = self.train_cfg.get("optimizer_params", {}) + opt_class = self._get_optimizer(opt_name) + + if train_branch == "classifier": + model = models["jaeger_classifier"] + optimizer = opt_class(model.parameters(), **opt_params) + loss = nn.CrossEntropyLoss() + return model, optimizer, loss + if train_branch == "reliability": + model = models["jaeger_reliability"] + optimizer = opt_class(model.parameters(), **opt_params) + loss = nn.BCEWithLogitsLoss() + return model, optimizer, loss + if train_branch == "pretrain": + model = models["jaeger_projection"] + optimizer = opt_class(model.parameters(), **opt_params) + loss = models["arcface_loss"] + return model, optimizer, loss + raise ValueError(f"Unknown train_branch: {train_branch}") + + def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: + from jaeger.nnlib.pytorch.metrics import PrecisionForClass, RecallForClass + + metrics = {} + out_dim = ( + self.classifier_out_dim + if branch == "classifier" + else self.reliability_out_dim + ) + for cls in range(out_dim): + metrics[f"precision_class_{cls}"] = PrecisionForClass(class_id=cls) + metrics[f"recall_class_{cls}"] = RecallForClass(class_id=cls) + return metrics + + def _get_optimizer(self, name: str) -> torch.optim.Optimizer: + optimizers = { + "adam": torch.optim.Adam, + "sgd": torch.optim.SGD, + "rmsprop": torch.optim.RMSprop, + } + if name not in optimizers: + raise ValueError(f"Unsupported optimizer: {name}") + return optimizers[name] diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index 16221ef..84beb3c 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from jaeger.nnlib.pytorch.layers import ( AxialAttention, @@ -54,7 +55,14 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: # x shape depends on input_type; caller must pass correct shape if self.input_type == "translated": - x = self.embed(x) + if self.use_embedding_layer: + x = self.embed(x) + else: + if not torch.is_floating_point(x): + x = F.one_hot(x, num_classes=self.embed.in_features).to( + self.embed.weight.dtype + ) + x = self.embed(x) if self.use_positional_embeddings: x = self.positional(x) return x @@ -153,6 +161,47 @@ def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str else: raise ValueError(f"Unknown layer type: {name}") self.pooler = self._build_pooler(pooling) + self.output_dim, self.nmd_dim = self._infer_dims() + + def _infer_dims(self): + output_dim = None + for block in reversed(self.blocks): + dim = self._block_output_dim(block) + if dim is not None: + output_dim = dim + break + + nmd_dim = None + for block in self.blocks: + dim = self._block_nmd_dim(block) + if dim is not None: + nmd_dim = dim + break + return output_dim, nmd_dim + + @staticmethod + def _block_output_dim(block: nn.Module) -> Optional[int]: + if isinstance(block, MaskedConv1D): + return int(block.filters) + if isinstance(block, (MaskedBatchNorm, MaskedLayerNorm)): + return int(block.num_features) + if isinstance(block, (AxialAttention, CrossFrameAttention)): + return int(block.attn.embed_dim) + if isinstance(block, TransformerEncoder): + return int(block.encoder.layers[0].self_attn.embed_dim) + if isinstance(block, nn.Linear): + return int(block.out_features) + if isinstance(block, ResidualBlock): + return RepresentationModel._block_output_dim(block.layer) + return None + + @staticmethod + def _block_nmd_dim(block: nn.Module) -> Optional[int]: + if isinstance(block, ResidualBlock): + return RepresentationModel._block_nmd_dim(block.layer) + if getattr(block, "return_nmd", False): + return RepresentationModel._block_output_dim(block) + return None def _build_pooler(self, pooling: str): if pooling == "average": diff --git a/tests/unit/nnlib/pytorch/test_builder.py b/tests/unit/nnlib/pytorch/test_builder.py new file mode 100644 index 0000000..ae5c05f --- /dev/null +++ b/tests/unit/nnlib/pytorch/test_builder.py @@ -0,0 +1,54 @@ +import torch + +from jaeger.nnlib.pytorch.builder import ModelBuilder + + +def test_builder_creates_jaeger_model(): + config = { + "model": { + "name": "test_model", + "classifier_out_dim": 3, + "reliability_out_dim": 1, + "class_label_map": [ + {"class": "phage", "label": 1}, + {"class": "bacteria", "label": 0}, + ], + "embedding": { + "input_type": "translated", + "use_embedding_layer": False, + "embedding_size": 32, + "input_shape": [6, None], + "vocab_size": 65, + "codon_depth": 1, + }, + "string_processor": {"codon": "CODON", "codon_id": "CODON_ID"}, + "representation_learner": { + "hidden_layers": [ + { + "name": "masked_conv1d", + "config": { + "filters": 16, + "kernel_size": 3, + "padding": "same", + }, + } + ], + "pooling": "average", + }, + "classifier": { + "input_shape": 16, + "hidden_layers": [{"name": "dense", "config": {"units": 3}}], + }, + }, + "training": { + "batch_size": 2, + "optimizer": "adam", + "optimizer_params": {"lr": 1e-3}, + }, + } + builder = ModelBuilder(config) + models = builder.build_fragment_classifier() + x = torch.randint(0, 65, (2, 6, 50)) + mask = torch.ones(2, 6, 50, dtype=torch.bool) + out = models["jaeger_model"](x, mask) + assert out["prediction"].shape == (2, 3) From 4bc86d744ee53897715463e130c976eb7fef9af5 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:00:00 +0200 Subject: [PATCH 18/64] feat(pytorch): fix ModelBuilder contract and mask propagation --- src/jaeger/nnlib/pytorch/builder.py | 59 +++++++++++++++--------- src/jaeger/nnlib/pytorch/models.py | 20 ++++++-- tests/unit/nnlib/pytorch/test_builder.py | 52 +++++++++++++++++++-- 3 files changed, 103 insertions(+), 28 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py index b4ced70..dfff3a7 100644 --- a/src/jaeger/nnlib/pytorch/builder.py +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -1,4 +1,4 @@ -from typing import Any, Dict +from typing import Any, Dict, Optional import torch import torch.nn as nn @@ -12,6 +12,21 @@ ) +class _ClassifierPipeline(nn.Module): + """Runs the representation model and classification head, propagating masks.""" + + def __init__(self, rep_model: nn.Module, head: nn.Module): + super().__init__() + self.rep_model = rep_model + self.head = head + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + embedding = self.rep_model(x, mask) + if isinstance(embedding, tuple): + embedding = embedding[0] + return self.head(embedding) + + class ModelBuilder: """Build PyTorch Jaeger models from YAML config.""" @@ -23,6 +38,9 @@ def __init__(self, config: Dict[str, Any]): self.reliability_out_dim = int(self.model_cfg.get("reliability_out_dim", 0)) def build_fragment_classifier(self) -> Dict[str, nn.Module]: + if "classifier" not in self.model_cfg: + raise ValueError("classifier config is required for build_fragment_classifier") + models: Dict[str, nn.Module] = {} embedding_cfg = self.model_cfg.get("embedding", {}) @@ -41,18 +59,18 @@ def build_fragment_classifier(self) -> Dict[str, nn.Module]: ) models["rep_model"] = rep_model - if "classifier" in self.model_cfg: - cls_cfg = self.model_cfg["classifier"] - head = ClassificationHead( - input_dim=cls_cfg.get("input_shape", rep_model.output_dim), - num_classes=self.classifier_out_dim, - hidden_units=[ - layer["config"]["units"] - for layer in cls_cfg.get("hidden_layers", [])[:-1] - ], - ) - models["classification_head"] = head - models["jaeger_classifier"] = nn.Sequential(rep_model, head) + cls_cfg = self.model_cfg["classifier"] + head = ClassificationHead( + input_dim=cls_cfg.get("input_shape", rep_model.output_dim), + num_classes=self.classifier_out_dim, + hidden_units=[ + layer["config"]["units"] + for layer in cls_cfg.get("hidden_layers", [])[:-1] + ], + dropout=cls_cfg.get("dropout", 0.0), + ) + models["classification_head"] = head + models["jaeger_classifier"] = _ClassifierPipeline(rep_model, head) if "reliability_model" in self.model_cfg: rel_cfg = self.model_cfg["reliability_model"] @@ -63,6 +81,7 @@ def build_fragment_classifier(self) -> Dict[str, nn.Module]: layer["config"]["units"] for layer in rel_cfg.get("hidden_layers", [])[:-1] ], + dropout=rel_cfg.get("dropout", 0.0), ) models["reliability_head"] = rel_head @@ -86,15 +105,13 @@ def compile_model( loss = nn.CrossEntropyLoss() return model, optimizer, loss if train_branch == "reliability": - model = models["jaeger_reliability"] - optimizer = opt_class(model.parameters(), **opt_params) - loss = nn.BCEWithLogitsLoss() - return model, optimizer, loss + raise NotImplementedError( + "reliability training branch is not implemented yet" + ) if train_branch == "pretrain": - model = models["jaeger_projection"] - optimizer = opt_class(model.parameters(), **opt_params) - loss = models["arcface_loss"] - return model, optimizer, loss + raise NotImplementedError( + "pretrain training branch is not implemented yet" + ) raise ValueError(f"Unknown train_branch: {train_branch}") def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index 84beb3c..08de46a 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -171,8 +171,9 @@ def _infer_dims(self): output_dim = dim break + # Match the forward loop: the last block that returns NMD wins. nmd_dim = None - for block in self.blocks: + for block in reversed(self.blocks): dim = self._block_nmd_dim(block) if dim is not None: nmd_dim = dim @@ -218,10 +219,21 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): nmd = None for block in self.blocks: if self._accepts_mask(block): - if hasattr(block, "return_nmd") and block.return_nmd: - x, nmd = block(x, mask) + out = block(x, mask) + if isinstance(out, tuple): + x, maybe_mask_or_nmd = out + if maybe_mask_or_nmd is None: + # mask unchanged + pass + elif getattr(block, "return_nmd", False): + nmd = maybe_mask_or_nmd + elif ( + isinstance(maybe_mask_or_nmd, torch.Tensor) + and maybe_mask_or_nmd.dtype == torch.bool + ): + mask = maybe_mask_or_nmd else: - x, mask = block(x, mask) + x = out else: x = block(x) pooled = self.pooler(x, mask) diff --git a/tests/unit/nnlib/pytorch/test_builder.py b/tests/unit/nnlib/pytorch/test_builder.py index ae5c05f..dbbf409 100644 --- a/tests/unit/nnlib/pytorch/test_builder.py +++ b/tests/unit/nnlib/pytorch/test_builder.py @@ -1,10 +1,11 @@ +import pytest import torch from jaeger.nnlib.pytorch.builder import ModelBuilder -def test_builder_creates_jaeger_model(): - config = { +def _minimal_config(): + return { "model": { "name": "test_model", "classifier_out_dim": 3, @@ -46,9 +47,54 @@ def test_builder_creates_jaeger_model(): "optimizer_params": {"lr": 1e-3}, }, } - builder = ModelBuilder(config) + + +def test_builder_creates_jaeger_model(): + builder = ModelBuilder(_minimal_config()) models = builder.build_fragment_classifier() x = torch.randint(0, 65, (2, 6, 50)) mask = torch.ones(2, 6, 50, dtype=torch.bool) out = models["jaeger_model"](x, mask) assert out["prediction"].shape == (2, 3) + + +def test_builder_compile_classifier(): + builder = ModelBuilder(_minimal_config()) + models = builder.build_fragment_classifier() + model, optimizer, loss = builder.compile_model(models, "classifier") + x = torch.randint(0, 65, (2, 6, 50)) + mask = torch.ones(2, 6, 50, dtype=torch.bool) + out = model(x, mask) + assert out.shape == (2, 3) + assert optimizer is not None + assert isinstance(loss, torch.nn.CrossEntropyLoss) + + +def test_builder_missing_classifier_raises(): + config = _minimal_config() + del config["model"]["classifier"] + builder = ModelBuilder(config) + with pytest.raises(ValueError, match="classifier config is required"): + builder.build_fragment_classifier() + + +def test_builder_mask_propagation(): + builder = ModelBuilder(_minimal_config()) + models = builder.build_fragment_classifier() + model, _, _ = builder.compile_model(models, "classifier") + model.eval() + + x = torch.zeros(2, 6, 50, dtype=torch.long) + x[0, 0, 25] = 50 + + mask_all = torch.ones(2, 6, 50, dtype=torch.bool) + mask_masked = mask_all.clone() + mask_masked[0, 0, 25] = False + + with torch.no_grad(): + out_all = model(x, mask_all) + out_masked = model(x, mask_masked) + + assert out_all.shape == (2, 3) + assert not torch.allclose(out_all[0], out_masked[0], atol=1e-6) + assert torch.allclose(out_all[1], out_masked[1]) From 93a1fd1eddae7acf6f74d526170c0878c00a58cb Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:22:49 +0200 Subject: [PATCH 19/64] test(pytorch): add TF vs PyTorch forward parity test --- src/jaeger/nnlib/pytorch/layers.py | 87 +++++++----- tests/integration/test_pytorch_parity.py | 174 +++++++++++++++++++++++ 2 files changed, 226 insertions(+), 35 deletions(-) create mode 100644 tests/integration/test_pytorch_parity.py diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index f5b9f09..537dca4 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -53,9 +53,7 @@ def __init__( super().__init__() padding_norm = padding.lower() if padding_norm not in ("same", "valid"): - raise ValueError( - f"padding must be 'same' or 'valid', got {padding!r}" - ) + raise ValueError(f"padding must be 'same' or 'valid', got {padding!r}") if kernel_initializer != "glorot_uniform": raise NotImplementedError( @@ -119,9 +117,7 @@ def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if x.ndim != 4: - raise ValueError( - f"Expected 4D input of shape (B, F, L, C), got {x.ndim}D" - ) + raise ValueError(f"Expected 4D input of shape (B, F, L, C), got {x.ndim}D") b, f, length, c = x.shape @@ -146,8 +142,10 @@ def forward( if mask is not None: x = x * mask.unsqueeze(-1).to(x.dtype) - # Merge batch and frame dims: (B*F, C, L) - x_2d = x.reshape(b * f, c, length) + # Merge batch and frame dims. The input layout is (B, F, L, C); Conv1d + # expects (N, C, L), so permute before reshaping to keep channels and + # length in the right order. + x_2d = x.permute(0, 1, 3, 2).reshape(b * f, c, length) if self.padding == "same": pad_total = self._resolve_padding(length) pad_left = pad_total // 2 @@ -156,30 +154,39 @@ def forward( out = self.conv(x_2d) _, _, length_out = out.shape - out = out.reshape(b, f, length_out, self.filters) + # Conv1d output is (B*F, C_out, L_out); restore to (B, F, L_out, C_out). + out = out.permute(0, 2, 1).reshape(b, f, length_out, self.filters) if self.activation is not None: out = self.activation(out) out_mask = None if mask is not None: - mask_f = mask.reshape(b * f, 1, length).to(x.dtype) - if self.padding == "same": - mask_f = F.pad(mask_f, (pad_left, pad_right)) - with torch.no_grad(): - kernel = torch.ones( - (1, 1, self.kernel_size), - dtype=mask_f.dtype, - device=mask_f.device, - ) - out_mask = F.conv1d( - mask_f, - kernel, - stride=self.strides, - dilation=self.dilation_rate, - ) - out_mask = (out_mask >= self.kernel_size).squeeze(1) - out_mask = out_mask.reshape(b, f, length_out) + if self.padding == "same" and self.strides == 1: + # With stride-1 "same" padding, output positions map one-to-one + # to input positions, so the center of the kernel falls on a + # valid position exactly when the input mask is True. This + # avoids marking boundary positions as invalid because of the + # zero padding used to implement "same". + out_mask = mask[:, :, :length_out] + else: + mask_f = mask.reshape(b * f, 1, length).to(x.dtype) + if self.padding == "same": + mask_f = F.pad(mask_f, (pad_left, pad_right)) + with torch.no_grad(): + kernel = torch.ones( + (1, 1, self.kernel_size), + dtype=mask_f.dtype, + device=mask_f.device, + ) + out_mask = F.conv1d( + mask_f, + kernel, + stride=self.strides, + dilation=self.dilation_rate, + ) + out_mask = (out_mask >= self.kernel_size).squeeze(1) + out_mask = out_mask.reshape(b, f, length_out) return out, out_mask @@ -269,7 +276,9 @@ def __init__(self, num_features: int, eps: float = 1e-3): self.gamma = nn.Parameter(torch.ones(num_features)) self.beta = nn.Parameter(torch.zeros(num_features)) - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: xf = x.to(torch.float32) if mask is not None: mask_f = mask.unsqueeze(-1).to(torch.float32) @@ -345,14 +354,19 @@ def forward( class MaskedGlobalAvgPooling(nn.Module): """Global average pooling over the sequence axis, respecting masks.""" - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: # x: (B, F, L, D) if mask is not None: mask_f = mask.unsqueeze(-1).to(x.dtype) masked_x = x * mask_f - count = mask_f.sum(dim=2, keepdim=True).clamp(min=1.0) - pooled = masked_x.sum(dim=2, keepdim=True) / count - return pooled.mean(dim=(1, 2)) + # Compute the global masked average in one step to match the + # TensorFlow GlobalAveragePooling2D behavior and reduce rounding + # differences from two-stage averaging. + count = mask_f.sum(dim=(1, 2), keepdim=True).clamp(min=1.0) + pooled = masked_x.sum(dim=(1, 2), keepdim=True) / count + return pooled.squeeze(dim=(1, 2)) return x.mean(dim=(1, 2)) @@ -381,7 +395,10 @@ def __init__(self, layer: nn.Module): self.layer = layer def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): - if hasattr(self.layer, "forward") and "mask" in self.layer.forward.__code__.co_varnames: + if ( + hasattr(self.layer, "forward") + and "mask" in self.layer.forward.__code__.co_varnames + ): out = self.layer(x, mask) if isinstance(out, tuple): out, out_mask = out @@ -392,7 +409,9 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): class TransformerEncoder(nn.Module): - def __init__(self, embed_dim: int, num_heads: int, num_layers: int = 1, dropout: float = 0.0): + def __init__( + self, embed_dim: int, num_heads: int, num_layers: int = 1, dropout: float = 0.0 + ): super().__init__() layer = nn.TransformerEncoderLayer( d_model=embed_dim, @@ -429,5 +448,3 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): out, _ = self.attn(x_t, x_t, x_t, key_padding_mask=key_mask, need_weights=False) out = self.norm(x_t + out).reshape(b, length, f, d).permute(0, 2, 1, 3) return out, mask - - diff --git a/tests/integration/test_pytorch_parity.py b/tests/integration/test_pytorch_parity.py new file mode 100644 index 0000000..0790db1 --- /dev/null +++ b/tests/integration/test_pytorch_parity.py @@ -0,0 +1,174 @@ +"""Forward parity test between the legacy TensorFlow builder and PyTorch builder. + +This test builds the same tiny fragment classifier with both backends, copies the +TensorFlow weights into the PyTorch model, and asserts that the forward outputs +agree within tolerance. +""" + +import os + +# Suppress TensorFlow logging before importing it. +os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") +os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0") +os.environ.setdefault("GRPC_VERBOSITY", "ERROR") +os.environ.setdefault("GLOG_minloglevel", "2") +os.environ.setdefault("ABSL_MIN_LOG_LEVEL", "2") + +import numpy as np +import torch + +from jaeger.nnlib.builder import DynamicModelBuilder as TFBuilder +from jaeger.nnlib.pytorch.builder import ModelBuilder as PTBuilder + + +CONFIG = { + "model": { + "name": "parity_test", + "classifier_out_dim": 3, + "reliability_out_dim": 1, + "class_label_map": [ + {"class": "phage", "label": 1}, + {"class": "bacteria", "label": 0}, + {"class": "eukarya", "label": 2}, + ], + "embedding": { + "input_type": "translated", + "use_embedding_layer": True, + "embedding_size": 32, + "input_shape": [6, None], + "vocab_size": 65, + "codon_depth": 1, + "embedding_regularizer": "l2", + "embedding_regularizer_w": 1e-5, + }, + "string_processor": { + "codon": "CODON", + "codon_id": "CODON_ID", + "crop_size": 50, + "input_type": "translated", + "use_embedding_layer": True, + }, + "representation_learner": { + "hidden_layers": [ + { + "name": "masked_conv1d", + "config": { + "filters": 16, + "kernel_size": 3, + "padding": "same", + }, + }, + ], + "pooling": "average", + }, + "classifier": { + "input_shape": 16, + "hidden_layers": [{"name": "dense", "config": {"units": 3}}], + }, + }, + "training": { + "batch_size": 2, + "optimizer": "adam", + "optimizer_params": {"learning_rate": 1e-3}, + }, +} + + +def _collect_tf_weights(layer, dest): + """Recursively collect leaf-layer TF weights keyed by (layer_name, weight_name).""" + has_sublayers = hasattr(layer, "layers") and layer.layers + if not has_sublayers and hasattr(layer, "weights"): + for weight in layer.weights: + # Weight names from top-level leaf layers are short (e.g. "kernel"), + # while nested leaf layers include the layer path + # (e.g. "classifier_dense_0/kernel"). + full_name = weight.name + if "/" in full_name: + layer_name, weight_name = full_name.rsplit("/", 1) + else: + layer_name = layer.name + weight_name = full_name + dest[(layer_name, weight_name)] = weight + if has_sublayers: + for sub in layer.layers: + _collect_tf_weights(sub, dest) + + +def _copy_tf_weights_to_pt(tf_model, pt_model): + """Copy weights from the TF model to the PyTorch model in-place.""" + # PyTorch MaskedConv1D is lazy: run one forward pass to materialize the real + # Conv1d with the correct in_channels before copying weights. + dummy_x = torch.zeros(2, 6, 50, dtype=torch.int64) + dummy_mask = torch.ones(2, 6, 50, dtype=torch.bool) + pt_model.eval() + with torch.no_grad(): + pt_model(dummy_x, mask=dummy_mask) + + tf_weights = {} + for layer in tf_model.layers: + _collect_tf_weights(layer, tf_weights) + + pt_state = pt_model.state_dict() + + def _get(name_pair): + weight = tf_weights.get(name_pair) + if weight is None: + raise ValueError(f"Missing TensorFlow weight for {name_pair}") + return torch.from_numpy(weight.numpy()) + + def _set(pt_key, tf_weight, transform=None): + if pt_key not in pt_state: + raise ValueError(f"Missing PyTorch state_dict key: {pt_key}") + tensor = _get(tf_weight) + if transform is not None: + tensor = transform(tensor) + pt_state[pt_key].copy_(tensor) + + # Embedding layer: direct copy. + _set("rep_model.embedding.embed.weight", ("embedding", "embeddings")) + + # MaskedConv1D: TF kernel shape (K, C_in, C_out) -> PyTorch (C_out, C_in, K). + _set( + "rep_model.blocks.0.conv.weight", + ("rep_masked_conv1d_0", "kernel"), + transform=lambda t: t.permute(2, 1, 0), + ) + _set("rep_model.blocks.0.conv.bias", ("rep_masked_conv1d_0", "bias")) + + # Classification dense head: TF kernel shape (C_in, C_out) -> PyTorch (C_out, C_in). + _set( + "classification_head.net.0.weight", + ("classifier_dense_0", "kernel"), + transform=lambda t: t.T, + ) + _set("classification_head.net.0.bias", ("classifier_dense_0", "bias")) + + +def test_forward_parity(): + tf_builder = TFBuilder(CONFIG) + tf_models = tf_builder.build_fragment_classifier() + tf_model = tf_models["jaeger_model"] + + pt_builder = PTBuilder(CONFIG) + pt_models = pt_builder.build_fragment_classifier() + pt_model = pt_models["jaeger_model"] + pt_model.eval() + + _copy_tf_weights_to_pt(tf_model, pt_model) + + # Avoid zeros so the TF model's automatic mask (x != 0) is all-ones and + # matches the all-ones PyTorch mask, sidestepping the fact that TF's + # GlobalAveragePooling2D does not consume masks while PyTorch's pooler does. + x = np.random.randint(1, 65, size=(2, 6, 50)).astype(np.int64) + tf_out = tf_model.predict({"translated": x}, verbose=0) + with torch.no_grad(): + pt_out = pt_model( + torch.from_numpy(x), mask=torch.ones(2, 6, 50, dtype=torch.bool) + ) + + np.testing.assert_allclose( + tf_out["prediction"], + pt_out["prediction"].numpy(), + atol=1e-4, + rtol=1e-4, + ) From a04bc625fe9c279c67b046964ea99a9e20866e77 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:29:58 +0200 Subject: [PATCH 20/64] feat(pytorch): add NumpyFullDataset and dataset builder --- src/jaeger/data/pytorch/__init__.py | 3 + src/jaeger/data/pytorch/builders.py | 68 +++++++++++++++++++ src/jaeger/data/pytorch/collate.py | 27 ++++++++ src/jaeger/data/pytorch/dataset_numpy.py | 40 +++++++++++ tests/unit/data/pytorch/test_dataset_numpy.py | 24 +++++++ 5 files changed, 162 insertions(+) create mode 100644 src/jaeger/data/pytorch/__init__.py create mode 100644 src/jaeger/data/pytorch/builders.py create mode 100644 src/jaeger/data/pytorch/collate.py create mode 100644 src/jaeger/data/pytorch/dataset_numpy.py create mode 100644 tests/unit/data/pytorch/test_dataset_numpy.py diff --git a/src/jaeger/data/pytorch/__init__.py b/src/jaeger/data/pytorch/__init__.py new file mode 100644 index 0000000..072a745 --- /dev/null +++ b/src/jaeger/data/pytorch/__init__.py @@ -0,0 +1,3 @@ +"""PyTorch data loading utilities for Jaeger.""" + +from __future__ import annotations diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py new file mode 100644 index 0000000..c52f988 --- /dev/null +++ b/src/jaeger/data/pytorch/builders.py @@ -0,0 +1,68 @@ +"""Dataset/DataLoader builders from Jaeger YAML training configs.""" + +from __future__ import annotations + +from typing import Any, Dict + +from torch.utils.data import DataLoader + +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset + + +def build_datasets( + config: Dict[str, Any], branch: str = "classifier" +) -> Dict[str, DataLoader]: + """Build training and validation DataLoaders for a config branch. + + Parameters + ---------- + config: + Parsed Jaeger training configuration. + branch: + Either ``classifier`` or ``reliability``. Determines which data + entries are read from ``config["training"]``. + + Returns + ------- + Dictionary mapping ``"train"`` and ``"validation"`` to + :class:`torch.utils.data.DataLoader` instances. + """ + model_cfg = config.get("model", {}) + train_cfg = config.get("training", {}) + string_cfg = model_cfg.get("string_processor", {}) + batch_size = train_cfg.get("batch_size", 32) + + data_key = ( + "fragment_classifier_data" + if branch == "classifier" + else "fragment_reliability_data" + ) + data_cfg = train_cfg.get(data_key, {}) + + datasets: Dict[str, NumpyFullDataset] = {} + for split in ["train", "validation"]: + entries = data_cfg.get(split, []) + paths = [p for entry in entries for p in entry.get("path", [])] + if not paths: + raise ValueError(f"No paths found for {branch}/{split}") + + data_format = string_cfg.get("data_format", "numpy_full") + if data_format == "numpy_full": + datasets[split] = NumpyFullDataset(paths[0]) + else: + raise ValueError(f"Unsupported data_format in Task 11: {data_format}") + + return { + "train": DataLoader( + datasets["train"], + batch_size=batch_size, + shuffle=True, + num_workers=0, + ), + "validation": DataLoader( + datasets["validation"], + batch_size=batch_size, + shuffle=False, + num_workers=0, + ), + } diff --git a/src/jaeger/data/pytorch/collate.py b/src/jaeger/data/pytorch/collate.py new file mode 100644 index 0000000..963f5e5 --- /dev/null +++ b/src/jaeger/data/pytorch/collate.py @@ -0,0 +1,27 @@ +"""Collators for variable-length PyTorch batches.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def pad_collate(batch: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]) -> tuple: + """Collate a list of ``(x, y, mask)`` tuples by padding to the max length. + + All tensors in the batch are stacked after padding the last dimension of + ``x`` and ``mask`` to ``max(x.shape[-1])``. + """ + xs, ys, masks = zip(*batch) + max_len = max(x.shape[-1] for x in xs) + padded_xs: list[torch.Tensor] = [] + padded_masks: list[torch.Tensor] = [] + for x, mask in zip(xs, masks): + pad = max_len - x.shape[-1] + if pad > 0: + padded_xs.append(F.pad(x, (0, pad))) + padded_masks.append(F.pad(mask, (0, pad))) + else: + padded_xs.append(x) + padded_masks.append(mask) + return torch.stack(padded_xs), torch.stack(ys), torch.stack(padded_masks) diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/data/pytorch/dataset_numpy.py new file mode 100644 index 0000000..100e6df --- /dev/null +++ b/src/jaeger/data/pytorch/dataset_numpy.py @@ -0,0 +1,40 @@ +"""PyTorch datasets backed by preprocessed NumPy arrays.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + + +class NumpyFullDataset(Dataset): + """Loads a fully-preprocessed ``.npz`` file. + + The archive is expected to contain at least an input array (e.g. + ``translated``) and a label array (``label``). Each sample is returned + together with a boolean mask indicating non-padding positions. + """ + + def __init__( + self, + path: str | Path, + input_key: str = "translated", + label_key: str = "label", + ): + data = np.load(path, allow_pickle=False) + self.inputs = torch.from_numpy(data[input_key]) + self.labels = torch.from_numpy(data[label_key]) + + def __len__(self) -> int: + return len(self.labels) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = self.inputs[idx] + if x.dim() == 3: + # (frames, length, channels) -> mask over frames and length + mask = (x != 0).any(dim=-1) + else: + mask = torch.ones(x.shape[-1], dtype=torch.bool) + return x, self.labels[idx], mask diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py new file mode 100644 index 0000000..60de353 --- /dev/null +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -0,0 +1,24 @@ +import numpy as np +import torch +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset + + +def test_numpy_full_dataset(tmp_path): + data = np.random.randint(0, 65, size=(10, 6, 50)).astype(np.int32) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] + path = tmp_path / "data.npz" + np.savez(path, translated=data, label=labels) + + ds = NumpyFullDataset(path, input_key="translated") + assert len(ds) == 10 + x, y, mask = ds[0] + assert x.shape == (6, 50) + assert y.shape == (3,) + assert mask.shape == (50,) + assert mask.dtype == torch.bool + + loader = torch.utils.data.DataLoader(ds, batch_size=2) + batch_x, batch_y, batch_mask = next(iter(loader)) + assert batch_x.shape == (2, 6, 50) + assert batch_y.shape == (2, 3) + assert batch_mask.shape == (2, 50) From 6f9ad13b3007c2053bdd5d72f64fa346d7f10b3b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:35:54 +0200 Subject: [PATCH 21/64] feat(pytorch): wire collator and add dataset builder tests --- src/jaeger/data/pytorch/builders.py | 18 ++++- tests/unit/data/pytorch/test_builders.py | 81 +++++++++++++++++++ tests/unit/data/pytorch/test_collate.py | 47 +++++++++++ tests/unit/data/pytorch/test_dataset_numpy.py | 15 ++++ 4 files changed, 158 insertions(+), 3 deletions(-) create mode 100644 tests/unit/data/pytorch/test_builders.py create mode 100644 tests/unit/data/pytorch/test_collate.py diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py index c52f988..cdfc823 100644 --- a/src/jaeger/data/pytorch/builders.py +++ b/src/jaeger/data/pytorch/builders.py @@ -4,8 +4,9 @@ from typing import Any, Dict -from torch.utils.data import DataLoader +from torch.utils.data import ConcatDataset, DataLoader, Dataset +from jaeger.data.pytorch.collate import pad_collate from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset @@ -14,6 +15,10 @@ def build_datasets( ) -> Dict[str, DataLoader]: """Build training and validation DataLoaders for a config branch. + The returned DataLoaders use :func:`pad_collate`, so they support + variable-length samples by padding each batch to the longest sequence + length in that batch. + Parameters ---------- config: @@ -39,7 +44,7 @@ def build_datasets( ) data_cfg = train_cfg.get(data_key, {}) - datasets: Dict[str, NumpyFullDataset] = {} + datasets: Dict[str, Dataset] = {} for split in ["train", "validation"]: entries = data_cfg.get(split, []) paths = [p for entry in entries for p in entry.get("path", [])] @@ -48,7 +53,12 @@ def build_datasets( data_format = string_cfg.get("data_format", "numpy_full") if data_format == "numpy_full": - datasets[split] = NumpyFullDataset(paths[0]) + split_datasets: list[Dataset] = [NumpyFullDataset(path) for path in paths] + datasets[split] = ( + split_datasets[0] + if len(split_datasets) == 1 + else ConcatDataset(split_datasets) + ) else: raise ValueError(f"Unsupported data_format in Task 11: {data_format}") @@ -58,11 +68,13 @@ def build_datasets( batch_size=batch_size, shuffle=True, num_workers=0, + collate_fn=pad_collate, ), "validation": DataLoader( datasets["validation"], batch_size=batch_size, shuffle=False, num_workers=0, + collate_fn=pad_collate, ), } diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/data/pytorch/test_builders.py new file mode 100644 index 0000000..7cd5e2b --- /dev/null +++ b/tests/unit/data/pytorch/test_builders.py @@ -0,0 +1,81 @@ +import numpy as np +import pytest + +from jaeger.data.pytorch.builders import build_datasets + + +def _make_npz(path, n_samples, length=50, channels=None): + if channels is None: + data = np.random.randint(0, 65, size=(n_samples, 6, length)).astype(np.int32) + else: + data = np.random.randint(0, 65, size=(n_samples, 6, length, channels)).astype( + np.int32 + ) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=n_samples)] + np.savez(path, translated=data, label=labels) + + +def _build_config(train_paths, val_paths, data_format="numpy_full", batch_size=4): + return { + "model": { + "string_processor": {"data_format": data_format}, + }, + "training": { + "batch_size": batch_size, + "fragment_classifier_data": { + "train": [{"path": [str(p) for p in train_paths]}], + "validation": [{"path": [str(p) for p in val_paths]}], + }, + }, + } + + +def test_build_datasets_numpy_full(tmp_path): + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_npz(train_path, n_samples=8) + _make_npz(val_path, n_samples=4) + + config = _build_config([train_path], [val_path], batch_size=4) + loaders = build_datasets(config, branch="classifier") + + assert set(loaders.keys()) == {"train", "validation"} + + batch_x, batch_y, batch_mask = next(iter(loaders["train"])) + assert batch_x.shape == (4, 6, 50) + assert batch_y.shape == (4, 3) + assert batch_mask.shape == (4, 50) + + batch_x, batch_y, batch_mask = next(iter(loaders["validation"])) + assert batch_x.shape == (4, 6, 50) + assert batch_y.shape == (4, 3) + assert batch_mask.shape == (4, 50) + + +def test_build_datasets_multiple_paths(tmp_path): + train_path1 = tmp_path / "train1.npz" + train_path2 = tmp_path / "train2.npz" + val_path = tmp_path / "val.npz" + _make_npz(train_path1, n_samples=3) + _make_npz(train_path2, n_samples=5) + _make_npz(val_path, n_samples=4) + + config = _build_config([train_path1, train_path2], [val_path], batch_size=8) + loaders = build_datasets(config, branch="classifier") + + assert len(loaders["train"].dataset) == 8 + batch_x, batch_y, batch_mask = next(iter(loaders["train"])) + assert batch_x.shape == (8, 6, 50) + assert batch_y.shape == (8, 3) + assert batch_mask.shape == (8, 50) + + +def test_build_datasets_unsupported_format(tmp_path): + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_npz(train_path, n_samples=4) + _make_npz(val_path, n_samples=4) + + config = _build_config([train_path], [val_path], data_format="csv") + with pytest.raises(ValueError, match="Unsupported data_format"): + build_datasets(config, branch="classifier") diff --git a/tests/unit/data/pytorch/test_collate.py b/tests/unit/data/pytorch/test_collate.py new file mode 100644 index 0000000..1f362f2 --- /dev/null +++ b/tests/unit/data/pytorch/test_collate.py @@ -0,0 +1,47 @@ +import torch + +from jaeger.data.pytorch.collate import pad_collate + + +def test_pad_collate_variable_length(): + x1 = torch.randn(6, 50) + y1 = torch.tensor([1.0, 0.0, 0.0]) + mask1 = torch.ones(50, dtype=torch.bool) + mask1[-10:] = False + + x2 = torch.randn(6, 30) + y2 = torch.tensor([0.0, 1.0, 0.0]) + mask2 = torch.ones(30, dtype=torch.bool) + + batch = [(x1, y1, mask1), (x2, y2, mask2)] + batch_x, batch_y, batch_mask = pad_collate(batch) + + assert batch_x.shape == (2, 6, 50) + assert batch_y.shape == (2, 3) + assert batch_mask.shape == (2, 50) + + # First sample should be unchanged except for stacking. + assert torch.equal(batch_x[0], x1) + assert torch.equal(batch_mask[0], mask1) + + # Second sample should be padded to length 50. + assert torch.equal(batch_x[1, :, :30], x2) + assert torch.equal(batch_x[1, :, 30:], torch.zeros(6, 20)) + assert torch.equal(batch_mask[1, :30], mask2) + assert not batch_mask[1, 30:].any() + + +def test_pad_collate_no_padding_needed(): + x1 = torch.randn(6, 40) + y1 = torch.tensor([1.0, 0.0, 0.0]) + mask1 = torch.ones(40, dtype=torch.bool) + x2 = torch.randn(6, 40) + y2 = torch.tensor([0.0, 1.0, 0.0]) + mask2 = torch.ones(40, dtype=torch.bool) + + batch = [(x1, y1, mask1), (x2, y2, mask2)] + batch_x, batch_y, batch_mask = pad_collate(batch) + + assert torch.equal(batch_x, torch.stack([x1, x2])) + assert torch.equal(batch_y, torch.stack([y1, y2])) + assert torch.equal(batch_mask, torch.stack([mask1, mask2])) diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py index 60de353..b73e069 100644 --- a/tests/unit/data/pytorch/test_dataset_numpy.py +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -22,3 +22,18 @@ def test_numpy_full_dataset(tmp_path): assert batch_x.shape == (2, 6, 50) assert batch_y.shape == (2, 3) assert batch_mask.shape == (2, 50) + + +def test_numpy_full_dataset_3d_input(tmp_path): + data = np.random.randint(0, 65, size=(5, 6, 50, 4)).astype(np.int32) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=5)] + path = tmp_path / "data3d.npz" + np.savez(path, translated=data, label=labels) + + ds = NumpyFullDataset(path, input_key="translated") + assert len(ds) == 5 + x, y, mask = ds[0] + assert x.shape == (6, 50, 4) + assert y.shape == (3,) + assert mask.shape == (6, 50) + assert mask.dtype == torch.bool From 2ece76ad3b1d0399816327f933163ff3e6585da4 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:43:31 +0200 Subject: [PATCH 22/64] feat(pytorch): fix dataset mask shapes and add model integration test --- src/jaeger/data/pytorch/__init__.py | 10 ++++ src/jaeger/data/pytorch/collate.py | 27 +++++++-- src/jaeger/data/pytorch/dataset_numpy.py | 9 ++- tests/unit/data/pytorch/test_builders.py | 56 ++++++++++++++++++- tests/unit/data/pytorch/test_collate.py | 44 ++++++++++++--- tests/unit/data/pytorch/test_dataset_numpy.py | 4 +- 6 files changed, 129 insertions(+), 21 deletions(-) diff --git a/src/jaeger/data/pytorch/__init__.py b/src/jaeger/data/pytorch/__init__.py index 072a745..4589889 100644 --- a/src/jaeger/data/pytorch/__init__.py +++ b/src/jaeger/data/pytorch/__init__.py @@ -1,3 +1,13 @@ """PyTorch data loading utilities for Jaeger.""" from __future__ import annotations + +from jaeger.data.pytorch.builders import build_datasets +from jaeger.data.pytorch.collate import pad_collate +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset + +__all__ = [ + "build_datasets", + "NumpyFullDataset", + "pad_collate", +] diff --git a/src/jaeger/data/pytorch/collate.py b/src/jaeger/data/pytorch/collate.py index 963f5e5..b4d6f98 100644 --- a/src/jaeger/data/pytorch/collate.py +++ b/src/jaeger/data/pytorch/collate.py @@ -9,17 +9,34 @@ def pad_collate(batch: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]) -> tuple: """Collate a list of ``(x, y, mask)`` tuples by padding to the max length. - All tensors in the batch are stacked after padding the last dimension of - ``x`` and ``mask`` to ``max(x.shape[-1])``. + Padding is applied to the sequence (length) dimension: + + - 2-D inputs ``(frames, length)`` are padded on the right of ``length``. + - 3-D inputs ``(frames, length, channels)`` are padded on the right of + ``length``; channels are left untouched. + + Masks are always ``(frames, length)`` and padded consistently. """ xs, ys, masks = zip(*batch) - max_len = max(x.shape[-1] for x in xs) + + def _seq_len(x: torch.Tensor) -> int: + if x.dim() == 2: + return x.shape[-1] + if x.dim() == 3: + return x.shape[-2] + raise ValueError(f"pad_collate supports 2-D and 3-D inputs, got rank {x.dim()}") + + max_len = max(_seq_len(x) for x in xs) padded_xs: list[torch.Tensor] = [] padded_masks: list[torch.Tensor] = [] for x, mask in zip(xs, masks): - pad = max_len - x.shape[-1] + pad = max_len - _seq_len(x) if pad > 0: - padded_xs.append(F.pad(x, (0, pad))) + if x.dim() == 2: + padded_xs.append(F.pad(x, (0, pad))) + else: + # Pad the length dimension (second-to-last) on the right. + padded_xs.append(F.pad(x, (0, 0, 0, pad))) padded_masks.append(F.pad(mask, (0, pad))) else: padded_xs.append(x) diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/data/pytorch/dataset_numpy.py index 100e6df..c3e9416 100644 --- a/src/jaeger/data/pytorch/dataset_numpy.py +++ b/src/jaeger/data/pytorch/dataset_numpy.py @@ -32,9 +32,14 @@ def __len__(self) -> int: def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: x = self.inputs[idx] - if x.dim() == 3: + if x.dim() == 2: + # (frames, length) + mask = x != 0 + elif x.dim() == 3: # (frames, length, channels) -> mask over frames and length mask = (x != 0).any(dim=-1) else: - mask = torch.ones(x.shape[-1], dtype=torch.bool) + raise ValueError( + f"NumpyFullDataset expects 2-D or 3-D inputs, got rank {x.dim()}" + ) return x, self.labels[idx], mask diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/data/pytorch/test_builders.py index 7cd5e2b..20bd9c4 100644 --- a/tests/unit/data/pytorch/test_builders.py +++ b/tests/unit/data/pytorch/test_builders.py @@ -1,7 +1,9 @@ import numpy as np import pytest +import torch from jaeger.data.pytorch.builders import build_datasets +from jaeger.nnlib.pytorch.builder import ModelBuilder def _make_npz(path, n_samples, length=50, channels=None): @@ -44,12 +46,12 @@ def test_build_datasets_numpy_full(tmp_path): batch_x, batch_y, batch_mask = next(iter(loaders["train"])) assert batch_x.shape == (4, 6, 50) assert batch_y.shape == (4, 3) - assert batch_mask.shape == (4, 50) + assert batch_mask.shape == (4, 6, 50) batch_x, batch_y, batch_mask = next(iter(loaders["validation"])) assert batch_x.shape == (4, 6, 50) assert batch_y.shape == (4, 3) - assert batch_mask.shape == (4, 50) + assert batch_mask.shape == (4, 6, 50) def test_build_datasets_multiple_paths(tmp_path): @@ -67,7 +69,7 @@ def test_build_datasets_multiple_paths(tmp_path): batch_x, batch_y, batch_mask = next(iter(loaders["train"])) assert batch_x.shape == (8, 6, 50) assert batch_y.shape == (8, 3) - assert batch_mask.shape == (8, 50) + assert batch_mask.shape == (8, 6, 50) def test_build_datasets_unsupported_format(tmp_path): @@ -79,3 +81,51 @@ def test_build_datasets_unsupported_format(tmp_path): config = _build_config([train_path], [val_path], data_format="csv") with pytest.raises(ValueError, match="Unsupported data_format"): build_datasets(config, branch="classifier") + + +def test_build_datasets_model_forward(tmp_path): + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_npz(train_path, n_samples=8) + _make_npz(val_path, n_samples=4) + + config = { + "model": { + "classifier_out_dim": 3, + "string_processor": {"data_format": "numpy_full"}, + "embedding": { + "input_type": "translated", + "vocab_size": 65, + "embedding_size": 4, + "use_embedding_layer": True, + }, + "representation_learner": { + "hidden_layers": [], + "pooling": "average", + }, + "classifier": { + "input_shape": 4, + "hidden_layers": [{"name": "dense", "config": {"units": 3}}], + }, + }, + "training": { + "batch_size": 4, + "fragment_classifier_data": { + "train": [{"path": [str(train_path)]}], + "validation": [{"path": [str(val_path)]}], + }, + }, + } + + loaders = build_datasets(config, branch="classifier") + batch_x, batch_y, batch_mask = next(iter(loaders["train"])) + + models = ModelBuilder(config).build_fragment_classifier() + model = models["jaeger_model"] + model.eval() + + with torch.no_grad(): + outputs = model(batch_x.long(), batch_mask) + + assert "prediction" in outputs + assert outputs["prediction"].shape == (4, 3) diff --git a/tests/unit/data/pytorch/test_collate.py b/tests/unit/data/pytorch/test_collate.py index 1f362f2..9ff2554 100644 --- a/tests/unit/data/pytorch/test_collate.py +++ b/tests/unit/data/pytorch/test_collate.py @@ -3,22 +3,22 @@ from jaeger.data.pytorch.collate import pad_collate -def test_pad_collate_variable_length(): +def test_pad_collate_variable_length_2d(): x1 = torch.randn(6, 50) y1 = torch.tensor([1.0, 0.0, 0.0]) - mask1 = torch.ones(50, dtype=torch.bool) - mask1[-10:] = False + mask1 = torch.ones(6, 50, dtype=torch.bool) + mask1[:, -10:] = False x2 = torch.randn(6, 30) y2 = torch.tensor([0.0, 1.0, 0.0]) - mask2 = torch.ones(30, dtype=torch.bool) + mask2 = torch.ones(6, 30, dtype=torch.bool) batch = [(x1, y1, mask1), (x2, y2, mask2)] batch_x, batch_y, batch_mask = pad_collate(batch) assert batch_x.shape == (2, 6, 50) assert batch_y.shape == (2, 3) - assert batch_mask.shape == (2, 50) + assert batch_mask.shape == (2, 6, 50) # First sample should be unchanged except for stacking. assert torch.equal(batch_x[0], x1) @@ -27,17 +27,43 @@ def test_pad_collate_variable_length(): # Second sample should be padded to length 50. assert torch.equal(batch_x[1, :, :30], x2) assert torch.equal(batch_x[1, :, 30:], torch.zeros(6, 20)) - assert torch.equal(batch_mask[1, :30], mask2) - assert not batch_mask[1, 30:].any() + assert torch.equal(batch_mask[1, :, :30], mask2) + assert not batch_mask[1, :, 30:].any() + + +def test_pad_collate_variable_length_3d(): + x1 = torch.randn(6, 50, 4) + y1 = torch.tensor([1.0, 0.0, 0.0]) + mask1 = torch.ones(6, 50, dtype=torch.bool) + mask1[:, -5:] = False + + x2 = torch.randn(6, 30, 4) + y2 = torch.tensor([0.0, 1.0, 0.0]) + mask2 = torch.ones(6, 30, dtype=torch.bool) + + batch = [(x1, y1, mask1), (x2, y2, mask2)] + batch_x, batch_y, batch_mask = pad_collate(batch) + + assert batch_x.shape == (2, 6, 50, 4) + assert batch_y.shape == (2, 3) + assert batch_mask.shape == (2, 6, 50) + + assert torch.equal(batch_x[0], x1) + assert torch.equal(batch_mask[0], mask1) + + assert torch.equal(batch_x[1, :, :30, :], x2) + assert torch.equal(batch_x[1, :, 30:, :], torch.zeros(6, 20, 4)) + assert torch.equal(batch_mask[1, :, :30], mask2) + assert not batch_mask[1, :, 30:].any() def test_pad_collate_no_padding_needed(): x1 = torch.randn(6, 40) y1 = torch.tensor([1.0, 0.0, 0.0]) - mask1 = torch.ones(40, dtype=torch.bool) + mask1 = torch.ones(6, 40, dtype=torch.bool) x2 = torch.randn(6, 40) y2 = torch.tensor([0.0, 1.0, 0.0]) - mask2 = torch.ones(40, dtype=torch.bool) + mask2 = torch.ones(6, 40, dtype=torch.bool) batch = [(x1, y1, mask1), (x2, y2, mask2)] batch_x, batch_y, batch_mask = pad_collate(batch) diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py index b73e069..95ae2fe 100644 --- a/tests/unit/data/pytorch/test_dataset_numpy.py +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -14,14 +14,14 @@ def test_numpy_full_dataset(tmp_path): x, y, mask = ds[0] assert x.shape == (6, 50) assert y.shape == (3,) - assert mask.shape == (50,) + assert mask.shape == (6, 50) assert mask.dtype == torch.bool loader = torch.utils.data.DataLoader(ds, batch_size=2) batch_x, batch_y, batch_mask = next(iter(loader)) assert batch_x.shape == (2, 6, 50) assert batch_y.shape == (2, 3) - assert batch_mask.shape == (2, 50) + assert batch_mask.shape == (2, 6, 50) def test_numpy_full_dataset_3d_input(tmp_path): From 0bd5a6460e23c2bd8ac6623e4d816e5d086540e4 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:48:29 +0200 Subject: [PATCH 23/64] feat(pytorch): add NumpyRawDataset and transforms --- src/jaeger/data/pytorch/__init__.py | 3 +- src/jaeger/data/pytorch/builders.py | 28 ++++++++- src/jaeger/data/pytorch/dataset_numpy.py | 59 +++++++++++++++++++ src/jaeger/data/pytorch/transforms.py | 53 +++++++++++++++++ tests/unit/data/pytorch/test_dataset_numpy.py | 26 +++++++- 5 files changed, 166 insertions(+), 3 deletions(-) create mode 100644 src/jaeger/data/pytorch/transforms.py diff --git a/src/jaeger/data/pytorch/__init__.py b/src/jaeger/data/pytorch/__init__.py index 4589889..fcd1ae7 100644 --- a/src/jaeger/data/pytorch/__init__.py +++ b/src/jaeger/data/pytorch/__init__.py @@ -4,10 +4,11 @@ from jaeger.data.pytorch.builders import build_datasets from jaeger.data.pytorch.collate import pad_collate -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset __all__ = [ "build_datasets", "NumpyFullDataset", + "NumpyRawDataset", "pad_collate", ] diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py index cdfc823..2a72ea5 100644 --- a/src/jaeger/data/pytorch/builders.py +++ b/src/jaeger/data/pytorch/builders.py @@ -7,7 +7,8 @@ from torch.utils.data import ConcatDataset, DataLoader, Dataset from jaeger.data.pytorch.collate import pad_collate -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset +from jaeger.seqops.maps import CODON_ID def build_datasets( @@ -36,6 +37,10 @@ def build_datasets( train_cfg = config.get("training", {}) string_cfg = model_cfg.get("string_processor", {}) batch_size = train_cfg.get("batch_size", 32) + out_dim_key = ( + "classifier_out_dim" if branch == "classifier" else "reliability_out_dim" + ) + num_classes = model_cfg.get(out_dim_key, 3) data_key = ( "fragment_classifier_data" @@ -44,6 +49,8 @@ def build_datasets( ) data_cfg = train_cfg.get(data_key, {}) + codon_table = {c: i for i, c in enumerate(CODON_ID)} + datasets: Dict[str, Dataset] = {} for split in ["train", "validation"]: entries = data_cfg.get(split, []) @@ -59,6 +66,25 @@ def build_datasets( if len(split_datasets) == 1 else ConcatDataset(split_datasets) ) + elif data_format == "numpy_raw": + split_datasets = [ + NumpyRawDataset( + path, + crop_size=string_cfg.get("crop_size", 500), + num_classes=num_classes, + codon_table=codon_table, + shuffle=string_cfg.get("shuffle", True), + mutate=string_cfg.get("mutate", False), + mutation_rate=string_cfg.get("mutation_rate", 0.1), + shuffle_frames=string_cfg.get("shuffle_frames", False), + ) + for path in paths + ] + datasets[split] = ( + split_datasets[0] + if len(split_datasets) == 1 + else ConcatDataset(split_datasets) + ) else: raise ValueError(f"Unsupported data_format in Task 11: {data_format}") diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/data/pytorch/dataset_numpy.py index c3e9416..bad8dd2 100644 --- a/src/jaeger/data/pytorch/dataset_numpy.py +++ b/src/jaeger/data/pytorch/dataset_numpy.py @@ -2,12 +2,20 @@ from __future__ import annotations +import random from pathlib import Path +from typing import Dict, Optional import numpy as np import torch from torch.utils.data import Dataset +from jaeger.data.pytorch.transforms import ( + apply_mutation, + shuffle_frames as shuffle_frames_fn, + translate_to_codons, +) + class NumpyFullDataset(Dataset): """Loads a fully-preprocessed ``.npz`` file. @@ -43,3 +51,54 @@ def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tenso f"NumpyFullDataset expects 2-D or 3-D inputs, got rank {x.dim()}" ) return x, self.labels[idx], mask + + +class NumpyRawDataset(Dataset): + """Loads raw int8 sequences and applies runtime preprocessing.""" + + def __init__( + self, + path: str | Path, + seq_key: str = "sequences", + label_key: str = "labels", + crop_size: int = 500, + num_classes: int = 3, + codon_table: Optional[Dict[str, int]] = None, + shuffle: bool = True, + mutate: bool = False, + mutation_rate: float = 0.1, + shuffle_frames: bool = False, + ): + data = np.load(path, allow_pickle=False) + self.seqs = data[seq_key] + self.labels = torch.from_numpy(data[label_key]) + self.crop_size = crop_size + self.num_classes = num_classes + self.codon_table = codon_table + self.shuffle = shuffle + self.mutate = mutate + self.mutation_rate = mutation_rate + self.shuffle_frames = shuffle_frames + + def __len__(self) -> int: + return len(self.labels) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + seq = self.seqs[idx] + if self.mutate: + seq = apply_mutation(seq, self.mutation_rate) + x = translate_to_codons(seq, self.codon_table) + length = x.shape[1] + if length > self.crop_size: + start = random.randint(0, length - self.crop_size) + x = x[:, start : start + self.crop_size] + elif length < self.crop_size: + pad = self.crop_size - length + x = torch.nn.functional.pad(x, (0, pad)) + mask = x != 0 + if self.shuffle_frames: + x = shuffle_frames_fn(x) + y = torch.nn.functional.one_hot( + self.labels[idx].long(), num_classes=self.num_classes + ).float() + return x, y, mask diff --git a/src/jaeger/data/pytorch/transforms.py b/src/jaeger/data/pytorch/transforms.py new file mode 100644 index 0000000..531fb65 --- /dev/null +++ b/src/jaeger/data/pytorch/transforms.py @@ -0,0 +1,53 @@ +"""Runtime preprocessing transforms for PyTorch datasets.""" + +from __future__ import annotations + +import random +from typing import Dict + +import numpy as np +import torch + + +_COMPLEMENT = {"A": "T", "T": "A", "G": "C", "C": "G"} +_NUCLEOTIDES = ["A", "T", "G", "C"] + + +def _reverse_complement(seq_str: str) -> str: + return "".join(_COMPLEMENT.get(base, "N") for base in reversed(seq_str.upper())) + + +def translate_to_codons(seq: np.ndarray, codon_table: Dict[str, int]) -> torch.Tensor: + """Translate an int8 nucleotide sequence into 6-frame codon indices.""" + seq_str = "".join(_NUCLEOTIDES[i] for i in seq) + frames = [] + for offset in range(3): + codons = [seq_str[i : i + 3] for i in range(offset, len(seq_str) - 2, 3)] + indices = [codon_table.get(c, 0) for c in codons] + rev = _reverse_complement(seq_str) + rev_codons = [rev[i : i + 3] for i in range(offset, len(rev) - 2, 3)] + rev_indices = [codon_table.get(c, 0) for c in rev_codons] + frames.append(torch.tensor(indices, dtype=torch.long)) + frames.append(torch.tensor(rev_indices, dtype=torch.long)) + return torch.stack(frames) + + +def dna_to_indices(seq: str) -> np.ndarray: + mapping = {"A": 0, "T": 1, "G": 2, "C": 3} + return np.array([mapping.get(base.upper(), 0) for base in seq], dtype=np.int8) + + +def apply_mutation(seq: np.ndarray, rate: float) -> np.ndarray: + if rate == 0: + return seq + mask = np.random.rand(len(seq)) < rate + mutated = seq.copy() + mutated[mask] = np.random.randint(0, 4, size=mask.sum()) + return mutated + + +def shuffle_frames(x: torch.Tensor) -> torch.Tensor: + """Shuffle the 6 reading frames of a translated tensor.""" + frames = list(range(x.shape[0])) + random.shuffle(frames) + return x[frames] diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py index 95ae2fe..c993827 100644 --- a/tests/unit/data/pytorch/test_dataset_numpy.py +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -1,6 +1,7 @@ import numpy as np import torch -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset +from jaeger.seqops.maps import CODON_ID def test_numpy_full_dataset(tmp_path): @@ -37,3 +38,26 @@ def test_numpy_full_dataset_3d_input(tmp_path): assert y.shape == (3,) assert mask.shape == (6, 50) assert mask.dtype == torch.bool + + +def test_numpy_raw_dataset(tmp_path): + seqs = np.random.randint(0, 4, size=(10, 500)).astype(np.int8) + labels = np.random.randint(0, 3, size=10).astype(np.int64) + path = tmp_path / "raw.npz" + np.savez(path, sequences=seqs, labels=labels) + + codon_table = {c: i for i, c in enumerate(CODON_ID)} + ds = NumpyRawDataset( + path, + seq_key="sequences", + label_key="labels", + crop_size=50, + num_classes=3, + codon_table=codon_table, + ) + x, y, mask = ds[0] + assert x.shape == (6, 50) + assert mask.shape == (6, 50) + assert y.shape == (3,) + assert mask.dtype == torch.bool + assert y.dtype == torch.float32 From 4bcd01e9228227b34af2450cd184c58be3d49f15 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 00:54:24 +0200 Subject: [PATCH 24/64] feat(pytorch): fix NumpyRawDataset labels, codon padding, and shuffle --- src/jaeger/data/pytorch/builders.py | 4 +-- src/jaeger/data/pytorch/dataset_numpy.py | 18 +++++++--- src/jaeger/data/pytorch/transforms.py | 4 +-- tests/unit/data/pytorch/test_builders.py | 33 +++++++++++++++++++ tests/unit/data/pytorch/test_dataset_numpy.py | 26 ++++++++++++++- 5 files changed, 75 insertions(+), 10 deletions(-) diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py index 2a72ea5..4b7b670 100644 --- a/src/jaeger/data/pytorch/builders.py +++ b/src/jaeger/data/pytorch/builders.py @@ -49,7 +49,7 @@ def build_datasets( ) data_cfg = train_cfg.get(data_key, {}) - codon_table = {c: i for i, c in enumerate(CODON_ID)} + codon_table = {c: i + 1 for i, c in enumerate(CODON_ID)} datasets: Dict[str, Dataset] = {} for split in ["train", "validation"]: @@ -86,7 +86,7 @@ def build_datasets( else ConcatDataset(split_datasets) ) else: - raise ValueError(f"Unsupported data_format in Task 11: {data_format}") + raise ValueError(f"Unsupported data_format: {data_format}") return { "train": DataLoader( diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/data/pytorch/dataset_numpy.py index bad8dd2..8fb8487 100644 --- a/src/jaeger/data/pytorch/dataset_numpy.py +++ b/src/jaeger/data/pytorch/dataset_numpy.py @@ -71,7 +71,8 @@ def __init__( ): data = np.load(path, allow_pickle=False) self.seqs = data[seq_key] - self.labels = torch.from_numpy(data[label_key]) + raw_labels = torch.from_numpy(data[label_key]) + self.labels = self._normalize_labels(raw_labels, num_classes) self.crop_size = crop_size self.num_classes = num_classes self.codon_table = codon_table @@ -80,11 +81,21 @@ def __init__( self.mutation_rate = mutation_rate self.shuffle_frames = shuffle_frames + @staticmethod + def _normalize_labels(labels: torch.Tensor, num_classes: int) -> torch.Tensor: + """Convert class-index labels to one-hot; leave one-hot labels unchanged.""" + if labels.dim() == 1 or (labels.dim() == 2 and labels.shape[1] == 1): + indices = labels.view(-1).long() + return torch.nn.functional.one_hot(indices, num_classes=num_classes).float() + return labels.float() + def __len__(self) -> int: return len(self.labels) def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: seq = self.seqs[idx] + if self.shuffle: + seq = seq[np.random.permutation(len(seq))] if self.mutate: seq = apply_mutation(seq, self.mutation_rate) x = translate_to_codons(seq, self.codon_table) @@ -98,7 +109,4 @@ def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tenso mask = x != 0 if self.shuffle_frames: x = shuffle_frames_fn(x) - y = torch.nn.functional.one_hot( - self.labels[idx].long(), num_classes=self.num_classes - ).float() - return x, y, mask + return x, self.labels[idx], mask diff --git a/src/jaeger/data/pytorch/transforms.py b/src/jaeger/data/pytorch/transforms.py index 531fb65..d7c4b31 100644 --- a/src/jaeger/data/pytorch/transforms.py +++ b/src/jaeger/data/pytorch/transforms.py @@ -23,10 +23,10 @@ def translate_to_codons(seq: np.ndarray, codon_table: Dict[str, int]) -> torch.T frames = [] for offset in range(3): codons = [seq_str[i : i + 3] for i in range(offset, len(seq_str) - 2, 3)] - indices = [codon_table.get(c, 0) for c in codons] + indices = [codon_table.get(c, 1) for c in codons] rev = _reverse_complement(seq_str) rev_codons = [rev[i : i + 3] for i in range(offset, len(rev) - 2, 3)] - rev_indices = [codon_table.get(c, 0) for c in rev_codons] + rev_indices = [codon_table.get(c, 1) for c in rev_codons] frames.append(torch.tensor(indices, dtype=torch.long)) frames.append(torch.tensor(rev_indices, dtype=torch.long)) return torch.stack(frames) diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/data/pytorch/test_builders.py index 20bd9c4..9b612ab 100644 --- a/tests/unit/data/pytorch/test_builders.py +++ b/tests/unit/data/pytorch/test_builders.py @@ -17,6 +17,12 @@ def _make_npz(path, n_samples, length=50, channels=None): np.savez(path, translated=data, label=labels) +def _make_raw_npz(path, n_samples, seq_length=500): + seqs = np.random.randint(0, 4, size=(n_samples, seq_length)).astype(np.int8) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=n_samples)] + np.savez(path, sequences=seqs, labels=labels) + + def _build_config(train_paths, val_paths, data_format="numpy_full", batch_size=4): return { "model": { @@ -54,6 +60,33 @@ def test_build_datasets_numpy_full(tmp_path): assert batch_mask.shape == (4, 6, 50) +def test_build_datasets_numpy_raw(tmp_path): + train_path = tmp_path / "train_raw.npz" + val_path = tmp_path / "val_raw.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=500) + _make_raw_npz(val_path, n_samples=4, seq_length=500) + + config = _build_config( + [train_path], [val_path], data_format="numpy_raw", batch_size=4 + ) + config["model"]["string_processor"]["crop_size"] = 50 + loaders = build_datasets(config, branch="classifier") + + assert set(loaders.keys()) == {"train", "validation"} + + batch_x, batch_y, batch_mask = next(iter(loaders["train"])) + assert batch_x.shape == (4, 6, 50) + assert batch_y.shape == (4, 3) + assert torch.allclose(batch_y.sum(dim=1), torch.ones(4)) + assert batch_mask.shape == (4, 6, 50) + + batch_x, batch_y, batch_mask = next(iter(loaders["validation"])) + assert batch_x.shape == (4, 6, 50) + assert batch_y.shape == (4, 3) + assert torch.allclose(batch_y.sum(dim=1), torch.ones(4)) + assert batch_mask.shape == (4, 6, 50) + + def test_build_datasets_multiple_paths(tmp_path): train_path1 = tmp_path / "train1.npz" train_path2 = tmp_path / "train2.npz" diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py index c993827..668f605 100644 --- a/tests/unit/data/pytorch/test_dataset_numpy.py +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -46,7 +46,7 @@ def test_numpy_raw_dataset(tmp_path): path = tmp_path / "raw.npz" np.savez(path, sequences=seqs, labels=labels) - codon_table = {c: i for i, c in enumerate(CODON_ID)} + codon_table = {c: i + 1 for i, c in enumerate(CODON_ID)} ds = NumpyRawDataset( path, seq_key="sequences", @@ -61,3 +61,27 @@ def test_numpy_raw_dataset(tmp_path): assert y.shape == (3,) assert mask.dtype == torch.bool assert y.dtype == torch.float32 + + +def test_numpy_raw_dataset_one_hot_labels(tmp_path): + seqs = np.random.randint(0, 4, size=(10, 500)).astype(np.int8) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] + path = tmp_path / "raw_onehot.npz" + np.savez(path, sequences=seqs, labels=labels) + + codon_table = {c: i + 1 for i, c in enumerate(CODON_ID)} + ds = NumpyRawDataset( + path, + seq_key="sequences", + label_key="labels", + crop_size=50, + num_classes=3, + codon_table=codon_table, + ) + x, y, mask = ds[0] + assert x.shape == (6, 50) + assert mask.shape == (6, 50) + assert y.shape == (3,) + assert torch.allclose(y, torch.tensor(labels[0])) + assert mask.dtype == torch.bool + assert y.dtype == torch.float32 From 293f0393c7f6549ef49272f2fa45d63c70f04910 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 01:05:22 +0200 Subject: [PATCH 25/64] feat(pytorch): fix codon table and raw crop semantics --- src/jaeger/data/pytorch/builders.py | 4 +- src/jaeger/data/pytorch/dataset_numpy.py | 20 +++++--- src/jaeger/data/pytorch/transforms.py | 22 +++++--- tests/unit/data/pytorch/test_builders.py | 19 ++++--- tests/unit/data/pytorch/test_dataset_numpy.py | 50 +++++++++++++++---- 5 files changed, 82 insertions(+), 33 deletions(-) diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py index 4b7b670..f3f0fe1 100644 --- a/src/jaeger/data/pytorch/builders.py +++ b/src/jaeger/data/pytorch/builders.py @@ -8,7 +8,7 @@ from jaeger.data.pytorch.collate import pad_collate from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset -from jaeger.seqops.maps import CODON_ID +from jaeger.seqops.maps import CODONS def build_datasets( @@ -49,7 +49,7 @@ def build_datasets( ) data_cfg = train_cfg.get(data_key, {}) - codon_table = {c: i + 1 for i, c in enumerate(CODON_ID)} + codon_table = {c: i + 1 for i, c in enumerate(CODONS)} datasets: Dict[str, Dataset] = {} for split in ["train", "validation"]: diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/data/pytorch/dataset_numpy.py index 8fb8487..5ba0c57 100644 --- a/src/jaeger/data/pytorch/dataset_numpy.py +++ b/src/jaeger/data/pytorch/dataset_numpy.py @@ -92,20 +92,26 @@ def _normalize_labels(labels: torch.Tensor, num_classes: int) -> torch.Tensor: def __len__(self) -> int: return len(self.labels) + def _crop_pad_sequence(self, seq: np.ndarray) -> np.ndarray: + """Crop or pad a raw nucleotide sequence to ``crop_size``.""" + seq = np.asarray(seq) + length = len(seq) + if length > self.crop_size: + start = random.randint(0, length - self.crop_size) + return seq[start : start + self.crop_size] + if length < self.crop_size: + pad = self.crop_size - length + return np.pad(seq, (0, pad), constant_values=4) + return seq + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: seq = self.seqs[idx] if self.shuffle: seq = seq[np.random.permutation(len(seq))] + seq = self._crop_pad_sequence(seq) if self.mutate: seq = apply_mutation(seq, self.mutation_rate) x = translate_to_codons(seq, self.codon_table) - length = x.shape[1] - if length > self.crop_size: - start = random.randint(0, length - self.crop_size) - x = x[:, start : start + self.crop_size] - elif length < self.crop_size: - pad = self.crop_size - length - x = torch.nn.functional.pad(x, (0, pad)) mask = x != 0 if self.shuffle_frames: x = shuffle_frames_fn(x) diff --git a/src/jaeger/data/pytorch/transforms.py b/src/jaeger/data/pytorch/transforms.py index d7c4b31..2f3af82 100644 --- a/src/jaeger/data/pytorch/transforms.py +++ b/src/jaeger/data/pytorch/transforms.py @@ -10,7 +10,7 @@ _COMPLEMENT = {"A": "T", "T": "A", "G": "C", "C": "G"} -_NUCLEOTIDES = ["A", "T", "G", "C"] +_NUCLEOTIDES = ["A", "T", "G", "C", "N"] def _reverse_complement(seq_str: str) -> str: @@ -18,23 +18,33 @@ def _reverse_complement(seq_str: str) -> str: def translate_to_codons(seq: np.ndarray, codon_table: Dict[str, int]) -> torch.Tensor: - """Translate an int8 nucleotide sequence into 6-frame codon indices.""" + """Translate an int8 nucleotide sequence into 6-frame codon indices. + + Valid codons are mapped using ``codon_table``; unknown or padded codons + map to ``0`` so that masks can be computed as ``x != 0``. The six reading + frames are trimmed to the same length before stacking. + """ seq_str = "".join(_NUCLEOTIDES[i] for i in seq) frames = [] for offset in range(3): codons = [seq_str[i : i + 3] for i in range(offset, len(seq_str) - 2, 3)] - indices = [codon_table.get(c, 1) for c in codons] + indices = [codon_table.get(c, 0) for c in codons] rev = _reverse_complement(seq_str) rev_codons = [rev[i : i + 3] for i in range(offset, len(rev) - 2, 3)] - rev_indices = [codon_table.get(c, 1) for c in rev_codons] + rev_indices = [codon_table.get(c, 0) for c in rev_codons] frames.append(torch.tensor(indices, dtype=torch.long)) frames.append(torch.tensor(rev_indices, dtype=torch.long)) + + max_len = max(len(f) for f in frames) + frames = [ + torch.nn.functional.pad(f, (0, max_len - len(f))) for f in frames + ] return torch.stack(frames) def dna_to_indices(seq: str) -> np.ndarray: - mapping = {"A": 0, "T": 1, "G": 2, "C": 3} - return np.array([mapping.get(base.upper(), 0) for base in seq], dtype=np.int8) + mapping = {"A": 0, "T": 1, "G": 2, "C": 3, "N": 4} + return np.array([mapping.get(base.upper(), 4) for base in seq], dtype=np.int8) def apply_mutation(seq: np.ndarray, rate: float) -> np.ndarray: diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/data/pytorch/test_builders.py index 9b612ab..14cf13f 100644 --- a/tests/unit/data/pytorch/test_builders.py +++ b/tests/unit/data/pytorch/test_builders.py @@ -61,30 +61,35 @@ def test_build_datasets_numpy_full(tmp_path): def test_build_datasets_numpy_raw(tmp_path): + crop_size = 50 train_path = tmp_path / "train_raw.npz" val_path = tmp_path / "val_raw.npz" - _make_raw_npz(train_path, n_samples=8, seq_length=500) - _make_raw_npz(val_path, n_samples=4, seq_length=500) + _make_raw_npz(train_path, n_samples=8, seq_length=crop_size) + _make_raw_npz(val_path, n_samples=4, seq_length=crop_size) config = _build_config( [train_path], [val_path], data_format="numpy_raw", batch_size=4 ) - config["model"]["string_processor"]["crop_size"] = 50 + config["model"]["string_processor"]["crop_size"] = crop_size loaders = build_datasets(config, branch="classifier") assert set(loaders.keys()) == {"train", "validation"} batch_x, batch_y, batch_mask = next(iter(loaders["train"])) - assert batch_x.shape == (4, 6, 50) + assert batch_x.shape[0] == 4 + assert batch_x.shape[1] == 6 assert batch_y.shape == (4, 3) assert torch.allclose(batch_y.sum(dim=1), torch.ones(4)) - assert batch_mask.shape == (4, 6, 50) + assert batch_mask.shape == batch_x.shape + assert batch_x.max() <= 64 + assert batch_x.min() >= 0 batch_x, batch_y, batch_mask = next(iter(loaders["validation"])) - assert batch_x.shape == (4, 6, 50) + assert batch_x.shape[0] == 4 + assert batch_x.shape[1] == 6 assert batch_y.shape == (4, 3) assert torch.allclose(batch_y.sum(dim=1), torch.ones(4)) - assert batch_mask.shape == (4, 6, 50) + assert batch_mask.shape == batch_x.shape def test_build_datasets_multiple_paths(tmp_path): diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py index 668f605..5e6550f 100644 --- a/tests/unit/data/pytorch/test_dataset_numpy.py +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -1,7 +1,8 @@ import numpy as np import torch from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset -from jaeger.seqops.maps import CODON_ID +from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons +from jaeger.seqops.maps import CODONS def test_numpy_full_dataset(tmp_path): @@ -41,47 +42,74 @@ def test_numpy_full_dataset_3d_input(tmp_path): def test_numpy_raw_dataset(tmp_path): - seqs = np.random.randint(0, 4, size=(10, 500)).astype(np.int8) + crop_size = 50 + seqs = np.random.randint(0, 4, size=(10, crop_size)).astype(np.int8) labels = np.random.randint(0, 3, size=10).astype(np.int64) path = tmp_path / "raw.npz" np.savez(path, sequences=seqs, labels=labels) - codon_table = {c: i + 1 for i, c in enumerate(CODON_ID)} + codon_table = {c: i + 1 for i, c in enumerate(CODONS)} ds = NumpyRawDataset( path, seq_key="sequences", label_key="labels", - crop_size=50, + crop_size=crop_size, num_classes=3, codon_table=codon_table, ) x, y, mask = ds[0] - assert x.shape == (6, 50) - assert mask.shape == (6, 50) + assert x.shape[0] == 6 + assert x.shape == mask.shape assert y.shape == (3,) assert mask.dtype == torch.bool assert y.dtype == torch.float32 + assert x.max() <= 64 + assert x.min() >= 0 def test_numpy_raw_dataset_one_hot_labels(tmp_path): - seqs = np.random.randint(0, 4, size=(10, 500)).astype(np.int8) + crop_size = 50 + seqs = np.random.randint(0, 4, size=(10, crop_size)).astype(np.int8) labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] path = tmp_path / "raw_onehot.npz" np.savez(path, sequences=seqs, labels=labels) - codon_table = {c: i + 1 for i, c in enumerate(CODON_ID)} + codon_table = {c: i + 1 for i, c in enumerate(CODONS)} ds = NumpyRawDataset( path, seq_key="sequences", label_key="labels", - crop_size=50, + crop_size=crop_size, num_classes=3, codon_table=codon_table, ) x, y, mask = ds[0] - assert x.shape == (6, 50) - assert mask.shape == (6, 50) + assert x.shape[0] == 6 + assert x.shape == mask.shape assert y.shape == (3,) assert torch.allclose(y, torch.tensor(labels[0])) assert mask.dtype == torch.bool assert y.dtype == torch.float32 + + +def test_translate_to_codons_known_sequence(): + codon_table = {c: i + 1 for i, c in enumerate(CODONS)} + seq = "ATGAAATTTCCC" + x = translate_to_codons(dna_to_indices(seq), codon_table) + + assert x.shape[0] == 6 + assert x.shape[1] == 4 + + expected = [codon_table["ATG"], codon_table["AAA"], codon_table["TTT"], codon_table["CCC"]] + assert torch.equal(x[0], torch.tensor(expected, dtype=torch.long)) + + +def test_translate_to_codons_pads_unknown_to_zero(): + codon_table = {c: i + 1 for i, c in enumerate(CODONS)} + seq = "ATGAAANNNCCC" + x = translate_to_codons(dna_to_indices(seq), codon_table) + + # The codon containing N should map to 0. + mask = x != 0 + assert not mask.all() + assert (x[~mask] == 0).all() From 8253c3c03b91ff89168d6aa6955f76b8616dcb24 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:22:50 +0200 Subject: [PATCH 26/64] feat(pytorch): add CSVDataset --- src/jaeger/data/pytorch/__init__.py | 2 + src/jaeger/data/pytorch/builders.py | 21 +++++ src/jaeger/data/pytorch/dataset_csv.py | 87 +++++++++++++++++++++ tests/unit/data/pytorch/test_builders.py | 2 +- tests/unit/data/pytorch/test_dataset_csv.py | 22 ++++++ 5 files changed, 133 insertions(+), 1 deletion(-) create mode 100644 src/jaeger/data/pytorch/dataset_csv.py create mode 100644 tests/unit/data/pytorch/test_dataset_csv.py diff --git a/src/jaeger/data/pytorch/__init__.py b/src/jaeger/data/pytorch/__init__.py index fcd1ae7..8797f46 100644 --- a/src/jaeger/data/pytorch/__init__.py +++ b/src/jaeger/data/pytorch/__init__.py @@ -4,10 +4,12 @@ from jaeger.data.pytorch.builders import build_datasets from jaeger.data.pytorch.collate import pad_collate +from jaeger.data.pytorch.dataset_csv import CSVDataset from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset __all__ = [ "build_datasets", + "CSVDataset", "NumpyFullDataset", "NumpyRawDataset", "pad_collate", diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py index f3f0fe1..5129836 100644 --- a/src/jaeger/data/pytorch/builders.py +++ b/src/jaeger/data/pytorch/builders.py @@ -7,6 +7,7 @@ from torch.utils.data import ConcatDataset, DataLoader, Dataset from jaeger.data.pytorch.collate import pad_collate +from jaeger.data.pytorch.dataset_csv import CSVDataset from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset from jaeger.seqops.maps import CODONS @@ -85,6 +86,26 @@ def build_datasets( if len(split_datasets) == 1 else ConcatDataset(split_datasets) ) + elif data_format == "csv": + split_datasets = [ + CSVDataset( + path, + crop_size=string_cfg.get("crop_size", 500), + num_classes=num_classes, + codon_table=codon_table, + shuffle=string_cfg.get("shuffle", False), + mutate=string_cfg.get("mutate", False), + mutation_rate=string_cfg.get("mutation_rate", 0.1), + shuffle_frames=string_cfg.get("shuffle_frames", False), + label_first=True, + ) + for path in paths + ] + datasets[split] = ( + split_datasets[0] + if len(split_datasets) == 1 + else ConcatDataset(split_datasets) + ) else: raise ValueError(f"Unsupported data_format: {data_format}") diff --git a/src/jaeger/data/pytorch/dataset_csv.py b/src/jaeger/data/pytorch/dataset_csv.py new file mode 100644 index 0000000..b91365f --- /dev/null +++ b/src/jaeger/data/pytorch/dataset_csv.py @@ -0,0 +1,87 @@ +"""PyTorch dataset backed by raw CSV files.""" + +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Dict, Optional + +import numpy as np +import torch +from torch.utils.data import Dataset + +from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons + + +class CSVDataset(Dataset): + """Reads label,sequence CSV files and applies runtime preprocessing.""" + + def __init__( + self, + path: str | Path, + crop_size: int = 500, + num_classes: int = 3, + codon_table: Optional[Dict[str, int]] = None, + shuffle: bool = False, + mutate: bool = False, + mutation_rate: float = 0.1, + shuffle_frames: bool = False, + label_first: bool = True, + ): + self.path = Path(path) + self.crop_size = crop_size + self.num_classes = num_classes + self.codon_table = codon_table + self.shuffle = shuffle + self.mutate = mutate + self.mutation_rate = mutation_rate + self.shuffle_frames = shuffle_frames + self.label_first = label_first + self.rows = [] + with self.path.open("r", newline="") as fh: + reader = csv.reader(fh) + for row in reader: + if len(row) < 2: + continue + if label_first: + label = int(row[0]) + seq = row[1] + else: + seq = row[0] + label = int(row[1]) + self.rows.append((label, seq)) + + def __len__(self): + return len(self.rows) + + def __getitem__(self, idx): + label, seq = self.rows[idx] + seq_indices = dna_to_indices(seq) + if self.mutate: + from jaeger.data.pytorch.transforms import apply_mutation + + seq_indices = apply_mutation(seq_indices, self.mutation_rate) + + # Crop/pad nucleotide sequence + if len(seq_indices) > self.crop_size: + start = np.random.randint(0, len(seq_indices) - self.crop_size + 1) + seq_indices = seq_indices[start : start + self.crop_size] + elif len(seq_indices) < self.crop_size: + pad = self.crop_size - len(seq_indices) + seq_indices = np.pad(seq_indices, (0, pad), constant_values=4) + + if self.shuffle: + seq_indices = np.random.permutation(seq_indices) + + x = translate_to_codons(seq_indices, self.codon_table) + mask = x != 0 + + if self.shuffle_frames: + from jaeger.data.pytorch.transforms import shuffle_frames + + x = shuffle_frames(x) + + y = torch.nn.functional.one_hot( + torch.tensor(label), num_classes=self.num_classes + ).float() + return x, y, mask diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/data/pytorch/test_builders.py index 14cf13f..ca8f187 100644 --- a/tests/unit/data/pytorch/test_builders.py +++ b/tests/unit/data/pytorch/test_builders.py @@ -116,7 +116,7 @@ def test_build_datasets_unsupported_format(tmp_path): _make_npz(train_path, n_samples=4) _make_npz(val_path, n_samples=4) - config = _build_config([train_path], [val_path], data_format="csv") + config = _build_config([train_path], [val_path], data_format="hdf5") with pytest.raises(ValueError, match="Unsupported data_format"): build_datasets(config, branch="classifier") diff --git a/tests/unit/data/pytorch/test_dataset_csv.py b/tests/unit/data/pytorch/test_dataset_csv.py new file mode 100644 index 0000000..bb7d24b --- /dev/null +++ b/tests/unit/data/pytorch/test_dataset_csv.py @@ -0,0 +1,22 @@ +import csv + +from jaeger.data.pytorch.dataset_csv import CSVDataset +from jaeger.seqops.maps import CODONS + + +def test_csv_dataset(tmp_path): + path = tmp_path / "data.csv" + with path.open("w", newline="") as fh: + writer = csv.writer(fh) + writer.writerow([1, "ATG" * 200]) + writer.writerow([0, "AAA" * 200]) + writer.writerow([2, "TTT" * 200]) + + codon_table = {c: i + 1 for i, c in enumerate(CODONS)} + ds = CSVDataset(path, crop_size=60, num_classes=3, codon_table=codon_table) + assert len(ds) == 3 + x, y, mask = ds[0] + assert x.shape[0] == 6 + assert y.shape == (3,) + assert mask.shape == x.shape + assert y[1] == 1.0 From 96909d8bf3140c152870f508e3a6018a219f49c3 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:27:45 +0200 Subject: [PATCH 27/64] feat(pytorch): add train_one_epoch and evaluate engine --- src/jaeger/training/__init__.py | 0 src/jaeger/training/pytorch/__init__.py | 0 src/jaeger/training/pytorch/engine.py | 117 +++++++++++++++++++++ tests/unit/training/__init__.py | 0 tests/unit/training/pytorch/__init__.py | 0 tests/unit/training/pytorch/test_engine.py | 49 +++++++++ 6 files changed, 166 insertions(+) create mode 100644 src/jaeger/training/__init__.py create mode 100644 src/jaeger/training/pytorch/__init__.py create mode 100644 src/jaeger/training/pytorch/engine.py create mode 100644 tests/unit/training/__init__.py create mode 100644 tests/unit/training/pytorch/__init__.py create mode 100644 tests/unit/training/pytorch/test_engine.py diff --git a/src/jaeger/training/__init__.py b/src/jaeger/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/jaeger/training/pytorch/__init__.py b/src/jaeger/training/pytorch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py new file mode 100644 index 0000000..15c059a --- /dev/null +++ b/src/jaeger/training/pytorch/engine.py @@ -0,0 +1,117 @@ +from typing import Any, Callable, Dict, Optional + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader + + +def train_one_epoch( + model: nn.Module, + dataloader: DataLoader, + loss_fn: Callable, + optimizer: torch.optim.Optimizer, + device: torch.device, + metrics: Optional[Dict[str, Any]] = None, + forward_key: str = "prediction", + label_key: str = "label", + branch: str = "classifier", +) -> Dict[str, float]: + """Train for one epoch and return averaged loss and metrics.""" + model.train() + total_loss = 0.0 + total_samples = 0 + if metrics: + for m in metrics.values(): + if hasattr(m, "reset"): + m.reset() + + for batch in dataloader: + if branch == "classifier": + x, y, mask = [b.to(device) for b in batch] + optimizer.zero_grad() + outputs = model(x, mask=mask) + if isinstance(outputs, dict): + preds = outputs[forward_key] + else: + preds = outputs + loss = loss_fn(preds, y) + elif branch == "reliability": + x, y, mask = [b.to(device) for b in batch] + optimizer.zero_grad() + outputs = model(x, mask=mask) + preds = outputs[forward_key] + loss = loss_fn(preds, y) + else: + raise ValueError(f"Unknown branch: {branch}") + + loss.backward() + optimizer.step() + + batch_size = y.size(0) + total_loss += loss.item() * batch_size + total_samples += batch_size + + if metrics: + for metric in metrics.values(): + if hasattr(metric, "update"): + metric.update(preds.detach().cpu(), y.detach().cpu()) + + avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 + result = {"loss": avg_loss} + if metrics: + result.update( + { + name: metric.compute() if hasattr(metric, "compute") else float(metric) + for name, metric in metrics.items() + } + ) + return result + + +def evaluate( + model: nn.Module, + dataloader: DataLoader, + loss_fn: Callable, + device: torch.device, + metrics: Optional[Dict[str, Any]] = None, + forward_key: str = "prediction", + branch: str = "classifier", +) -> Dict[str, float]: + """Evaluate and return averaged loss and metrics.""" + model.eval() + total_loss = 0.0 + total_samples = 0 + if metrics: + for m in metrics.values(): + if hasattr(m, "reset"): + m.reset() + + with torch.no_grad(): + for batch in dataloader: + x, y, mask = [b.to(device) for b in batch] + outputs = model(x, mask=mask) + if isinstance(outputs, dict): + preds = outputs[forward_key] + else: + preds = outputs + loss = loss_fn(preds, y) + + batch_size = y.size(0) + total_loss += loss.item() * batch_size + total_samples += batch_size + + if metrics: + for metric in metrics.values(): + if hasattr(metric, "update"): + metric.update(preds.cpu(), y.cpu()) + + avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 + result = {"loss": avg_loss} + if metrics: + result.update( + { + name: metric.compute() if hasattr(metric, "compute") else float(metric) + for name, metric in metrics.items() + } + ) + return result diff --git a/tests/unit/training/__init__.py b/tests/unit/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/training/pytorch/__init__.py b/tests/unit/training/pytorch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py new file mode 100644 index 0000000..7290649 --- /dev/null +++ b/tests/unit/training/pytorch/test_engine.py @@ -0,0 +1,49 @@ +import torch + +from jaeger.training.pytorch.engine import evaluate, train_one_epoch + + +class DummyClassifier(torch.nn.Module): + def __init__(self): + super().__init__() + self.net = torch.nn.Sequential( + torch.nn.Linear(10, 5), + torch.nn.ReLU(), + torch.nn.Linear(5, 2), + ) + + def forward(self, x, mask=None): + return self.net(x) + + +def _make_dummy_classifier(): + return DummyClassifier() + + +def _make_dummy_loader(n=16): + xs = torch.randn(n, 10) + ys = torch.randint(0, 2, (n,)) + masks = torch.ones(n, dtype=torch.bool) + dataset = list(zip(xs, ys, masks)) + return torch.utils.data.DataLoader(dataset, batch_size=4) + + +def test_train_one_epoch(): + model = _make_dummy_classifier() + loader = _make_dummy_loader() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + metrics = {} + history = train_one_epoch( + model, loader, loss_fn, optimizer, torch.device("cpu"), metrics=metrics + ) + assert "loss" in history + + +def test_evaluate(): + model = _make_dummy_classifier() + loader = _make_dummy_loader() + loss_fn = torch.nn.CrossEntropyLoss() + metrics = {} + history = evaluate(model, loader, loss_fn, torch.device("cpu"), metrics=metrics) + assert "loss" in history From ad572d8e07b5d9ac26fbf6b2ea1e61f6a2144e6b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:30:39 +0200 Subject: [PATCH 28/64] feat(pytorch): add Trainer class --- src/jaeger/training/pytorch/trainer.py | 114 ++++++++++++++++++++ tests/unit/training/pytorch/test_trainer.py | 63 +++++++++++ 2 files changed, 177 insertions(+) create mode 100644 src/jaeger/training/pytorch/trainer.py create mode 100644 tests/unit/training/pytorch/test_trainer.py diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py new file mode 100644 index 0000000..6694a66 --- /dev/null +++ b/src/jaeger/training/pytorch/trainer.py @@ -0,0 +1,114 @@ +import json +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader + +from jaeger.training.pytorch.engine import evaluate, train_one_epoch + + +class Trainer: + """High-level training loop for Jaeger PyTorch models.""" + + def __init__( + self, + model: nn.Module, + train_loader: DataLoader, + val_loader: DataLoader, + loss_fn: Callable, + optimizer: torch.optim.Optimizer, + epochs: int, + device: torch.device, + metrics: Optional[Dict[str, Any]] = None, + callbacks: Optional[List[Any]] = None, + checkpoint_dir: Optional[str] = None, + history_path: Optional[str] = None, + branch: str = "classifier", + ): + self.model = model + self.train_loader = train_loader + self.val_loader = val_loader + self.loss_fn = loss_fn + self.optimizer = optimizer + self.epochs = epochs + self.device = device + self.metrics = metrics or {} + self.callbacks = callbacks or [] + self.checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir else None + self.history_path = Path(history_path) if history_path else None + self.branch = branch + self.history: List[Dict[str, float]] = [] + + def fit(self) -> List[Dict[str, float]]: + self.model.to(self.device) + for callback in self.callbacks: + if hasattr(callback, "on_train_begin"): + callback.on_train_begin(self) + + for epoch in range(1, self.epochs + 1): + for callback in self.callbacks: + if hasattr(callback, "on_epoch_begin"): + callback.on_epoch_begin(self, epoch) + + train_metrics = train_one_epoch( + self.model, + self.train_loader, + self.loss_fn, + self.optimizer, + self.device, + metrics=self.metrics, + branch=self.branch, + ) + val_metrics = evaluate( + self.model, + self.val_loader, + self.loss_fn, + self.device, + metrics=self.metrics, + branch=self.branch, + ) + + epoch_log = {"epoch": epoch} + for k, v in train_metrics.items(): + epoch_log[f"train_{k}"] = v + for k, v in val_metrics.items(): + epoch_log[f"val_{k}"] = v + self.history.append(epoch_log) + + for callback in self.callbacks: + if hasattr(callback, "on_epoch_end"): + callback.on_epoch_end(self, epoch, epoch_log) + + if self.checkpoint_dir: + self._save_checkpoint(epoch) + if self.history_path: + self._save_history() + + for callback in self.callbacks: + if hasattr(callback, "on_train_end"): + callback.on_train_end(self) + + return self.history + + def _save_checkpoint(self, epoch: int): + if self.checkpoint_dir is None: + return + self.checkpoint_dir.mkdir(parents=True, exist_ok=True) + path = self.checkpoint_dir / f"checkpoint_epoch_{epoch}.pt" + torch.save( + { + "epoch": epoch, + "model_state_dict": self.model.state_dict(), + "optimizer_state_dict": self.optimizer.state_dict(), + }, + path, + ) + + def _save_history(self): + if self.history_path is None: + return + self.history_path.parent.mkdir(parents=True, exist_ok=True) + with self.history_path.open("w") as fh: + json.dump(self.history, fh, indent=2) diff --git a/tests/unit/training/pytorch/test_trainer.py b/tests/unit/training/pytorch/test_trainer.py new file mode 100644 index 0000000..87c8a60 --- /dev/null +++ b/tests/unit/training/pytorch/test_trainer.py @@ -0,0 +1,63 @@ +import tempfile +from pathlib import Path + +import torch +from torch.utils.data import DataLoader, TensorDataset + +from jaeger.training.pytorch.trainer import Trainer + + +class _MaskedTensorDataset(TensorDataset): + def __getitem__(self, index): + x, y = super().__getitem__(index) + return x, y, torch.ones(x.shape[0], dtype=torch.bool) + + +class _DummyClassifier(torch.nn.Module): + def __init__(self): + super().__init__() + self.net = torch.nn.Sequential( + torch.nn.Linear(10, 5), + torch.nn.ReLU(), + torch.nn.Linear(5, 2), + ) + + def forward(self, x, mask=None): + return self.net(x) + + +def _make_dummy_data(n=16): + xs = torch.randn(n, 10) + ys = torch.randint(0, 2, (n,)) + return DataLoader(_MaskedTensorDataset(xs, ys), batch_size=4) + + +def _make_dummy_model(): + return _DummyClassifier() + + +def test_trainer_fit(): + model = _make_dummy_model() + train_loader = _make_dummy_data() + val_loader = _make_dummy_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + checkpoint_dir=tmpdir, + history_path=Path(tmpdir) / "history.json", + ) + history = trainer.fit() + assert len(history) == 2 + assert "train_loss" in history[0] + assert "val_loss" in history[0] + assert (Path(tmpdir) / "history.json").exists() + assert len(list(Path(tmpdir).glob("checkpoint_epoch_*.pt"))) == 2 From b50e44f21c565b8672d0cd97ec06864ea74d5c9b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:35:50 +0200 Subject: [PATCH 29/64] feat(pytorch): add EarlyStopping and ModelCheckpoint callbacks --- src/jaeger/training/pytorch/callbacks.py | 109 ++++++++++++++++++ src/jaeger/training/pytorch/trainer.py | 4 + tests/unit/training/pytorch/test_callbacks.py | 83 +++++++++++++ 3 files changed, 196 insertions(+) create mode 100644 src/jaeger/training/pytorch/callbacks.py create mode 100644 tests/unit/training/pytorch/test_callbacks.py diff --git a/src/jaeger/training/pytorch/callbacks.py b/src/jaeger/training/pytorch/callbacks.py new file mode 100644 index 0000000..3b2b160 --- /dev/null +++ b/src/jaeger/training/pytorch/callbacks.py @@ -0,0 +1,109 @@ +from pathlib import Path +from typing import Any, Dict, Optional + +import torch + + +class EarlyStopping: + """Stop training when a monitored metric has stopped improving.""" + + def __init__( + self, + monitor: str = "val_loss", + patience: int = 5, + mode: str = "min", + min_delta: float = 0.0, + restore_best_weights: bool = False, + ): + self.monitor = monitor + self.patience = patience + self.mode = mode + self.min_delta = min_delta + self.restore_best_weights = restore_best_weights + self.best_value = float("inf") if mode == "min" else float("-inf") + self.best_epoch = 0 + self.wait = 0 + self.best_state: Optional[Dict[str, Any]] = None + self.stopped_epoch = 0 + + def on_train_begin(self, trainer): + pass + + def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): + current = logs.get(self.monitor) + if current is None: + return + + improved = (self.mode == "min" and current < self.best_value - self.min_delta) or ( + self.mode == "max" and current > self.best_value + self.min_delta + ) + + if improved: + self.best_value = current + self.best_epoch = epoch + self.wait = 0 + if self.restore_best_weights: + self.best_state = {k: v.cpu().clone() for k, v in trainer.model.state_dict().items()} + else: + self.wait += 1 + if self.wait >= self.patience: + trainer.should_stop = True + self.stopped_epoch = epoch + + def on_train_end(self, trainer): + if self.restore_best_weights and self.best_state is not None: + trainer.model.load_state_dict(self.best_state) + + +class ModelCheckpoint: + """Save the model whenever a monitored metric improves.""" + + def __init__( + self, + filepath: str, + monitor: str = "val_loss", + mode: str = "min", + save_best_only: bool = True, + verbose: int = 0, + ): + self.filepath = Path(filepath) + self.monitor = monitor + self.mode = mode + self.save_best_only = save_best_only + self.verbose = verbose + self.best_value = float("inf") if mode == "min" else float("-inf") + + def on_train_begin(self, trainer): + pass + + def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): + current = logs.get(self.monitor) + if current is None: + return + + improved = (self.mode == "min" and current < self.best_value) or ( + self.mode == "max" and current > self.best_value + ) + + if improved: + self.best_value = current + if self.save_best_only: + self._save(trainer, epoch) + if self.verbose: + print(f"Epoch {epoch}: {self.monitor} improved to {current:.4f}; saving model") + elif not self.save_best_only: + self._save(trainer, epoch) + + def _save(self, trainer, epoch: int): + self.filepath.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "epoch": epoch, + "model_state_dict": trainer.model.state_dict(), + "optimizer_state_dict": trainer.optimizer.state_dict(), + }, + self.filepath, + ) + + def on_train_end(self, trainer): + pass diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 6694a66..27bab82 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -40,6 +40,7 @@ def __init__( self.history_path = Path(history_path) if history_path else None self.branch = branch self.history: List[Dict[str, float]] = [] + self.should_stop = False def fit(self) -> List[Dict[str, float]]: self.model.to(self.device) @@ -81,6 +82,9 @@ def fit(self) -> List[Dict[str, float]]: if hasattr(callback, "on_epoch_end"): callback.on_epoch_end(self, epoch, epoch_log) + if getattr(self, "should_stop", False): + break + if self.checkpoint_dir: self._save_checkpoint(epoch) if self.history_path: diff --git a/tests/unit/training/pytorch/test_callbacks.py b/tests/unit/training/pytorch/test_callbacks.py new file mode 100644 index 0000000..e2a837b --- /dev/null +++ b/tests/unit/training/pytorch/test_callbacks.py @@ -0,0 +1,83 @@ +import tempfile +from pathlib import Path + +import torch +from torch.utils.data import DataLoader, TensorDataset + +from jaeger.training.pytorch.callbacks import EarlyStopping, ModelCheckpoint +from jaeger.training.pytorch.trainer import Trainer + + +class _MaskedTensorDataset(TensorDataset): + def __getitem__(self, index): + x, y = super().__getitem__(index) + return x, y, torch.ones(x.shape[0], dtype=torch.bool) + + +class _DummyClassifier(torch.nn.Module): + def __init__(self): + super().__init__() + self.net = torch.nn.Sequential( + torch.nn.Linear(10, 5), + torch.nn.ReLU(), + torch.nn.Linear(5, 2), + ) + + def forward(self, x, mask=None): + return self.net(x) + + +def _make_data(n=16): + xs = torch.randn(n, 10) + ys = torch.randint(0, 2, (n,)) + return DataLoader(_MaskedTensorDataset(xs, ys), batch_size=4) + + +def _make_model(): + return _DummyClassifier() + + +def test_early_stopping(): + model = _make_model() + train_loader = _make_data() + val_loader = _make_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.0) + early_stop = EarlyStopping(monitor="val_loss", patience=1, restore_best_weights=True) + + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=10, + device=torch.device("cpu"), + callbacks=[early_stop], + ) + history = trainer.fit() + assert len(history) < 10 + + +def test_model_checkpoint(): + model = _make_model() + train_loader = _make_data() + val_loader = _make_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt_path = Path(tmpdir) / "best.pt" + checkpoint = ModelCheckpoint(filepath=ckpt_path, monitor="val_loss", mode="min") + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + callbacks=[checkpoint], + ) + trainer.fit() + assert ckpt_path.exists() From 682f40947116fedcd8546e0572022447f0e4ff83 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:38:13 +0200 Subject: [PATCH 30/64] feat(pytorch): add distributed training helper --- src/jaeger/training/pytorch/distributed.py | 42 +++++++++++++++++++ .../unit/training/pytorch/test_distributed.py | 19 +++++++++ 2 files changed, 61 insertions(+) create mode 100644 src/jaeger/training/pytorch/distributed.py create mode 100644 tests/unit/training/pytorch/test_distributed.py diff --git a/src/jaeger/training/pytorch/distributed.py b/src/jaeger/training/pytorch/distributed.py new file mode 100644 index 0000000..ce31914 --- /dev/null +++ b/src/jaeger/training/pytorch/distributed.py @@ -0,0 +1,42 @@ +import os + +import torch +import torch.distributed as dist + + +def setup_distributed(backend: str = "nccl") -> bool: + """Initialize process group if multi-GPU / SLURM environment is detected.""" + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + return True + + if "SLURM_PROCID" in os.environ: + rank = int(os.environ["SLURM_PROCID"]) + world_size = int(os.environ["SLURM_NTASKS"]) + local_rank = int(os.environ.get("SLURM_LOCALID", 0)) + dist.init_process_group(backend=backend, init_method="env://", rank=rank, world_size=world_size) + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + return True + + return False + + +def cleanup_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def get_device() -> torch.device: + if dist.is_initialized(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + return torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def is_main_process() -> bool: + if not dist.is_initialized(): + return True + return dist.get_rank() == 0 diff --git a/tests/unit/training/pytorch/test_distributed.py b/tests/unit/training/pytorch/test_distributed.py new file mode 100644 index 0000000..e7ca931 --- /dev/null +++ b/tests/unit/training/pytorch/test_distributed.py @@ -0,0 +1,19 @@ +import os + +from jaeger.training.pytorch.distributed import get_device, is_main_process, setup_distributed + + +def test_get_device_without_distributed(): + device = get_device() + assert device.type in ("cpu", "cuda") + + +def test_is_main_process_without_distributed(): + assert is_main_process() is True + + +def test_setup_distributed_returns_false_outside_slurm(): + # Ensure no SLURM / torchrun env vars are present + for key in ["RANK", "WORLD_SIZE", "SLURM_PROCID", "SLURM_NTASKS"]: + os.environ.pop(key, None) + assert setup_distributed() is False From 31adbae452f7f01ccba7ac63b6474806ca96cf77 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:51:41 +0200 Subject: [PATCH 31/64] feat(pytorch): update jaeger train command for PyTorch --- src/jaeger/commands/train.py | 844 ++++++++-------------- tests/unit/commands/test_train_pytorch.py | 147 ++++ 2 files changed, 466 insertions(+), 525 deletions(-) create mode 100644 tests/unit/commands/test_train_pytorch.py diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 0d9bbda..6f20935 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -1,8 +1,8 @@ -"""Training CLI entry point for Jaeger. +"""Training CLI entry point for Jaeger (PyTorch). -All core model-building logic lives in `jaeger.nnlib.builder`. +All core model-building logic lives in `jaeger.nnlib.pytorch.builder`. This module is a thin Click wrapper that orchestrates: -- strategy selection (GPU/CPU, mixed precision) +- strategy selection (GPU/CPU) - data pipeline construction - model training loops - saving @@ -11,24 +11,28 @@ from __future__ import annotations import os +import shutil +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple # temporary fix os.environ["WRAPT_DISABLE_EXTENSIONS"] = "true" import click -import tensorflow as tf -from pathlib import Path - -from jaeger.nnlib.builder import DynamicModelBuilder, check_files -from jaeger.seqops.encode import process_string_train -from jaeger.utils.misc import load_model_config, numerize -from jaeger.utils.logging import get_logger -from jaeger.data.tfrecord import _make_parse_tfrecord_fn -from jaeger.data.loaders import ( - _load_numpy_raw_dataset, - _load_numpy_raw_variable_dataset, - _load_numpy_full_dataset, +import torch +import torch.nn as nn + +from jaeger.data.pytorch.builders import build_datasets +from jaeger.nnlib.pytorch.builder import ModelBuilder +from jaeger.training.pytorch.callbacks import EarlyStopping, ModelCheckpoint +from jaeger.training.pytorch.distributed import ( + cleanup_distributed, + get_device, + setup_distributed, ) +from jaeger.training.pytorch.trainer import Trainer +from jaeger.utils.logging import get_logger +from jaeger.utils.misc import load_model_config, numerize try: from icecream import ic @@ -43,6 +47,153 @@ def ic(*args, **kwargs): logger = get_logger(log_file=None, log_path=None, level=3) +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +def _count_parameters(model: nn.Module) -> Tuple[int, int]: + """Return (total_params, trainable_params) for *model*.""" + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return total, trainable + + +def _load_checkpoint_if_requested( + model: nn.Module, + optimizer: Optional[torch.optim.Optimizer], + checkpoint_path: Optional[str | Path], +) -> None: + """Load model and optimizer state dicts from *checkpoint_path* if provided.""" + if checkpoint_path is None: + return + path = Path(checkpoint_path) + if not path.exists(): + logger.warning("Checkpoint path does not exist: %s", path) + return + logger.info("Loading checkpoint from %s", path) + checkpoint = torch.load(path, map_location="cpu", weights_only=True) + model.load_state_dict(checkpoint["model_state_dict"]) + if optimizer is not None and "optimizer_state_dict" in checkpoint: + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + + +def _normalize_optimizer_params(params: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Convert TF-style optimizer params to PyTorch-compatible ones.""" + if not params: + return {} + normalized = dict(params) + if "learning_rate" in normalized: + normalized["lr"] = normalized.pop("learning_rate") + for key in ("clipnorm", "clipvalue", "global_clipnorm"): + if key in normalized: + logger.warning("Ignoring unsupported optimizer parameter: %s", key) + normalized.pop(key) + return normalized + + +def _callback_path_from_config( + callbacks_cfg: List[Dict[str, Any]], name: str +) -> Optional[Dict[str, Any]]: + """Return the first callback entry matching *name* from the callbacks list.""" + for entry in callbacks_cfg or []: + if entry.get("name") == name: + return entry.get("params", {}) or {} + return None + + +def _build_callbacks(train_cfg: Dict[str, Any]) -> List[Any]: + """Build PyTorch callbacks from the training config.""" + callbacks: List[Any] = [] + callbacks_cfg = train_cfg.get("callbacks", {}).get("classifier", []) + + # EarlyStopping + early_params = _callback_path_from_config(callbacks_cfg, "EarlyStopping") + patience = train_cfg.get("patience") + if early_params is not None and patience is None: + patience = early_params.get("patience") + if patience is not None: + callbacks.append( + EarlyStopping( + monitor=(early_params or {}).get("monitor", "val_loss") + or train_cfg.get("early_stopping_monitor", "val_loss"), + patience=int(patience), + mode=(early_params or {}).get("mode", "min") + or train_cfg.get("early_stopping_mode", "min"), + restore_best_weights=(early_params or {}).get( + "restore_best_weights", True + ), + ) + ) + + # ModelCheckpoint + checkpoint_params = _callback_path_from_config(callbacks_cfg, "ModelCheckpoint") + checkpoint_path = train_cfg.get("checkpoint_path") + if checkpoint_params is not None and checkpoint_path is None: + checkpoint_path = checkpoint_params.get("filepath") + if checkpoint_path is not None: + callbacks.append( + ModelCheckpoint( + filepath=checkpoint_path, + monitor=(checkpoint_params or {}).get("monitor", "val_loss") + or train_cfg.get("checkpoint_monitor", "val_loss"), + mode=(checkpoint_params or {}).get("mode", "min") + or train_cfg.get("checkpoint_mode", "min"), + save_best_only=(checkpoint_params or {}).get("save_best_only", True), + verbose=int((checkpoint_params or {}).get("verbose", 1) or 0), + ) + ) + + return callbacks + + +def _find_latest_checkpoint(directory: Path) -> Optional[Path]: + """Return the most recently modified ``.pt`` file in *directory*, if any.""" + if not directory.exists(): + return None + checkpoints = sorted(directory.glob("*.pt"), key=lambda p: p.stat().st_mtime) + return checkpoints[-1] if checkpoints else None + + +def _save_model_checkpoint( + model: nn.Module, + save_dir: Path, + filename: str, + metadata: Optional[str] = None, + num_params: Optional[str] = None, +) -> None: + """Persist *model*'s state dict to ``save_dir / filename``.""" + save_dir.mkdir(parents=True, exist_ok=True) + payload: Dict[str, Any] = { + "model_state_dict": model.state_dict(), + } + if metadata is not None: + payload["metadata"] = str(metadata) + if num_params is not None: + payload["num_params"] = num_params + torch.save(payload, save_dir / filename) + logger.info("Saved checkpoint to %s", save_dir / filename) + + +class _ReliabilityPipeline(nn.Module): + """Runs the representation model and reliability head.""" + + def __init__(self, rep_model: nn.Module, head: nn.Module): + super().__init__() + self.rep_model = rep_model + self.head = head + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + outputs = self.rep_model(x, mask) + if isinstance(outputs, tuple): + nmd = outputs[1] if len(outputs) > 1 else outputs[0] + else: + nmd = outputs + return self.head(nmd) + + # ------------------------------------------------------------------ # CLI command # ------------------------------------------------------------------ @@ -103,535 +254,178 @@ def train_fragment( def train_fragment_core(**kwargs): - """Train fragment classification and reliability models.""" - gpus = tf.config.list_physical_devices("GPU") - num_gpus = len(gpus) - logger.info(f"Physical GPUs detected: {num_gpus}") + """Train fragment classification and reliability models with PyTorch.""" + config = load_model_config(Path(kwargs["config"])) + config["use_pytorch"] = True + config["mix_precision"] = kwargs.get("mixed_precision", False) + config["from_last_checkpoint"] = kwargs.get("from_last_checkpoint") + config["force"] = kwargs.get("force") + config["use_xla"] = kwargs.get("xla", False) - if num_gpus > 1: - strategy = tf.distribute.MirroredStrategy() - logger.info( - f"Using MirroredStrategy with {strategy.num_replicas_in_sync} replicas" + if kwargs.get("mixed_precision", False): + logger.warning( + "mixed_precision is not implemented for PyTorch training yet; ignoring" ) - elif num_gpus == 1: - strategy = tf.distribute.OneDeviceStrategy(device="/GPU:0") - logger.info("Using OneDeviceStrategy on /GPU:0") - else: - strategy = tf.distribute.get_strategy() - logger.info("No GPU detected, using default CPU strategy") - if kwargs.get("mixed_precision", False): - logger.info("experimental: using mix precision floats for faster training") - policy = tf.keras.mixed_precision.Policy("mixed_float16") - tf.keras.mixed_precision.set_global_policy(policy) - - multi_gpu = strategy.num_replicas_in_sync > 1 - - with strategy.scope(): - logger.info("initializing model") - config = load_model_config(Path(kwargs.get("config"))) - config["mix_precision"] = kwargs.get("mixed_precision", False) - config["from_last_checkpoint"] = kwargs.get("from_last_checkpoint") - config["force"] = kwargs.get("force") - config["use_xla"] = kwargs.get("xla", False) - if config["use_xla"]: - logger.info("Using XLA JIT compilation for training") - - builder = DynamicModelBuilder(config) + if kwargs.get("xla", False): + logger.warning("XLA is not supported for PyTorch training; ignoring") + + if kwargs.get("self_supervised_pretraining", False): + logger.warning("self_supervised_pretraining is not implemented yet; skipping") + + setup_distributed() + device = get_device() + logger.info("Using device: %s", device) + + try: + builder = ModelBuilder(config) models = builder.build_fragment_classifier() - models.get("rep_model").summary() - model_num_params = numerize(models.get("rep_model").count_params(), decimal=1) + rep_model = models["rep_model"] + + total_params, trainable_params = _count_parameters(rep_model) + logger.info( + "rep_model parameters: total=%s, trainable=%s", + numerize(total_params, decimal=1), + numerize(trainable_params, decimal=1), + ) + model_num_params = numerize(total_params, decimal=1) + + train_cfg = config.get("training", {}) + classifier_dir = Path(train_cfg.get("classifier_dir", "checkpoints/classifier")) + reliability_dir = Path( + train_cfg.get("reliability_dir", "checkpoints/reliability") + ) + model_save_cfg = train_cfg.get("model_saving", {}) or {} + model_save_path = model_save_cfg.get("path") + if model_save_path: + model_save_path = Path(model_save_path) + + if kwargs.get("force", False): + for directory in (classifier_dir, reliability_dir): + if directory.exists(): + shutil.rmtree(directory) + logger.info("Removed existing checkpoint directory: %s", directory) # ================= train classifier ====================== - builder.compile_model(models, train_branch="classifier") + if not kwargs.get("only_reliability_head", False): + classifier_loaders = build_datasets(config, branch="classifier") - string_processor_config = builder._get_string_processor_config() - _train_data = builder._get_fragment_paths() - train_data = {"train": None, "validation": None} + # Normalize optimizer params before compiling so the builder uses them. + opt_params = train_cfg.get("optimizer_params", {}) or {} + train_cfg["optimizer_params"] = _normalize_optimizer_params(opt_params) - data_format = string_processor_config.get("data_format", "csv") - _buffer_size = string_processor_config.get("buffer_size") + model, optimizer, loss_fn = builder.compile_model( + models, train_branch="classifier" + ) - for k, v in _train_data.items(): - paths = check_files(v.get("paths")) - if not paths: - logger.warning("no valid files in paths=%r %r", k, v.get("paths")) - exit(1) + checkpoint_path = None + if kwargs.get("from_last_checkpoint", False): + checkpoint_path = _find_latest_checkpoint(classifier_dir) + _load_checkpoint_if_requested(model, optimizer, checkpoint_path) - if data_format == "csv": - _data = tf.data.TextLineDataset( - paths, num_parallel_reads=len(paths), buffer_size=200 - ) - if string_processor_config.get("input_type") == "translated": - padded_shape = { - "translated": [6, None] - if string_processor_config.get("use_embedding_layer") is True - else [ - 6, - None, - string_processor_config.get("codon_depth"), - ] - } - elif string_processor_config.get("input_type") == "nucleotide": - padded_shape = { - "nucleotide": [2, string_processor_config.get("crop_size"), 4] - } - train_data[k] = ( - _data.map( - process_string_train( - codons=string_processor_config.get("codon"), - codon_num=string_processor_config.get("codon_id"), - codon_depth=string_processor_config.get("codon_depth"), - label_original=string_processor_config.get( - "classifier_labels", None - ), - label_alternative=string_processor_config.get( - "classifier_labels_map", None - ), - ngram_width=string_processor_config.get("ngram_width"), - seq_onehot=string_processor_config.get("seq_onehot"), - crop_size=string_processor_config.get("crop_size"), - input_type=string_processor_config.get("input_type"), - masking=string_processor_config.get("masking"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - num_classes=builder.classifier_out_dim, - class_label_onehot=False - if "binary" - in builder.train_cfg.get( - "loss_classifier", "categorical_crossentropy" - ).lower() - else True, - shuffle=string_processor_config.get("shuffle"), - shuffle_frames=string_processor_config.get( - "shuffle_frames", False - ), - ), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .shuffle( - buffer_size=_data.cardinality() - if _buffer_size == -1 - else _buffer_size, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - padded_shape, - [builder.classifier_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - elif data_format == "tfrecord": - logger.info(f"Loading {k} data from TFRecord: {paths}") - _data = tf.data.TFRecordDataset( - paths, num_parallel_reads=tf.data.AUTOTUNE - ) - parse_fn = _make_parse_tfrecord_fn( - input_type=string_processor_config.get("input_type"), - use_embedding_layer=string_processor_config.get( - "use_embedding_layer", False - ), - codon_depth=string_processor_config.get("codon_depth"), - crop_size=string_processor_config.get("crop_size"), - num_classes=builder.classifier_out_dim, - ) - train_data[k] = ( - _data.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE) - .cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .batch( - batch_size=builder.train_cfg.get("batch_size"), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - elif data_format == "numpy_raw": - logger.info(f"Loading {k} data from NumPy raw: {paths}") - if len(paths) > 1: - logger.warning( - "NumPy raw format only supports a single .npz file per split; using first: %s", - paths[0], - ) - _data = _load_numpy_raw_dataset( - paths[0], - crop_size=string_processor_config.get("crop_size"), - ngram_width=string_processor_config.get("ngram_width"), - num_classes=builder.classifier_out_dim, - shuffle=string_processor_config.get("shuffle"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle_frames=string_processor_config.get("shuffle_frames", False), - ) - train_data[k] = ( - _data.cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - {"translated": [6, None]}, - [builder.classifier_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - elif data_format == "numpy_raw_variable": - logger.info(f"Loading {k} data from NumPy raw variable: {paths}") - if len(paths) > 1: - logger.warning( - "NumPy raw variable format only supports a single .npz file per split; using first: %s", - paths[0], - ) - _data = _load_numpy_raw_variable_dataset( - paths[0], - ngram_width=string_processor_config.get("ngram_width"), - num_classes=builder.classifier_out_dim, - shuffle=string_processor_config.get("shuffle"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle_frames=string_processor_config.get("shuffle_frames", False), - ) - train_data[k] = ( - _data.cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - {"translated": [6, None]}, - [builder.classifier_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - elif data_format == "numpy_full": - logger.info(f"Loading {k} data from NumPy full: {paths}") - if len(paths) > 1: - logger.warning( - "NumPy full format only supports a single .npz file per split; using first: %s", - paths[0], - ) - _data = _load_numpy_full_dataset(paths[0]) - train_data[k] = ( - _data.cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - {"translated": [6, None]}, - [builder.classifier_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - else: - raise ValueError( - f"Unsupported data_format: {data_format}. Use 'csv', 'tfrecord', 'numpy_raw', 'numpy_raw_variable', or 'numpy_full'." - ) + if kwargs.get("only_classification_head", False) or kwargs.get( + "only_heads", False + ): + for param in rep_model.parameters(): + param.requires_grad = False - # ============ check convergence & train classifier =========== - checkpoint = builder._checkpoints - cls_converged = checkpoint and checkpoint.get("classifier", {}).get( - "is_converged", False - ) - if kwargs.get("only_save", False) is False: - if not cls_converged and not kwargs.get("only_reliability_head", False): - train_args = { - "validation_data": train_data.get("validation").take( - builder.train_cfg.get("classifier_validation_steps") - ), - "epochs": builder.train_cfg.get("classifier_epochs"), - "callbacks": builder.get_callbacks(branch="classifier"), - } - if checkpoint: - train_args["initial_epoch"] = checkpoint.get("classifier", {}).get( - "epoch", 0 - ) - - if kwargs.get("only_classification_head", False) or kwargs.get( - "only_heads", False - ): - models.get("rep_model").trainable = False - - # self-supervised pre-training - if kwargs.get("self_supervised_pretraining", False): - builder.compile_model(models, train_branch="pretrain") - models.get("jaeger_projection").summary() - self_supervised_train_args = { - "validation_data": train_data.get("validation").take( - builder.train_cfg.get("classifier_validation_steps") - ), - "epochs": builder.train_cfg.get("projection_epochs"), - "callbacks": builder.get_callbacks(branch="projection"), - } - if checkpoint: - self_supervised_train_args["initial_epoch"] = checkpoint.get( - "projection", {} - ).get("epoch", 0) - models.get("jaeger_projection").fit( - train_data.get("train").take( - builder.train_cfg.get("classifier_train_steps") - ), - **self_supervised_train_args, - ) - - # train classification model - models.get("jaeger_classifier").fit( - train_data.get("train").take( - builder.train_cfg.get("classifier_train_steps"), - ), - class_weight=builder.train_cfg.get("classifier_class_weights"), - **train_args, - ) + callbacks = _build_callbacks(train_cfg) + epochs = int(train_cfg.get("classifier_epochs", 1)) + + if kwargs.get("only_save", False): + logger.info("Skipping classifier training (--only_save)") else: - logger.info("Skipping training — classification model") - - # ============== reliability model ======================== - builder.compile_model(models, train_branch="reliability") - - _rel_train_data = builder._get_reliability_fragment_paths() - rel_train_data = {"train": None, "validation": None} - for k, v in _rel_train_data.items(): - paths = check_files(v.get("paths")) - if not paths: - logger.warning("no valid files in paths=%r", k, v.get("paths")) - exit(1) - - if data_format == "csv": - _data = tf.data.TextLineDataset( - v.get("paths"), - num_parallel_reads=len(v.get("paths")), - buffer_size=200, - ) - if string_processor_config.get("input_type") == "translated": - padded_shape = { - "translated": [6, None] - if string_processor_config.get("use_embedding_layer") is True - else [ - 6, - string_processor_config.get("crop_size") // 3 - 1, - string_processor_config.get("codon_depth"), - ] - } - elif string_processor_config.get("input_type") == "nucleotide": - padded_shape = { - "nucleotide": [2, string_processor_config.get("crop_size"), 4] - } - rel_train_data[k] = ( - _data.map( - process_string_train( - codons=string_processor_config.get("codon"), - codon_num=string_processor_config.get("codon_id"), - codon_depth=string_processor_config.get("codon_depth"), - ngram_width=string_processor_config.get("ngram_width"), - seq_onehot=string_processor_config.get("seq_onehot"), - crop_size=string_processor_config.get("crop_size"), - input_type=string_processor_config.get("input_type"), - masking=string_processor_config.get("masking"), - num_classes=builder.reliability_out_dim, - class_label_onehot=False, - shuffle_frames=string_processor_config.get( - "shuffle_frames", False - ), - ), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .shuffle( - buffer_size=_data.cardinality() - if _buffer_size == -1 - else _buffer_size, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - padded_shape, - [builder.reliability_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - elif data_format == "tfrecord": - logger.info(f"Loading reliability {k} data from TFRecord: {paths}") - _data = tf.data.TFRecordDataset( - paths, num_parallel_reads=tf.data.AUTOTUNE - ) - parse_fn = _make_parse_tfrecord_fn( - input_type=string_processor_config.get("input_type"), - use_embedding_layer=string_processor_config.get( - "use_embedding_layer", False - ), - codon_depth=string_processor_config.get("codon_depth"), - crop_size=string_processor_config.get("crop_size"), - num_classes=builder.reliability_out_dim, + trainer = Trainer( + model=model, + train_loader=classifier_loaders["train"], + val_loader=classifier_loaders["validation"], + loss_fn=loss_fn, + optimizer=optimizer, + epochs=epochs, + device=device, + metrics=builder.get_metrics(branch="classifier"), + callbacks=callbacks, + checkpoint_dir=str(classifier_dir), + history_path=str(classifier_dir / "history.json"), + branch="classifier", ) - rel_train_data[k] = ( - _data.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE) - .cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .batch( - batch_size=builder.train_cfg.get("batch_size"), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) + trainer.fit() + + # ================= train reliability ====================== + if ( + "reliability_model" in config.get("model", {}) + and not kwargs.get("only_classification_head", False) + and models.get("reliability_head") is not None + ): + try: + reliability_loaders = build_datasets(config, branch="reliability") + except (KeyError, ValueError) as exc: + logger.warning( + "Reliability data not configured; skipping reliability training: %s", + exc, ) - elif data_format == "numpy_raw": - logger.info(f"Loading reliability {k} data from NumPy raw: {paths}") - if len(paths) > 1: - logger.warning( - "NumPy raw format only supports a single .npz file per split; using first: %s", - paths[0], - ) - _data = _load_numpy_raw_dataset( - paths[0], - crop_size=string_processor_config.get("crop_size"), - ngram_width=string_processor_config.get("ngram_width"), - num_classes=builder.reliability_out_dim, - shuffle=string_processor_config.get("shuffle"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle_frames=string_processor_config.get("shuffle_frames", False), - ) - rel_train_data[k] = ( - _data.cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - {"translated": [6, None]}, - [builder.reliability_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - elif data_format == "numpy_raw_variable": - logger.info( - f"Loading reliability {k} data from NumPy raw variable: {paths}" - ) - if len(paths) > 1: - logger.warning( - "NumPy raw variable format only supports a single .npz file per split; using first: %s", - paths[0], - ) - _data = _load_numpy_raw_variable_dataset( - paths[0], - ngram_width=string_processor_config.get("ngram_width"), - num_classes=builder.reliability_out_dim, - shuffle=string_processor_config.get("shuffle"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle_frames=string_processor_config.get("shuffle_frames", False), + reliability_loaders = None + + if reliability_loaders is not None and not kwargs.get("only_save", False): + rel_pipeline = _ReliabilityPipeline( + rep_model, models["reliability_head"] ) - rel_train_data[k] = ( - _data.cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - {"translated": [6, None]}, - [builder.reliability_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) + rel_optimizer = torch.optim.Adam( + rel_pipeline.parameters(), + lr=(train_cfg.get("optimizer_params", {}) or {}).get("lr", 1e-3), ) - elif data_format == "numpy_full": - logger.info(f"Loading reliability {k} data from NumPy full: {paths}") - if len(paths) > 1: - logger.warning( - "NumPy full format only supports a single .npz file per split; using first: %s", - paths[0], - ) - _data = _load_numpy_full_dataset(paths[0]) - rel_train_data[k] = ( - _data.cache() - .shuffle( - buffer_size=_buffer_size if _buffer_size != -1 else 100000, - ) - .padded_batch( - batch_size=builder.train_cfg.get("batch_size"), - padded_shapes=( - {"translated": [6, None]}, - [builder.reliability_out_dim], - ), - drop_remainder=multi_gpu, - ) - .prefetch(tf.data.AUTOTUNE) - ) - else: - raise ValueError( - f"Unsupported data_format: {data_format}. Use 'csv', 'tfrecord', 'numpy_raw', 'numpy_raw_variable', or 'numpy_full'." + rel_out_dim = builder.reliability_out_dim + rel_loss = ( + nn.BCEWithLogitsLoss() + if rel_out_dim == 1 + else nn.CrossEntropyLoss() ) - # ============== check convergence & train reliability ======== - checkpoint = builder._checkpoints - rel_converged = checkpoint and checkpoint.get("reliability", {}).get( - "is_converged", False - ) - if kwargs.get("only_save", False) is False: - if ( - not rel_converged - and not kwargs.get("only_classification_head", False) - and models.get("jaeger_reliability") is not None - ): - train_args = { - "validation_data": rel_train_data.get("validation").take( - builder.train_cfg.get("reliability_validation_steps") - ), - "epochs": builder.train_cfg.get("reliability_epochs"), - "callbacks": builder.get_callbacks(branch="reliability"), - } - if checkpoint: - train_args["initial_epoch"] = checkpoint.get("reliability", {}).get( - "epoch", 0 - ) - - models.get("jaeger_reliability").fit( - rel_train_data.get("train").take( - builder.train_cfg.get("reliability_train_steps") - ), - class_weight=builder.train_cfg.get("reliability_class_weights"), - **train_args, + rel_checkpoint_path = None + if kwargs.get("from_last_checkpoint", False): + rel_checkpoint_path = _find_latest_checkpoint(reliability_dir) + _load_checkpoint_if_requested( + rel_pipeline, rel_optimizer, rel_checkpoint_path ) - else: - logger.info("Skipping training — reliability model") - # ============= test final model ========================= - logger.info("testing the final model") - models.get("jaeger_model").trainable = False - models.get("jaeger_model").predict(train_data.get("validation").take(100)) - logger.info("training completed!") - - # ============= saving =================================== - if kwargs.get("save_model", False): - builder.save_model( - model=models.get("jaeger_model"), - num_params=model_num_params, - suffix="fragment", - metadata=kwargs.get("meta", None), - ) - builder.save_embedding_model( - models=models, + rel_callbacks = _build_callbacks(train_cfg) + rel_epochs = int(train_cfg.get("reliability_epochs", 1)) + + rel_trainer = Trainer( + model=rel_pipeline, + train_loader=reliability_loaders["train"], + val_loader=reliability_loaders["validation"], + loss_fn=rel_loss, + optimizer=rel_optimizer, + epochs=rel_epochs, + device=device, + metrics=builder.get_metrics(branch="reliability"), + callbacks=rel_callbacks, + checkpoint_dir=str(reliability_dir), + history_path=str(reliability_dir / "history.json"), + branch="reliability", + ) + rel_trainer.fit() + + # ================= saving ====================== + if kwargs.get("save_model", False) or kwargs.get("only_save", False): + save_dir = model_save_path if model_save_path else classifier_dir + _save_model_checkpoint( + models["jaeger_classifier"], + save_dir, + "classifier.pt", + metadata=kwargs.get("meta"), num_params=model_num_params, - suffix="fragment", - metadata=kwargs.get("meta", None), ) + if models.get("reliability_head") is not None: + _save_model_checkpoint( + _ReliabilityPipeline(rep_model, models["reliability_head"]), + save_dir, + "reliability.pt", + metadata=kwargs.get("meta"), + num_params=model_num_params, + ) + + logger.info("training completed!") + finally: + cleanup_distributed() diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py new file mode 100644 index 0000000..da75421 --- /dev/null +++ b/tests/unit/commands/test_train_pytorch.py @@ -0,0 +1,147 @@ +"""Smoke tests for the PyTorch-based ``jaeger train`` command.""" + +from __future__ import annotations + +import numpy as np +import yaml + +from jaeger.commands.train import train_fragment_core + + +def _make_raw_npz(path, n_samples, seq_length=50): + """Create a minimal numpy_raw dataset archive.""" + seqs = np.random.randint(0, 4, size=(n_samples, seq_length)).astype(np.int8) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=n_samples)] + np.savez(path, sequences=seqs, labels=labels) + + +def _build_config(tmp_path, train_path, val_path, data_format="numpy_raw"): + return { + "model": { + "name": "test_pytorch_train", + "classifier_out_dim": 3, + "reliability_out_dim": 0, + "embedding": { + "input_type": "translated", + "use_embedding_layer": True, + "vocab_size": 65, + "embedding_size": 4, + }, + "string_processor": { + "data_format": data_format, + "crop_size": seq_length, + "shuffle": False, + "mutate": False, + "shuffle_frames": False, + }, + "representation_learner": { + "hidden_layers": [], + "pooling": "average", + }, + "classifier": { + "input_shape": 4, + "hidden_layers": [{"name": "dense", "config": {"units": 3}}], + }, + }, + "training": { + "batch_size": 4, + "classifier_epochs": 1, + "classifier_dir": str(tmp_path / "checkpoints" / "classifier"), + "optimizer": "adam", + "optimizer_params": {"lr": 1e-3}, + "fragment_classifier_data": { + "train": [{"path": [str(train_path)]}], + "validation": [{"path": [str(val_path)]}], + }, + "model_saving": { + "path": str(tmp_path / "model"), + }, + }, + } + + +seq_length = 50 + + +def test_train_fragment_core_runs(tmp_path): + """``train_fragment_core`` should complete a single epoch without error.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + +def test_train_fragment_core_only_save(tmp_path): + """``train_fragment_core`` should save a checkpoint without training.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=True, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + checkpoint_path = tmp_path / "model" / "classifier.pt" + assert checkpoint_path.exists() + + +def test_train_fragment_core_only_classification_head(tmp_path): + """``train_fragment_core`` should freeze the representation learner when asked.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=True, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) From 53a2ca1e1821e519ffb60a7eecbc08bcb4eed78b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 11:55:16 +0200 Subject: [PATCH 32/64] feat(pytorch): add PyTorch inference runner --- src/jaeger/inference/pytorch/__init__.py | 1 + src/jaeger/inference/pytorch/runner.py | 73 ++++++++++++++++++++ tests/unit/inference/pytorch/test_runner.py | 74 +++++++++++++++++++++ 3 files changed, 148 insertions(+) create mode 100644 src/jaeger/inference/pytorch/__init__.py create mode 100644 src/jaeger/inference/pytorch/runner.py create mode 100644 tests/unit/inference/pytorch/test_runner.py diff --git a/src/jaeger/inference/pytorch/__init__.py b/src/jaeger/inference/pytorch/__init__.py new file mode 100644 index 0000000..4f6e2b1 --- /dev/null +++ b/src/jaeger/inference/pytorch/__init__.py @@ -0,0 +1 @@ +"""PyTorch inference backend for Jaeger.""" diff --git a/src/jaeger/inference/pytorch/runner.py b/src/jaeger/inference/pytorch/runner.py new file mode 100644 index 0000000..c50a302 --- /dev/null +++ b/src/jaeger/inference/pytorch/runner.py @@ -0,0 +1,73 @@ +from pathlib import Path +from typing import Any, Dict, Optional, Union + +import torch + +from jaeger.nnlib.pytorch.builder import ModelBuilder + + +class PyTorchInferenceRunner: + """Load and run a PyTorch Jaeger model for inference.""" + + def __init__( + self, + config: Dict[str, Any], + checkpoint_path: Optional[Union[str, Path]] = None, + device: Optional[torch.device] = None, + ): + self.config = config + self.device = device or torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + builder = ModelBuilder(config) + self.models = builder.build_fragment_classifier() + self.model = self.models["jaeger_model"] + if checkpoint_path is not None: + self.load_checkpoint(checkpoint_path) + self.model.to(self.device) + self.model.eval() + + def load_checkpoint(self, path: Union[str, Path]): + path = Path(path) + try: + state = torch.load(path, map_location=self.device, weights_only=True) + except TypeError: + state = torch.load(path, map_location=self.device, weights_only=False) + if "model_state_dict" in state: + self.model.load_state_dict(state["model_state_dict"]) + else: + self.model.load_state_dict(state) + + def predict( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + batch_size: int = 32, + ) -> Dict[str, torch.Tensor]: + self.model.eval() + outputs = [] + with torch.no_grad(): + for i in range(0, len(x), batch_size): + batch_x = x[i : i + batch_size].to(self.device) + batch_mask = ( + None if mask is None else mask[i : i + batch_size].to(self.device) + ) + out = self.model(batch_x, mask=batch_mask) + outputs.append({k: v.cpu() for k, v in out.items()}) + + return {k: torch.cat([o[k] for o in outputs], dim=0) for k in outputs[0]} + + def predict_from_dataset( + self, dataset: torch.utils.data.Dataset, batch_size: int = 32 + ) -> Dict[str, torch.Tensor]: + loader = torch.utils.data.DataLoader( + dataset, batch_size=batch_size, shuffle=False + ) + outputs = [] + with torch.no_grad(): + for batch in loader: + x = batch[0].to(self.device) + mask = batch[2].to(self.device) if len(batch) > 2 else None + out = self.model(x, mask=mask) + outputs.append({k: v.cpu() for k, v in out.items()}) + return {k: torch.cat([o[k] for o in outputs], dim=0) for k in outputs[0]} diff --git a/tests/unit/inference/pytorch/test_runner.py b/tests/unit/inference/pytorch/test_runner.py new file mode 100644 index 0000000..43f4af3 --- /dev/null +++ b/tests/unit/inference/pytorch/test_runner.py @@ -0,0 +1,74 @@ +import tempfile + +import torch + +from jaeger.inference.pytorch.runner import PyTorchInferenceRunner +from jaeger.nnlib.pytorch.builder import ModelBuilder + + +def _make_config(tmp_path): + return { + "model": { + "name": "test_model", + "classifier_out_dim": 3, + "reliability_out_dim": 1, + "class_label_map": [ + {"class": "phage", "label": 1}, + {"class": "bacteria", "label": 0}, + ], + "embedding": { + "input_type": "translated", + "use_embedding_layer": False, + "embedding_size": 32, + "input_shape": [6, None], + "vocab_size": 65, + "codon_depth": 1, + }, + "string_processor": {"codon": "CODON", "codon_id": "CODON_ID"}, + "representation_learner": { + "hidden_layers": [ + { + "name": "masked_conv1d", + "config": {"filters": 16, "kernel_size": 3, "padding": "same"}, + } + ], + "pooling": "average", + }, + "classifier": { + "input_shape": 16, + "hidden_layers": [{"name": "dense", "config": {"units": 3}}], + }, + }, + "training": { + "batch_size": 2, + "optimizer": "adam", + "optimizer_params": {"lr": 1e-3}, + }, + } + + +def test_inference_runner_predict(): + with tempfile.TemporaryDirectory() as tmpdir: + config = _make_config(tmpdir) + runner = PyTorchInferenceRunner(config, device=torch.device("cpu")) + x = torch.randint(0, 65, (4, 6, 50)) + mask = torch.ones(4, 6, 50, dtype=torch.bool) + out = runner.predict(x, mask=mask, batch_size=2) + assert out["prediction"].shape == (4, 3) + + +def test_inference_runner_from_checkpoint(tmp_path): + config = _make_config(tmp_path) + builder = ModelBuilder(config) + models = builder.build_fragment_classifier() + checkpoint_path = tmp_path / "model.pt" + torch.save( + {"model_state_dict": models["jaeger_model"].state_dict()}, checkpoint_path + ) + + runner = PyTorchInferenceRunner( + config, checkpoint_path=checkpoint_path, device=torch.device("cpu") + ) + x = torch.randint(0, 65, (2, 6, 50)) + out = runner.predict(x, batch_size=2) + assert out["prediction"].shape == (2, 3) From 4833ba2f78379ee9cbc5a35aa6054bed49aefe3b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 12:08:26 +0200 Subject: [PATCH 33/64] feat(pytorch): update jaeger predict command for PyTorch --- src/jaeger/commands/predict.py | 716 ++++++++++---------- tests/unit/commands/test_predict_pytorch.py | 106 +++ 2 files changed, 461 insertions(+), 361 deletions(-) create mode 100644 tests/unit/commands/test_predict_pytorch.py diff --git a/src/jaeger/commands/predict.py b/src/jaeger/commands/predict.py index f6bcc02..3110b63 100644 --- a/src/jaeger/commands/predict.py +++ b/src/jaeger/commands/predict.py @@ -1,75 +1,161 @@ -import traceback +"""PyTorch-based ``jaeger predict`` command.""" + +from __future__ import annotations + import sys -import psutil import time -from importlib.resources import files +import traceback from importlib.metadata import version from pathlib import Path -import tensorflow as tf -from jaeger.nnlib.inference import InferModel, TFLiteInferModel, ONNXEngine -from jaeger.seqops.io import fragment_generator -from jaeger.utils.gpu import get_device_name -from jaeger.utils.termini import scan_for_terminal_repeats -from jaeger.utils.fs import validate_fasta_entries -from jaeger.utils.misc import json_to_dict, AvailableModels, get_model_id +import numpy as np +import torch + +from jaeger.inference.pytorch.runner import PyTorchInferenceRunner +from jaeger.postprocess.collect import pred_to_dict, write_output +from jaeger.seqops.io import fragment_generator, validate_fasta_entries +from jaeger.seqops.maps import CODON_ID, CODONS from jaeger.utils.logging import description, get_logger +from jaeger.utils.misc import load_model_config +from jaeger.utils.termini import scan_for_terminal_repeats -# from jaeger.utils.tandem import split_fasta_with_pyfastx, run_batch, merge_masked_files GB_BYTES = 1024**3 -def run_core(**kwargs): - current_process = psutil.Process() +def _resolve_codon_value(value): + """Resolve a codon table configuration value to a list.""" + if isinstance(value, str): + if value == "CODONS": + return CODONS + if value == "CODON_ID": + return CODON_ID + raise ValueError(f"Unsupported codon table name: {value!r}") + return value + + +def _build_codon_table(codon_cfg, codon_id_cfg): + """Build a codon -> integer ID lookup from config values.""" + codons = _resolve_codon_value(codon_cfg) + codon_ids = _resolve_codon_value(codon_id_cfg) + return dict(zip(codons, codon_ids)) + + +def _translate_fragment(seq: str, codon_table: dict, crop_size: int) -> torch.Tensor: + """Translate a DNA fragment into a ``(6, L)`` codon index tensor. + + The six reading frames are produced using the same offset/striding logic as + the training-time raw sequence processor. + """ + seq = seq.upper()[:crop_size] + comp = {"A": "T", "T": "A", "G": "C", "C": "G", "N": "N"} + rev = "".join(comp.get(base, "N") for base in reversed(seq)) + + def _codon_ids(s: str): + return [codon_table.get(s[i : i + 3], -1) for i in range(len(s) - 2)] + + forward = _codon_ids(seq) + reverse = _codon_ids(rev) - USER_MODEL_PATH = kwargs.get("model_path") - CONFIG_PATH = kwargs.get("config") or files("jaeger.data") / "config.json" + frames = [ + forward[offset::3] for offset in range(3) + ] + [reverse[offset::3] for offset in range(3)] - if not USER_MODEL_PATH: - # Use default model from config - model_name = kwargs.get("model") - model_id = get_model_id(model_name) + target_len = max(0, len(seq) // 3 - 1) + frames = [f[:target_len] + [-1] * (target_len - len(f)) for f in frames] - model_paths = json_to_dict(CONFIG_PATH).get("model_paths") - info = AvailableModels(path=model_paths).info + arr = np.array(frames, dtype=np.int64) + 1 # 0 is reserved for padding/masking + return torch.from_numpy(arr) - model_info = info[model_name] + +def _aggregate_predictions(outputs: dict) -> dict: + """Average each tensor in ``outputs`` over the batch dimension.""" + return {k: v.mean(dim=0) for k, v in outputs.items()} + + +def _append_outputs(container: dict, outputs: dict): + """Append CPU numpy arrays from a runner output dict to a container.""" + for k, v in outputs.items(): + container.setdefault(k, []).append(v.cpu().numpy()) + + +def _concat_outputs(container: dict) -> dict: + """Concatenate lists of numpy arrays per output key.""" + return {k: np.concatenate(vs, axis=0) for k, vs in container.items()} + + +def _build_class_map(config: dict) -> dict: + """Build a class_map compatible with ``jaeger.postprocess.collect``.""" + class_label_map = config.get("model", {}).get("class_label_map", []) + classes = [entry["class"] for entry in class_label_map] + indices = [entry["label"] for entry in class_label_map] + return {"num_classes": len(classes), "class": classes, "index": indices} + + +def run_core(**kwargs): + """Run the PyTorch Jaeger inference pipeline.""" + try: + import psutil + + current_process = psutil.Process() + except Exception: + current_process = None + + model_path = kwargs.get("model_path") + config_path = kwargs.get("config") + checkpoint_path = kwargs.get("checkpoint") + + if model_path: + model_path = Path(model_path) + config_file = None + for name in ("config.yaml", "config.yml", "config.json"): + candidate = model_path / name + if candidate.exists(): + config_file = candidate + break + if config_file is None: + raise FileNotFoundError(f"No config file found in {model_path}") + config = load_model_config(config_file) + checkpoint_path = model_path / "model.pt" + model_name = config.get("model", {}).get("name", model_path.name) + elif config_path and checkpoint_path: + config = load_model_config(Path(config_path)) + checkpoint_path = Path(checkpoint_path) + model_name = config.get("model", {}).get("name", "jaeger_pytorch") else: - # Use provided model path - info = AvailableModels(path=USER_MODEL_PATH).info - model_paths = USER_MODEL_PATH - model_name = next(iter(info)) - model_id = get_model_id(model_name) - model_info = info[model_name] - - MEMORY_LIMIT = 1024 * kwargs.get("mem", 4) - THREADS = kwargs.get("workers") - input_file_path = Path(kwargs.get("input")) - input_file = input_file_path.name - file_base = input_file_path.stem - # INPUT_FILE_MASKED = input_file_path.with_name(file_base + "_masked" + input_file_path.suffix) + raise NotImplementedError( + "Default TensorFlow model loading is not supported in PyTorch mode. " + "Please provide --model_path (with config.yaml and model.pt) or both " + "--config and --checkpoint." + ) - OUTPUT_DIR = Path(kwargs.get("output")) / model_id - OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - # CHUNKS_DIR = OUTPUT_DIR / "chunks" - # MAKSED_DIR = OUTPUT_DIR / "masked" + input_file_path = Path(kwargs["input"]) + file_base = input_file_path.stem + model_id = model_name.replace(" ", "_") + output_dir = Path(kwargs["output"]) / model_id + output_dir.mkdir(parents=True, exist_ok=True) log_file = Path(f"{file_base}_jaeger.log") - logger = get_logger(OUTPUT_DIR, log_file, level=kwargs.get("verbose")) + logger = get_logger(output_dir, log_file, level=kwargs.get("verbose", 1)) logger.info( - description(version("jaeger-bio")) + "\n{:-^80}".format("validating parameters") + description(version("jaeger-bio")) + + "\n{:-^80}".format("validating parameters") ) - logger.debug(info) - logger.debug(model_info) + logger.debug(config) + + fsize = kwargs.get("fsize", 2000) + stride = kwargs.get("stride", 2000) + threads = kwargs.get("workers", 4) + batch_size = kwargs.get("batch", 96) + try: - num = validate_fasta_entries(str(input_file_path), min_len=kwargs.get("fsize")) + num = validate_fasta_entries(str(input_file_path), min_len=fsize) except Exception as e: logger.error(e) logger.debug(traceback.format_exc()) sys.exit(1) - output_table_path = OUTPUT_DIR / f"{file_base}.tsv" - output_phage_table_path = OUTPUT_DIR / f"{file_base}_phages.tsv" + output_table_path = output_dir / f"{file_base}.tsv" + output_phage_table_path = output_dir / f"{file_base}_phages.tsv" if output_table_path.exists() and not kwargs.get("overwrite"): logger.error( @@ -77,368 +163,276 @@ def run_core(**kwargs): ) sys.exit(1) - if not model_info["graph"].exists(): - logger.error(f"could not find model graph. please check {model_paths}") - sys.exit(1) - tf.config.threading.set_inter_op_parallelism_threads(THREADS) - tf.config.threading.set_intra_op_parallelism_threads(THREADS) - tf.config.set_soft_device_placement(True) - gpus = tf.config.list_physical_devices("GPU") - mode = None + for opt in ("quantized", "xla", "onnx", "int8"): + if kwargs.get(opt): + logger.warning( + f"--{opt} is not supported by the PyTorch inference path and will be ignored." + ) + precision = kwargs.get("precision", "fp32") + if precision and precision != "fp32": + logger.warning( + f"--precision {precision} is not supported and will be ignored." + ) - if kwargs.get("cpu"): + if kwargs.get("cpu") or not torch.cuda.is_available(): + device = torch.device("cpu") mode = "CPU" - tf.config.set_visible_devices([], "GPU") - logger.info("CPU only mode selected") - elif gpus: - mode = "GPU" - tf.config.set_visible_devices([gpus[kwargs.get("physicalid")]], "GPU") - - # Set mixed precision policy if requested - precision = kwargs.get("precision", "fp32") - if precision == "fp16": - tf.keras.mixed_precision.set_global_policy("mixed_float16") - logger.info("GPU precision: mixed_float16 (FP16 compute, FP32 variables)") - elif precision == "bf16": - tf.keras.mixed_precision.set_global_policy("mixed_bfloat16") - logger.info("GPU precision: mixed_bfloat16 (BF16 compute, FP32 variables)") - else: - logger.info("GPU precision: float32 (FP32)") - - try: - tf.config.set_logical_device_configuration( - gpus[kwargs.get("physicalid")], - [ - tf.config.LogicalDeviceConfiguration( - memory_limit=MEMORY_LIMIT, experimental_device_ordinal=10 - ) - ], + if not kwargs.get("cpu") and not torch.cuda.is_available(): + logger.warning( + "could not find a GPU on the system. " + "For optimal performance run Jaeger on a GPU." ) - except Exception as e: - logger.error(f"an error {e} occurred during virtual device initialization ") - logger.debug(traceback.format_exc()) else: - mode = "CPU" - logger.warning( - "could not find a GPU on the system. For optimal performance run Jaeger on a GPU." - ) + device = torch.device("cuda") + mode = "GPU" - logger.info(f"tensorflow: {version('tensorflow')}") - logger.info(f"input file: {input_file}") + logger.info(f"pytorch: {version('torch')}") + logger.info(f"input file: {input_file_path.name}") logger.info(f"log file: {log_file.name}") - logger.info(f"outpath: {OUTPUT_DIR.resolve()}") - logger.info(f"fragment size: {kwargs.get('fsize')}") - logger.info(f"stride: {kwargs.get('stride')}") - logger.info(f"batch size: {kwargs.get('batch')}") + logger.info(f"outpath: {output_dir.resolve()}") + logger.info(f"fragment size: {fsize}") + logger.info(f"stride: {stride}") + logger.info(f"batch size: {batch_size}") logger.info(f"mode: {mode}") logger.info(f"model: {model_id}") - logger.info(f"avail mem: {psutil.virtual_memory().available / (GB_BYTES):.2f}GB") - logger.info( - f"intra threads: {tf.config.threading.get_intra_op_parallelism_threads()}" - ) - logger.info( - f"inter threads: {tf.config.threading.get_inter_op_parallelism_threads()}" - ) - logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") - logger.info( - f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB ({current_process.memory_percent():.2f}%)" - ) - device = tf.config.list_logical_devices(mode) - device_names = [get_device_name(i) for i in device] - logger.debug(f"{device}, {device_names}") - if len(device) > 1: - logger.info(f"Using MirroredStrategy {device_names}") - strategy = tf.distribute.MirroredStrategy(device_names) - else: - logger.info(f"Using OneDeviceStrategy {device_names}") - strategy = tf.distribute.OneDeviceStrategy(device_names[0]) - # # 1. Split the input FASTA - # chunk_files = split_fasta_with_pyfastx(input_file_path, CHUNKS_DIR, chunks=16) - # # Or: chunk_files = split_fasta_with_pyfastx(input_fasta, chunks_dir, chunks=10) - - # # 2. Run TRF on chunks in parallel - # masked_files = run_batch(chunk_files, out_dir=MAKSED_DIR, n_threads=16) + if current_process is not None: + logger.info( + f"avail mem: {psutil.virtual_memory().available / GB_BYTES:.2f}GB" + ) + logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") + logger.info( + f"wall time(s) : {time.time() - current_process.create_time():.2f}" + ) + logger.info( + f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB " + f"({current_process.memory_percent():.2f}%)" + ) - # # 3. Merge all masked files into one - # merge_masked_files(masked_files, INPUT_FILE_MASKED) + torch.set_num_threads(threads) - # # 4. Cleanup temporary directories - # shutil.rmtree(CHUNKS_DIR) - # shutil.rmtree(MAKSED_DIR) + runner = PyTorchInferenceRunner( + config, checkpoint_path=checkpoint_path, device=device + ) + class_map = _build_class_map(config) term_repeats = scan_for_terminal_repeats( - # file_path=str(INPUT_FILE_MASKED), file_path=str(input_file_path), num=num, - workers=THREADS, - fsize=kwargs.get("fsize"), - ) - - # Select inference engine based on options - quantized_mode = kwargs.get("quantized") - use_xla = kwargs.get("xla", False) - use_onnx = kwargs.get("onnx", False) - use_onnx_int8 = kwargs.get("int8", False) - - if quantized_mode: - # Look for quantized model in the same directory as the original model - graph_dir = Path(model_info["graph"]) - quantized_dir = graph_dir.parent / f"{model_name}_{quantized_mode}" - tflite_path = quantized_dir / f"{model_name}_{quantized_mode}.tflite" - - if not tflite_path.exists(): - logger.error( - f"Quantized model not found at {tflite_path}. " - f"Run 'jaeger utils quantize -m {model_name} -o {graph_dir.parent} --mode {quantized_mode}' first." - ) - sys.exit(1) - - logger.info(f"Using quantized model: {tflite_path}") - model_info_tflite = model_info.copy() - model_info_tflite["tflite"] = tflite_path - model = TFLiteInferModel(model_info_tflite) - elif use_onnx: - # Look for ONNX model (FP32 or INT8) - graph_dir = Path(model_info["graph"]) - if use_onnx_int8: - onnx_dir = graph_dir.parent / f"{model_name}_onnx_int8" - onnx_path = onnx_dir / f"{model_name}_int8.onnx" - if not onnx_path.exists(): - logger.error( - f"INT8 ONNX model not found at {onnx_path}. " - f"Run 'jaeger utils convert-graph -m {model_name} -o {graph_dir.parent} --mode onnx --int8' first." - ) - sys.exit(1) - logger.info(f"Using INT8 ONNX model: {onnx_path}") - else: - onnx_dir = graph_dir.parent / f"{model_name}_onnx" - onnx_path = onnx_dir / f"{model_name}.onnx" - - if not onnx_path.exists(): - logger.error( - f"ONNX model not found at {onnx_path}. " - f"Run 'jaeger utils convert-graph -m {model_name} -o {graph_dir.parent} --mode onnx' first." - ) - sys.exit(1) - - logger.info(f"Using ONNX model: {onnx_path}") - - model_info_onnx = model_info.copy() - model_info_onnx["onnx"] = onnx_path - - model = ONNXEngine(model_info_onnx) - else: - if use_xla: - logger.info( - "Using XLA-compiled inference (first batch may be slow due to compilation)" - ) - model = InferModel(model_info, use_xla=use_xla) - - string_processor_config = model.string_processor_config - input_dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - # str(INPUT_FILE_MASKED), - file_path=str(input_file_path), - no_progress=False, - fragsize=kwargs.get("fsize"), - stride=kwargs.get("stride"), - num=num, - ), - output_signature=(tf.TensorSpec(shape=(), dtype=tf.string)), + workers=threads, + fsize=fsize, ) - from jaeger.seqops.encode import process_string_inference - - # from icecream import ic - # ic(string_processor_config) - idataset = ( - input_dataset.map( - process_string_inference( - codons=string_processor_config.get("codon"), - codon_num=string_processor_config.get("codon_id"), - codon_depth=string_processor_config.get("codon_depth"), - ngram_width=string_processor_config.get("ngram_width"), - seq_onehot=string_processor_config.get("seq_onehot"), - crop_size=string_processor_config.get("crop_size"), - input_type=string_processor_config.get("input_type"), - masking=string_processor_config.get("masking"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle=string_processor_config.get("shuffle"), - ), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .batch(kwargs.get("batch"), num_parallel_calls=tf.data.AUTOTUNE) - .prefetch(25) + sp_config = config.get("model", {}).get("string_processor", {}) + crop_size = sp_config.get("crop_size", fsize) + codon_table = _build_codon_table( + sp_config.get("codon", CODONS), sp_config.get("codon_id", CODON_ID) ) - with strategy.scope(): - try: - logger.info("starting model inference") - y_pred = model.predict(idataset, no_progress=True) - except Exception as e: - logger.debug(traceback.format_exc()) - logger.error( - f"an error {e} occured during inference on {'|'.join(device_names)}! check {log_file} for traceback." - ) - sys.exit(1) + meta = { + "headers": [], + "index": [], + "contig_end": [], + "window_idx": [], + "length": [], + "c_count": [], + "g_count": [], + "a_count": [], + "t_count": [], + "gc_skew": [], + } + inputs: list[torch.Tensor] = [] + outputs_container: dict = {} + + for frag_str in fragment_generator( + file_path=str(input_file_path), + fragsize=fsize, + stride=stride, + num=num, + no_progress=True, + ): + parts = frag_str.split(",") + seq = parts[0] + meta["headers"].append(parts[1]) + meta["index"].append(int(parts[2])) + meta["contig_end"].append(int(parts[3])) + meta["window_idx"].append(int(parts[4])) + meta["length"].append(int(parts[5])) + meta["g_count"].append(float(parts[6])) + meta["c_count"].append(float(parts[7])) + meta["a_count"].append(float(parts[8])) + meta["t_count"].append(float(parts[9])) + meta["gc_skew"].append(float(parts[10])) + + inputs.append(_translate_fragment(seq, codon_table, crop_size)) + if len(inputs) >= batch_size: + batch_x = torch.stack(inputs) + out = runner.predict(batch_x, batch_size=batch_size) + _append_outputs(outputs_container, out) + inputs.clear() + + if inputs: + batch_x = torch.stack(inputs) + out = runner.predict(batch_x, batch_size=batch_size) + _append_outputs(outputs_container, out) + inputs.clear() + + if not outputs_container: + logger.error("no fragments were generated for inference") + sys.exit(1) - from jaeger.postprocess.collect import pred_to_dict, write_output + y_pred = _concat_outputs(outputs_container) + y_pred["meta_0"] = np.array(meta["headers"], dtype=str) + y_pred["meta_1"] = np.array(meta["index"], dtype=np.int32) + y_pred["meta_2"] = np.array(meta["contig_end"], dtype=np.int32) + y_pred["meta_3"] = np.array(meta["window_idx"], dtype=np.int32) + y_pred["meta_4"] = np.array(meta["length"], dtype=np.int32) + y_pred["meta_5"] = np.array(meta["c_count"], dtype=np.float32) + y_pred["meta_6"] = np.array(meta["g_count"], dtype=np.float32) + y_pred["meta_7"] = np.array(meta["a_count"], dtype=np.float32) + y_pred["meta_8"] = np.array(meta["t_count"], dtype=np.float32) + y_pred["meta_9"] = np.array(meta["gc_skew"], dtype=np.float32) - if kwargs.get("getalllabels"): - pass + try: data, data_full = pred_to_dict( y_pred, - class_map=model.class_map, - fsize=kwargs.get("fsize"), + class_map=class_map, + fsize=fsize, term_repeats=term_repeats, ) - # print(len(data['headers'])) num_written = write_output( data, - labels=model.class_map.get("class"), - indices=model.class_map.get("index"), + labels=class_map["class"], + indices=class_map["index"], output_table_path=output_table_path, output_phage_table_path=output_phage_table_path, reliability_cutoff=kwargs.get("rc", 0.5), phage_score=kwargs.get("pc", 1), ) - logger.info(f"processed {num_written}/{num} sequences") + except Exception as e: + logger.error(f"an error {e} occurred during postprocessing") + logger.debug(traceback.format_exc()) + sys.exit(1) - # --- Prophage extraction --- - if kwargs.get("prophage"): - from jaeger.postprocess.prophages import ( - logits_to_df_v2, - plot_scores, - plot_scores_linear, - prophage_report, - segment, - ) + if kwargs.get("prophage"): + from jaeger.postprocess.prophages import ( + logits_to_df_v2, + plot_scores, + plot_scores_linear, + prophage_report, + segment, + ) - try: - if logits_df := logits_to_df_v2( - class_map=model.class_map, - cmdline_kwargs=kwargs, - headers=data_full["headers"], - predictions=data_full["predictions"], - lengths=data_full["lengths"], - gc_skews=data_full["gc_skews"], - gcs=data_full["gcs"], - ): - logger.info("identifying prophages") - pro_dir = OUTPUT_DIR / f"{file_base}_prophages" - plots_dir = pro_dir / "plots" - for d in [pro_dir, plots_dir]: - d.mkdir(parents=True, exist_ok=True) - - phage_cord = segment( + try: + if logits_df := logits_to_df_v2( + class_map=class_map, + cmdline_kwargs=kwargs, + headers=data_full["headers"], + predictions=data_full["predictions"], + lengths=data_full["lengths"], + gc_skews=data_full["gc_skews"], + gcs=data_full["gcs"], + ): + logger.info("identifying prophages") + pro_dir = output_dir / f"{file_base}_prophages" + plots_dir = pro_dir / "plots" + for d in [pro_dir, plots_dir]: + d.mkdir(parents=True, exist_ok=True) + + phage_cord = segment( + logits_df, + outdir=plots_dir, + cutoff_length=kwargs.get("lc", 500_000), + sensitivity=kwargs.get("sensitivity", 1.5), + identifier="phage", + ) + plot_type = kwargs.get("plot_type", "circular") + config_labels = { + i: c for i, c in enumerate(class_map.get("class", [])) + } + if plot_type in ("circular", "both"): + plot_scores( logits_df, + config={"all_labels": config_labels}, + model=model_name, + fsize=fsize, + infile_base=file_base, outdir=plots_dir, - cutoff_length=kwargs.get("lc"), - sensitivity=kwargs.get("sensitivity"), - identifier="phage", + phage_cordinates=phage_cord, ) - plot_type = kwargs.get("plot_type", "circular") - if plot_type in ("circular", "both"): - plot_scores( - logits_df, - config={ - "all_labels": { - i: c - for i, c in enumerate( - model.class_map.get("class", []) - ) - } - }, - model=model_name, - fsize=kwargs.get("fsize"), - infile_base=file_base, - outdir=plots_dir, - phage_cordinates=phage_cord, - ) - if plot_type in ("linear", "both"): - plot_scores_linear( - logits_df, - config={ - "all_labels": { - i: c - for i, c in enumerate( - model.class_map.get("class", []) - ) - } - }, - model=model_name, - fsize=kwargs.get("fsize"), - infile_base=file_base, - outdir=plots_dir, - phage_cordinates=phage_cord, - ) - prophage_report( - fsize=kwargs.get("fsize"), - filehandle=str(input_file_path), - prophage_cordinates=phage_cord, - outdir=pro_dir, + if plot_type in ("linear", "both"): + plot_scores_linear( + logits_df, + config={"all_labels": config_labels}, + model=model_name, + fsize=fsize, + infile_base=file_base, + outdir=plots_dir, + phage_cordinates=phage_cord, ) - else: - logger.info("no prophage regions found") - except Exception as e: - logger.error( - f"an error {e} occurred during the prophage prediction step" + prophage_report( + fsize=fsize, + filehandle=str(input_file_path), + prophage_cordinates=phage_cord, + outdir=pro_dir, ) - logger.debug(traceback.format_exc()) + else: + logger.info("no prophage regions found") + except Exception as e: + logger.error( + f"an error {e} occurred during the prophage prediction step" + ) + logger.debug(traceback.format_exc()) - # --- Write phage sequences to FASTA --- - if kwargs.get("getsequences"): - from jaeger.postprocess.collect import write_fasta_from_results + if kwargs.get("getsequences"): + from jaeger.postprocess.collect import write_fasta_from_results - output_fasta_file = f"{file_base}_phages_jaeger.fasta" - output_fasta_file_path = OUTPUT_DIR / output_fasta_file + output_fasta_file_path = output_dir / f"{file_base}_phages_jaeger.fasta" + if output_phage_table_path.exists(): write_fasta_from_results( - input_fasta=input_file_path, - output_tsv=output_phage_table_path, - output_fasta=output_fasta_file_path, + input_fasta=str(input_file_path), + output_tsv=str(output_phage_table_path), + output_fasta=str(output_fasta_file_path), ) - logger.info(f"{output_fasta_file} created") - - # --- Write window-wise scores + metadata to NPZ --- - if kwargs.get("window_scores"): - import numpy as np + logger.info(f"{output_fasta_file_path.name} created") + else: + logger.info("no phage sequences to write") - output_scores = f"{file_base}_window_scores.npz" - output_scores_path = OUTPUT_DIR / output_scores - logger.info( - f"writing window-wise scores and metadata to {output_scores_path}" - ) - np.savez( - output_scores_path, - headers=data_full["headers"], - lengths=data_full["lengths"], - predictions=np.array(data_full["predictions"], dtype=object), - gc_skews=np.array(data_full["gc_skews"], dtype=object), - gcs=np.array(data_full["gcs"], dtype=object), - ) - logger.info(f"{output_scores} created") + if kwargs.get("window_scores"): + output_scores_path = output_dir / f"{file_base}_window_scores.npz" + logger.info( + f"writing window-wise scores and metadata to {output_scores_path}" + ) + np.savez( + output_scores_path, + headers=data_full["headers"], + lengths=data_full["lengths"], + predictions=np.array(data_full["predictions"], dtype=object), + gc_skews=np.array(data_full["gc_skews"], dtype=object), + gcs=np.array(data_full["gcs"], dtype=object), + ) + logger.info(f"{output_scores_path.name} created") + if current_process is not None: logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") logger.info( - f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB ({current_process.memory_percent():.2f}%)" + f"wall time(s) : {time.time() - current_process.create_time():.2f}" + ) + logger.info( + f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB " + f"({current_process.memory_percent():.2f}%)" ) - import numpy as np - if "embedding" in y_pred: - np.savez( - OUTPUT_DIR / f"{file_base}_embedding.npz", - embedding=y_pred["embedding"], - headers=y_pred["meta_0"], - ) - if "nmd" in y_pred: - np.savez( - OUTPUT_DIR / f"{file_base}_nmd.npz", - embedding=y_pred["nmd"], - headers=y_pred["meta_0"], - ) - # INPUT_FILE_MASKED.unlink() + if "embedding" in y_pred: + np.savez( + output_dir / f"{file_base}_embedding.npz", + embedding=y_pred["embedding"], + headers=y_pred["meta_0"], + ) + if "nmd" in y_pred: + np.savez( + output_dir / f"{file_base}_nmd.npz", + embedding=y_pred["nmd"], + headers=y_pred["meta_0"], + ) diff --git a/tests/unit/commands/test_predict_pytorch.py b/tests/unit/commands/test_predict_pytorch.py new file mode 100644 index 0000000..276d556 --- /dev/null +++ b/tests/unit/commands/test_predict_pytorch.py @@ -0,0 +1,106 @@ +"""Tests for the PyTorch-based ``jaeger predict`` command.""" + +from __future__ import annotations + +import random + +import torch +import yaml + +from jaeger.commands.predict import run_core +from jaeger.nnlib.pytorch.builder import ModelBuilder + + +def _build_config(tmp_path, seq_length: int): + """Create a minimal PyTorch model config.""" + return { + "model": { + "name": "test_predict_model", + "classifier_out_dim": 3, + "reliability_out_dim": 0, + "class_label_map": [ + {"class": "chromosome", "label": 0}, + {"class": "virus", "label": 1}, + {"class": "plasmid", "label": 2}, + ], + "embedding": { + "input_type": "translated", + "use_embedding_layer": True, + "vocab_size": 65, + "embedding_size": 4, + }, + "string_processor": { + "codon": "CODONS", + "codon_id": "CODON_ID", + "crop_size": seq_length, + "seq_onehot": False, + "shuffle": False, + "mutate": False, + "shuffle_frames": False, + }, + "representation_learner": { + "hidden_layers": [], + "pooling": "average", + }, + "classifier": { + "input_shape": 4, + "hidden_layers": [{"name": "dense", "config": {"units": 3}}], + }, + }, + "training": { + "batch_size": 2, + "optimizer": "adam", + "optimizer_params": {"lr": 1e-3}, + }, + } + + +def _write_fasta(path, records: dict[str, str]): + """Write a small FASTA file from a header -> sequence mapping.""" + with open(path, "w") as fh: + for header, seq in records.items(): + fh.write(f">{header}\n{seq}\n") + + +def test_predict_pytorch_run_core(tmp_path): + """``run_core`` should produce TSV outputs for a tiny FASTA file.""" + seq_length = 100 + config = _build_config(tmp_path, seq_length) + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(config)) + + builder = ModelBuilder(config) + models = builder.build_fragment_classifier() + # Bias the classifier toward the virus class so the phage TSV is non-empty. + with torch.no_grad(): + models["jaeger_model"].classification_head.net[-1].bias[:] = torch.tensor( + [-10.0, 10.0, -10.0] + ) + checkpoint_path = tmp_path / "model.pt" + torch.save( + {"model_state_dict": models["jaeger_model"].state_dict()}, + checkpoint_path, + ) + + fasta_path = tmp_path / "input.fasta" + bases = ["A", "C", "G", "T"] + seq = "".join(random.choices(bases, k=200)) # 200 bp, longer than fsize + _write_fasta(fasta_path, {"contig_1": seq, "contig_2": seq}) + + output_dir = tmp_path / "output" + run_core( + input=str(fasta_path), + output=str(output_dir), + config=str(config_path), + checkpoint=str(checkpoint_path), + cpu=True, + fsize=seq_length, + stride=50, + batch=2, + overwrite=True, + pc=-100, + ) + + out_model_dir = output_dir / config["model"]["name"] + assert (out_model_dir / "input.tsv").exists() + assert (out_model_dir / "input_phages.tsv").exists() From 38c62c23315f2ceaa9390c65aa70ea1768042425 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 12:14:13 +0200 Subject: [PATCH 34/64] feat(pytorch): update SLURM scripts and Singularity container for PyTorch --- singularity/jaeger_singularity.def | 18 +++---- slurm/axial_500bp.slurm | 8 +-- slurm/baseline_500bp.slurm | 8 +-- ...rt_tfrecords.slurm => convert_numpy.slurm} | 36 +++++++------ slurm/crossframe_500bp.slurm | 8 +-- slurm/pytorch_ddp_500bp.slurm | 52 +++++++++++++++++++ 6 files changed, 87 insertions(+), 43 deletions(-) rename slurm/{convert_tfrecords.slurm => convert_numpy.slurm} (53%) create mode 100644 slurm/pytorch_ddp_500bp.slurm diff --git a/singularity/jaeger_singularity.def b/singularity/jaeger_singularity.def index 8f11012..1deb041 100644 --- a/singularity/jaeger_singularity.def +++ b/singularity/jaeger_singularity.def @@ -1,5 +1,5 @@ Bootstrap: docker -From: tensorflow/tensorflow:2.15.0-gpu # Build upon the latest TensorFlow GPU-enabled base image +From: pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime # PyTorch GPU-enabled base image %post # Update and install system dependencies @@ -9,17 +9,17 @@ From: tensorflow/tensorflow:2.15.0-gpu # Build upon the latest TensorFlow GPU-e python3-pip \ && rm -rf /var/lib/apt/lists/* - # Install Jaeger and other dependencies - pip install --root-user-action=ignore --no-cache-dir "jaeger-bio==1.26.2" # install jaeger without tensorflow support + # Install Jaeger with PyTorch/GPU extras + pip install --root-user-action=ignore --no-cache-dir "jaeger-bio[gpu]" %environment - # Set environment variables for TensorFlow GPU usage - export TF_FORCE_GPU_ALLOW_GROWTH=true - export NVIDIA_VISIBLE_DEVICES=all # Makes all GPUs visible to the container - export NVIDIA_DRIVER_CAPABILITIES=compute,utility # Required for GPU support in TensorFlow + # Make all GPUs visible to the container + export NVIDIA_VISIBLE_DEVICES=all + # Required for GPU support in PyTorch + export NVIDIA_DRIVER_CAPABILITIES=compute,utility %labels Author: "Yasas Wijesekara" - Version: "1.4" - Description: "Apptainer container with Jaeger and TensorFlow GPU support" + Version: "1.5" + Description: "Apptainer container with Jaeger and PyTorch GPU support" diff --git a/slurm/axial_500bp.slurm b/slurm/axial_500bp.slurm index 4ffc844..cde0e5c 100755 --- a/slurm/axial_500bp.slurm +++ b/slurm/axial_500bp.slurm @@ -9,17 +9,13 @@ #SBATCH --output=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/axial_%j.log #SBATCH --error=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/axial_%j.err -echo "=== Jaeger 500bp axial training (mixed precision + XLA) ===" +echo "=== Jaeger 500bp axial training (PyTorch) ===" echo "Running on: $(hostname)" echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" echo "Started at: $(date)" -export TF_FORCE_GPU_ALLOW_GROWTH=true export CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES -# Enable XLA for faster training -export TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" - # Job-specific install directory to avoid conflicts INSTALL_DIR=/mnt/beegfs/bioinf/wijesekara/jaeger/.local_${SLURM_JOB_ID}_axial rm -rf $INSTALL_DIR @@ -37,6 +33,6 @@ apptainer exec --nv \ --bind /mnt/beegfs/bioinf/wijesekara/jaeger/experiments:/data/experiments \ --bind $INSTALL_DIR:/home/wijesekara/.local \ /mnt/beegfs/bioinf/wijesekara/jaeger/jaeger_dev.sif \ - bash -c "cd /jaeger && pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . && jaeger train -c /jaeger/train_config/nn_config_500bp_axial_zeus.yaml --force --mixed_precision" + bash -c "cd /jaeger && pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . && jaeger train -c /jaeger/train_config/nn_config_500bp_axial_zeus.yaml --force" echo "Finished at: $(date)" diff --git a/slurm/baseline_500bp.slurm b/slurm/baseline_500bp.slurm index c296601..cabcea5 100755 --- a/slurm/baseline_500bp.slurm +++ b/slurm/baseline_500bp.slurm @@ -9,17 +9,13 @@ #SBATCH --output=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/baseline_%j.log #SBATCH --error=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/baseline_%j.err -echo "=== Jaeger 500bp baseline training (mixed precision + XLA) ===" +echo "=== Jaeger 500bp baseline training (PyTorch) ===" echo "Running on: $(hostname)" echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" echo "Started at: $(date)" -export TF_FORCE_GPU_ALLOW_GROWTH=true export CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES -# Enable XLA for faster training -export TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" - # Job-specific install directory to avoid conflicts INSTALL_DIR=/mnt/beegfs/bioinf/wijesekara/jaeger/.local_${SLURM_JOB_ID}_baseline rm -rf $INSTALL_DIR @@ -37,6 +33,6 @@ apptainer exec --nv \ --bind /mnt/beegfs/bioinf/wijesekara/jaeger/experiments:/data/experiments \ --bind $INSTALL_DIR:/home/wijesekara/.local \ /mnt/beegfs/bioinf/wijesekara/jaeger/jaeger_dev.sif \ - bash -c "cd /jaeger && pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . && jaeger train -c /jaeger/train_config/nn_config_500bp_baseline_zeus.yaml --force --mixed_precision" + bash -c "cd /jaeger && pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . && jaeger train -c /jaeger/train_config/nn_config_500bp_baseline_zeus.yaml --force" echo "Finished at: $(date)" diff --git a/slurm/convert_tfrecords.slurm b/slurm/convert_numpy.slurm similarity index 53% rename from slurm/convert_tfrecords.slurm rename to slurm/convert_numpy.slurm index 944f2c2..1620b52 100755 --- a/slurm/convert_tfrecords.slurm +++ b/slurm/convert_numpy.slurm @@ -1,19 +1,19 @@ #!/bin/bash -#SBATCH --job-name=convert_tfrecords +#SBATCH --job-name=convert_numpy #SBATCH --partition=gpu #SBATCH --nodelist=node030 #SBATCH --gpus=0 #SBATCH --cpus-per-task=32 #SBATCH --mem=64G #SBATCH --time=04:00:00 -#SBATCH --output=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/convert_tfrecords_%j.log -#SBATCH --error=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/convert_tfrecords_%j.err +#SBATCH --output=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/convert_numpy_%j.log +#SBATCH --error=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/convert_numpy_%j.err -echo "=== TFRecord Conversion ===" +echo "=== NumPy conversion (PyTorch training data) ===" echo "Started at: $(date)" echo "CPUs: $(nproc)" -# Run conversion inside container where TensorFlow is installed +# Run conversion inside container where PyTorch/NumPy are installed apptainer exec \ --bind /mnt/beegfs/bioinf/wijesekara/jaeger/Jaeger_current:/jaeger \ --bind /mnt/beegfs/bioinf/wijesekara/jaeger/data:/data \ @@ -21,19 +21,23 @@ apptainer exec \ bash -c " cd /jaeger pip install -e . --quiet 2>/dev/null - + echo 'Converting training data...' - python3 /jaeger/convert_to_tfrecord.py \ - --csv /data/train_shuffled.csv \ - --output /data/train_shuffled.tfrecord \ - --crop-size 500 - + jaeger utils optimize-data \ + -i /data/train_shuffled.csv \ + -o /data/train_shuffled_full.npz \ + --format numpy_full \ + --crop-size 500 \ + --num-classes 3 + echo 'Converting validation data...' - python3 /jaeger/convert_to_tfrecord.py \ - --csv /data/val_shuffled.csv \ - --output /data/val_shuffled.tfrecord \ - --crop-size 500 + jaeger utils optimize-data \ + -i /data/val_shuffled.csv \ + -o /data/val_shuffled_full.npz \ + --format numpy_full \ + --crop-size 500 \ + --num-classes 3 " echo "Finished at: $(date)" -ls -lh /mnt/beegfs/bioinf/wijesekara/jaeger/data/*.tfrecord +ls -lh /mnt/beegfs/bioinf/wijesekara/jaeger/data/*.npz diff --git a/slurm/crossframe_500bp.slurm b/slurm/crossframe_500bp.slurm index b3bdfc3..d31378c 100755 --- a/slurm/crossframe_500bp.slurm +++ b/slurm/crossframe_500bp.slurm @@ -9,17 +9,13 @@ #SBATCH --output=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/crossframe_%j.log #SBATCH --error=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/crossframe_%j.err -echo "=== Jaeger 500bp crossframe training (mixed precision + XLA) ===" +echo "=== Jaeger 500bp crossframe training (PyTorch) ===" echo "Running on: $(hostname)" echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" echo "Started at: $(date)" -export TF_FORCE_GPU_ALLOW_GROWTH=true export CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES -# Enable XLA for faster training -export TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" - # Job-specific install directory to avoid conflicts INSTALL_DIR=/mnt/beegfs/bioinf/wijesekara/jaeger/.local_${SLURM_JOB_ID}_crossframe rm -rf $INSTALL_DIR @@ -37,6 +33,6 @@ apptainer exec --nv \ --bind /mnt/beegfs/bioinf/wijesekara/jaeger/experiments:/data/experiments \ --bind $INSTALL_DIR:/home/wijesekara/.local \ /mnt/beegfs/bioinf/wijesekara/jaeger/jaeger_dev.sif \ - bash -c "cd /jaeger && pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . && jaeger train -c /jaeger/train_config/nn_config_500bp_crossframe_zeus.yaml --force --mixed_precision" + bash -c "cd /jaeger && pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . && jaeger train -c /jaeger/train_config/nn_config_500bp_crossframe_zeus.yaml --force" echo "Finished at: $(date)" diff --git a/slurm/pytorch_ddp_500bp.slurm b/slurm/pytorch_ddp_500bp.slurm new file mode 100644 index 0000000..9fb49b9 --- /dev/null +++ b/slurm/pytorch_ddp_500bp.slurm @@ -0,0 +1,52 @@ +#!/bin/bash +#SBATCH --job-name=jaeger_500bp_ddp +#SBATCH --partition=gpu +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=2 +#SBATCH --gres=gpu:2 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=128G +#SBATCH --time=12:00:00 +#SBATCH --output=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/ddp_%j.log +#SBATCH --error=/mnt/beegfs/bioinf/wijesekara/jaeger/logs/ddp_%j.err + +echo "=== Jaeger 500bp PyTorch DDP training ===" +echo "Running on: $(hostname)" +echo "Nodes: $SLURM_NNODES" +echo "Tasks per node: $SLURM_NTASKS_PER_NODE" +echo "GPUs per node: $SLURM_GPUS_PER_NODE" +echo "Started at: $(date)" + +export CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES + +# Job-specific install directory to avoid conflicts +INSTALL_DIR=/mnt/beegfs/bioinf/wijesekara/jaeger/.local_${SLURM_JOB_ID}_ddp +rm -rf $INSTALL_DIR +mkdir -p $INSTALL_DIR/lib/python3.12/site-packages +mkdir -p $INSTALL_DIR/bin +export PYTHONPATH="$INSTALL_DIR/lib/python3.12/site-packages:$PYTHONPATH" +export PATH="$INSTALL_DIR/bin:$PATH" + +# Ensure experiments directory exists +mkdir -p /mnt/beegfs/bioinf/wijesekara/jaeger/experiments + +apptainer exec --nv \ + --bind /mnt/beegfs/bioinf/wijesekara/jaeger/Jaeger_current:/jaeger \ + --bind /mnt/beegfs/bioinf/wijesekara/jaeger/data:/data \ + --bind /mnt/beegfs/bioinf/wijesekara/jaeger/experiments:/data/experiments \ + --bind $INSTALL_DIR:/home/wijesekara/.local \ + /mnt/beegfs/bioinf/wijesekara/jaeger/jaeger_dev.sif \ + bash -c " + cd /jaeger + pip install --target=/home/wijesekara/.local/lib/python3.12/site-packages --no-deps -e . + + srun torchrun \ + --nnodes=$SLURM_NNODES \ + --nproc_per_node=$SLURM_GPUS_PER_NODE \ + --rdzv_id=jaeger_ddp_$SLURM_JOB_ID \ + --rdzv_backend=c10d \ + --rdzv_endpoint=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1):29500 \ + jaeger train -c /jaeger/train_config/nn_config_500bp_baseline_zeus.yaml --force + " + +echo "Finished at: $(date)" From 6a9f878c62801c367f2a4347d48331f6032e27e4 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 12:30:57 +0200 Subject: [PATCH 35/64] feat(pytorch): remove legacy TensorFlow training and inference code --- scripts/convert_preprocessed_data.py | 245 --- scripts/eval_clean.py | 201 -- scripts/eval_length_generalization_fast.py | 173 -- scripts/eval_simple.py | 138 -- .../train_and_eval_length_generalization.py | 198 -- src/jaeger/__init__.py | 15 - src/jaeger/cli.py | 570 +---- src/jaeger/commands/convert_graph.py | 84 - src/jaeger/commands/health.py | 145 +- src/jaeger/commands/predict_legacy.py | 342 --- src/jaeger/commands/quantize.py | 199 -- src/jaeger/commands/taxonomy.py | 673 ------ src/jaeger/commands/test.py | 146 -- src/jaeger/commands/utils.py | 10 +- src/jaeger/commands/utils_models.py | 166 -- src/jaeger/data/__init__.py | 17 - src/jaeger/data/loaders.py | 138 -- src/jaeger/data/tfrecord.py | 82 - src/jaeger/dataops/convert.py | 126 +- src/jaeger/nnlib/builder.py | 969 --------- src/jaeger/nnlib/conversion.py | 610 ------ src/jaeger/nnlib/inference.py | 1047 ---------- src/jaeger/nnlib/metrics.py | 106 - src/jaeger/nnlib/v1/__init__.py | 0 src/jaeger/nnlib/v1/layers.py | 784 ------- src/jaeger/nnlib/v2/__init__.py | 0 src/jaeger/nnlib/v2/layers.py | 1841 ----------------- src/jaeger/nnlib/v2/losses.py | 157 -- src/jaeger/nnlib/v2/maps.py | 539 ----- src/jaeger/preprocess/v1/convert.py | 222 -- src/jaeger/preprocess/v2/convert.py | 178 -- src/jaeger/seqops/encode.py | 521 ----- src/jaeger/utils/gpu.py | 36 - src/jaeger/utils/test.py | 22 +- tests/conftest.py | 6 - tests/integration/test_inference_pipeline.py | 114 - tests/integration/test_pytorch_parity.py | 174 -- tests/pytest/test_health.py | 246 --- tests/pytest/test_inference.py | 455 ---- tests/smoke/convert_test.py | 40 - tests/smoke/data_format_test.py | 301 --- tests/smoke/data_test.py | 275 --- tests/smoke/infer_test.py | 109 - tests/smoke/layers_test.py | 127 -- tests/smoke/loss_test.py | 60 - tests/smoke/train_test.py | 186 -- tests/test_frame_shuffle.py | 150 -- tests/unit/test_cli.py | 8 +- tests/unit/test_data_loaders.py | 60 - tests/unit/test_data_tfrecord.py | 46 - tests/unit/test_dataops_convert.py | 11 - tests/unit/test_dataops_ood.py | 2 - tests/unit/test_nnlib_conversion.py | 56 - tests/unit/test_nnlib_inference.py | 102 - tests/unit/test_nnlib_metrics.py | 52 - tests/unit/test_nnlib_v2_layers.py | 103 - tests/unit/test_nnlib_v2_losses.py | 61 - tests/unit/test_seqops_encode.py | 107 - tests/unit/test_utils_misc.py | 3 - 59 files changed, 60 insertions(+), 13494 deletions(-) delete mode 100644 scripts/convert_preprocessed_data.py delete mode 100644 scripts/eval_clean.py delete mode 100644 scripts/eval_length_generalization_fast.py delete mode 100644 scripts/eval_simple.py delete mode 100644 scripts/train_and_eval_length_generalization.py delete mode 100644 src/jaeger/commands/convert_graph.py delete mode 100644 src/jaeger/commands/predict_legacy.py delete mode 100644 src/jaeger/commands/quantize.py delete mode 100644 src/jaeger/commands/taxonomy.py delete mode 100644 src/jaeger/commands/test.py delete mode 100644 src/jaeger/commands/utils_models.py delete mode 100644 src/jaeger/data/loaders.py delete mode 100644 src/jaeger/data/tfrecord.py delete mode 100644 src/jaeger/nnlib/builder.py delete mode 100644 src/jaeger/nnlib/conversion.py delete mode 100644 src/jaeger/nnlib/inference.py delete mode 100644 src/jaeger/nnlib/metrics.py delete mode 100644 src/jaeger/nnlib/v1/__init__.py delete mode 100644 src/jaeger/nnlib/v1/layers.py delete mode 100644 src/jaeger/nnlib/v2/__init__.py delete mode 100644 src/jaeger/nnlib/v2/layers.py delete mode 100644 src/jaeger/nnlib/v2/losses.py delete mode 100644 src/jaeger/nnlib/v2/maps.py delete mode 100644 src/jaeger/preprocess/v1/convert.py delete mode 100644 src/jaeger/preprocess/v2/convert.py delete mode 100644 src/jaeger/seqops/encode.py delete mode 100644 src/jaeger/utils/gpu.py delete mode 100644 tests/integration/test_inference_pipeline.py delete mode 100644 tests/integration/test_pytorch_parity.py delete mode 100644 tests/pytest/test_health.py delete mode 100644 tests/pytest/test_inference.py delete mode 100644 tests/smoke/convert_test.py delete mode 100644 tests/smoke/data_format_test.py delete mode 100644 tests/smoke/data_test.py delete mode 100644 tests/smoke/infer_test.py delete mode 100644 tests/smoke/layers_test.py delete mode 100644 tests/smoke/loss_test.py delete mode 100644 tests/smoke/train_test.py delete mode 100644 tests/test_frame_shuffle.py delete mode 100644 tests/unit/test_data_loaders.py delete mode 100644 tests/unit/test_data_tfrecord.py delete mode 100644 tests/unit/test_nnlib_conversion.py delete mode 100644 tests/unit/test_nnlib_inference.py delete mode 100644 tests/unit/test_nnlib_metrics.py delete mode 100644 tests/unit/test_nnlib_v2_layers.py delete mode 100644 tests/unit/test_nnlib_v2_losses.py delete mode 100644 tests/unit/test_seqops_encode.py diff --git a/scripts/convert_preprocessed_data.py b/scripts/convert_preprocessed_data.py deleted file mode 100644 index d525b4e..0000000 --- a/scripts/convert_preprocessed_data.py +++ /dev/null @@ -1,245 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert Jaeger CSV training data to TFRecord or NumPy format. - -This pre-processes the data once, so training can skip live preprocessing -and load tensors directly — achieving 20-30x faster data loading. - -Usage: - python scripts/convert_preprocessed_data.py \\ - --csv train_shuffled.csv \\ - --output train_shuffled.npz \\ - --format numpy \\ - --crop-size 500 - - python scripts/convert_preprocessed_data.py \\ - --csv train_shuffled.csv \\ - --output train_shuffled.tfrecord \\ - --format tfrecord \\ - --crop-size 500 -""" - -import sys -from pathlib import Path - -# Add src to path for imports -sys.path.insert(0, str(Path(__file__).parent.parent / "src")) - -import argparse -import time -import numpy as np -import tensorflow as tf - -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - - -def _bytes_feature(value): - """Returns a bytes_list from a string / byte.""" - if isinstance(value, type(tf.constant(0))): - value = value.numpy() - return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) - - -def _float_feature(value): - """Returns a float_list from a float / double.""" - return tf.train.Feature(float_list=tf.train.FloatList(value=value)) - - -def _int64_feature(value): - """Returns an int64_list from a bool / enum / int / uint.""" - return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) - - -def serialize_tfrecord_embedding(translated, label): - """Serialize a single example for embedding layer (int32 indices).""" - translated_flat = tf.reshape(translated, [-1]) - feature = { - "translated": _int64_feature(translated_flat.numpy().tolist()), - "label": _float_feature(label.numpy().flatten().tolist()), - } - return tf.train.Example( - features=tf.train.Features(feature=feature) - ).SerializeToString() - - -def serialize_tfrecord_onehot(translated, label): - """Serialize a single example for one-hot encoded input (float32).""" - translated_flat = tf.reshape(translated, [-1]) - feature = { - "translated": _float_feature(translated_flat.numpy().flatten().tolist()), - "label": _float_feature(label.numpy().flatten().tolist()), - } - return tf.train.Example( - features=tf.train.Features(feature=feature) - ).SerializeToString() - - -def convert_csv_to_tfrecord( - csv_path, output_path, crop_size=500, use_embedding_layer=True -): - """Convert CSV to TFRecord with preprocessed tensors.""" - print(f"Converting {csv_path} -> {output_path}") - - total = sum(1 for _ in open(csv_path)) - print(f"Total samples: {total}") - - preprocess_fn = process_string_train( - crop_size=crop_size, - seq_onehot=False, - input_type="translated", - class_label_onehot=True, - num_classes=3, - shuffle=False, - ngram_width=3, - codons=CODONS, - codon_num=CODON_ID, - ) - - serialize_fn = ( - serialize_tfrecord_embedding - if use_embedding_layer - else serialize_tfrecord_onehot - ) - - Path(output_path).parent.mkdir(parents=True, exist_ok=True) - - with tf.io.TFRecordWriter(output_path) as writer: - start = time.time() - - with open(csv_path) as f: - for i, line in enumerate(f): - outputs, label = preprocess_fn(line.strip().encode()) - translated = outputs["translated"] - example = serialize_fn(translated, label) - writer.write(example) - - if (i + 1) % 5000 == 0: - elapsed = time.time() - start - rate = (i + 1) / elapsed - print( - f" {i + 1}/{total} ({100 * (i + 1) / total:.1f}%) - {rate:.1f} samples/sec" - ) - - elapsed = time.time() - start - print( - f"Done! {total} samples in {elapsed:.1f}s ({total / elapsed:.1f} samples/sec)" - ) - - -def convert_csv_to_numpy( - csv_path, output_path, crop_size=500, use_embedding_layer=True -): - """Convert CSV to NumPy .npz with preprocessed tensors.""" - print(f"Converting {csv_path} -> {output_path}") - - total = sum(1 for _ in open(csv_path)) - print(f"Total samples: {total}") - - preprocess_fn = process_string_train( - crop_size=crop_size, - seq_onehot=False, - input_type="translated", - class_label_onehot=True, - num_classes=3, - shuffle=False, - ngram_width=3, - codons=CODONS, - codon_num=CODON_ID, - ) - - # For fixed crop_size, sequences have fixed length: crop_size // 3 - 1 - seq_len = crop_size // 3 - 1 - - if use_embedding_layer: - sequences = np.zeros((total, 6, seq_len), dtype=np.int32) - else: - # Need to determine codon_depth from first sample - with open(csv_path) as f: - first_line = next(f) - outputs, _ = preprocess_fn(first_line.strip().encode()) - codon_depth = outputs["translated"].shape[-1] - sequences = np.zeros((total, 6, seq_len, codon_depth), dtype=np.float32) - print(f"Codon depth: {codon_depth}") - - labels = np.zeros((total, 3), dtype=np.float32) - - start = time.time() - with open(csv_path) as f: - for i, line in enumerate(f): - outputs, label = preprocess_fn(line.strip().encode()) - seq = outputs["translated"].numpy() - sequences[i] = seq - labels[i] = label.numpy() - - if (i + 1) % 5000 == 0: - elapsed = time.time() - start - rate = (i + 1) / elapsed - print( - f" {i + 1}/{total} ({100 * (i + 1) / total:.1f}%) - {rate:.1f} samples/sec" - ) - - elapsed = time.time() - start - print( - f"Done! {total} samples in {elapsed:.1f}s ({total / elapsed:.1f} samples/sec)" - ) - print( - f"Memory: sequences={sequences.nbytes / (1024**3):.2f}GB, labels={labels.nbytes / (1024**3):.3f}GB" - ) - - print(f"Saving to {output_path}...") - np.savez_compressed(output_path, translated=sequences, label=labels) - print("Saved!") - - -def main(): - parser = argparse.ArgumentParser( - description="Convert Jaeger CSV training data to TFRecord or NumPy format" - ) - parser.add_argument("--csv", required=True, help="Input CSV file path") - parser.add_argument("--output", required=True, help="Output file path") - parser.add_argument( - "--format", - required=True, - choices=["tfrecord", "numpy"], - help="Output format", - ) - parser.add_argument( - "--crop-size", - type=int, - default=500, - help="Sequence crop size (default: 500)", - ) - parser.add_argument( - "--use-embedding-layer", - action="store_true", - default=True, - help="Use embedding layer (int indices) vs one-hot (default: True)", - ) - parser.add_argument( - "--no-embedding-layer", - action="store_false", - dest="use_embedding_layer", - help="Use one-hot encoding instead of embedding layer", - ) - - args = parser.parse_args() - - if args.format == "tfrecord": - convert_csv_to_tfrecord( - args.csv, - args.output, - crop_size=args.crop_size, - use_embedding_layer=args.use_embedding_layer, - ) - elif args.format == "numpy": - convert_csv_to_numpy( - args.csv, - args.output, - crop_size=args.crop_size, - use_embedding_layer=args.use_embedding_layer, - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/eval_clean.py b/scripts/eval_clean.py deleted file mode 100644 index d75b2bb..0000000 --- a/scripts/eval_clean.py +++ /dev/null @@ -1,201 +0,0 @@ -#!/usr/bin/env python3 -"""Clean evaluation without checkpoint interference.""" - -import sys - -sys.path.insert(0, "src") -import os - -os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" -os.environ["GRPC_VERBOSITY"] = "ERROR" - -import tensorflow as tf -import numpy as np -import pandas as pd -from pathlib import Path -from sklearn.metrics import accuracy_score, classification_report - -from jaeger.utils.misc import load_model_config -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - -CLASS_NAMES = ["chromosome", "virus", "plasmid"] - - -def load_model(config_path, checkpoint_path): - """Load model without touching checkpoint dirs.""" - config = load_model_config(Path(config_path)) - - # Manually build model without DynamicModelBuilder's checkpoint logic - from jaeger.nnlib.v2.layers import ( - MaskedConv1D, - MaskedBatchNorm, - ResidualBlock_wrapper, - ) - - # Get model config - model_cfg = config["model"] - embedding_cfg = model_cfg.get("embedding", {}) - rep_cfg = model_cfg.get("representation_learner", {}) - cls_cfg = model_cfg.get("classifier", {}) - - # Build inputs - input_type = embedding_cfg.get("input_type", "translated") - use_embedding = embedding_cfg.get("use_embedding_layer", True) - - if input_type == "translated" and use_embedding: - inputs = tf.keras.Input(shape=(6, None), dtype=tf.int32, name="translated") - x = tf.keras.layers.Embedding( - 65, # codon_depth + 1 (since seq+1) - embedding_cfg.get("embedding_size", 64), - name="embedding", - )(inputs) - else: - raise NotImplementedError("Only translated + embedding supported") - - # Build representation learner - mask = tf.keras.layers.Lambda(lambda x: tf.not_equal(x, 0), name="mask")(inputs) - - for layer_cfg in rep_cfg.get("hidden_layers", []): - name = layer_cfg.get("name", "").lower() - cfg = dict(layer_cfg.get("config", {})) - - if name == "masked_conv1d": - cfg.pop("kernel_regularizer_w", None) - cfg.pop("kernel_regularizer", None) - x = MaskedConv1D(**cfg)(x, mask=mask) - elif name == "masked_batchnorm": - x = MaskedBatchNorm(**cfg)(x) - elif name == "activation": - x = tf.keras.layers.Activation(**cfg)(x) - elif name == "residual_block": - block_size = cfg.pop("block_size", 1) - filters = cfg.get("filters", 32) - cfg.pop("kernel_regularizer_w", None) - cfg.pop("kernel_regularizer", None) - shape = (6, None, filters) - x = ResidualBlock_wrapper(block_size, shape, **cfg)(x) - elif name == "cross_frame_attention": - from jaeger.nnlib.v2.layers import CrossFrameAttention - - x = CrossFrameAttention(**cfg)(x) - elif name == "axial_attention": - from jaeger.nnlib.v2.layers import AxialAttention - - x = AxialAttention(**cfg)(x) - - # Pooling - pooling = rep_cfg.get("pooling", "average").lower() - if pooling == "average": - x = tf.keras.layers.GlobalAveragePooling2D()(x) - elif pooling == "max": - x = tf.keras.layers.GlobalMaxPooling2D()(x) - - # Classifier - for layer_cfg in cls_cfg.get("hidden_layers", []): - name = layer_cfg.get("name", "").lower() - cfg = dict(layer_cfg.get("config", {})) - - if name == "dropout": - x = tf.keras.layers.Dropout(**cfg)(x) - elif name == "dense": - cfg.pop("kernel_regularizer_w", None) - cfg.pop("kernel_regularizer", None) - cfg.pop("bias_initializer", None) - x = tf.keras.layers.Dense(**cfg)(x) - - outputs = {"prediction": x, "embedding": x} - model = tf.keras.Model(inputs=inputs, outputs=outputs, name="jaeger_model") - - # Load weights - if checkpoint_path and Path(checkpoint_path).exists(): - print(f" Loading weights from: {checkpoint_path}") - model.load_weights(checkpoint_path) - - return model - - -def create_dataset(csv_path, crop_size, max_samples=5000, batch_size=256): - df = pd.read_csv(csv_path, header=None, names=["label", "seq"]) - if max_samples and len(df) > max_samples: - df = df.sample(n=max_samples, random_state=42) - - csv_strings = [f"{label},{seq}" for label, seq in zip(df["label"], df["seq"])] - labels = df["label"].values.astype(np.int32) - - processor = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - seq_onehot=False, - crop_size=crop_size, - timesteps=False, - mutate=False, - mutation_rate=0.0, - shuffle=False, - masking=False, - class_label_onehot=False, - ngram_width=3, - ) - - def _parse(csv_str): - out, _ = processor(csv_str) - return out["translated"] - - dataset = tf.data.Dataset.from_tensor_slices(csv_strings) - dataset = dataset.map(_parse, num_parallel_calls=tf.data.AUTOTUNE) - dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) - - return dataset, labels - - -def evaluate(model, dataset, true_labels, length, model_name): - print(f"\nEvaluating {model_name} on {length}bp") - preds = [] - for batch in dataset: - out = model(batch, training=False) - logits = out["prediction"].numpy() - preds.extend(np.argmax(logits, axis=-1)) - preds = np.array(preds) - acc = accuracy_score(true_labels, preds) - print(f" Accuracy: {acc:.4f}") - print( - f" Report:\n{classification_report(true_labels, preds, target_names=CLASS_NAMES, digits=4, zero_division=0)}" - ) - return {"model": model_name, "length": length, "accuracy": acc} - - -def main(): - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--config", required=True) - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--test-csv", required=True) - parser.add_argument("--model-name", required=True) - parser.add_argument("--lengths", type=int, nargs="+", default=[500, 1000, 2000]) - parser.add_argument("--max-samples", type=int, default=5000) - parser.add_argument("--batch-size", type=int, default=256) - args = parser.parse_args() - - print(f"Loading model: {args.model_name}") - model = load_model(args.config, args.checkpoint) - - results = [] - for length in args.lengths: - dataset, labels = create_dataset( - args.test_csv, length, args.max_samples, args.batch_size - ) - result = evaluate(model, dataset, labels, length, args.model_name) - results.append(result) - - print(f"\n{'=' * 50}") - print("SUMMARY") - print(f"{'=' * 50}") - print(f"{'Model':<15} {'500bp':<10} {'1000bp':<10} {'2000bp':<10}") - print("-" * 50) - accs = [r["accuracy"] for r in results] - print(f"{args.model_name:<15} {accs[0]:<10.4f} {accs[1]:<10.4f} {accs[2]:<10.4f}") - - -if __name__ == "__main__": - main() diff --git a/scripts/eval_length_generalization_fast.py b/scripts/eval_length_generalization_fast.py deleted file mode 100644 index 190ec85..0000000 --- a/scripts/eval_length_generalization_fast.py +++ /dev/null @@ -1,173 +0,0 @@ -#!/usr/bin/env python3 -""" -Fast evaluation of length generalization using tf.data pipeline. -""" - -import sys - -sys.path.insert(0, "src") - -import os - -os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" -os.environ["GRPC_VERBOSITY"] = "ERROR" - -import tensorflow as tf -import numpy as np -import pandas as pd -from pathlib import Path -from sklearn.metrics import accuracy_score, classification_report, confusion_matrix -import argparse - -from jaeger.utils.misc import load_model_config -from jaeger.commands.train import DynamicModelBuilder -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - -CLASS_NAMES = ["chromosome", "virus", "plasmid"] - - -def load_model(config_path, checkpoint_path=None): - import shutil - - config = load_model_config(Path(config_path)) - config["from_last_checkpoint"] = False - config["force"] = False - - # Temporarily move checkpoints to avoid the warning - checkpoint_dir = None - if "training" in config: - checkpoint_dir = config["training"].get("classifier_dir", "") - checkpoint_dir = checkpoint_dir.replace( - "{{ model.base_dir }}", config["model"]["base_dir"] - ) - checkpoint_dir = checkpoint_dir.replace( - "{{ training.experiment_root }}", - f"experiments/experiment_{config['model']['experiment']}_{config['model']['seed']}", - ) - - if checkpoint_dir and Path(checkpoint_dir).exists(): - backup_dir = str(checkpoint_dir) + "_backup" - if Path(backup_dir).exists(): - shutil.rmtree(backup_dir) - shutil.move(checkpoint_dir, backup_dir) - else: - backup_dir = None - - builder = DynamicModelBuilder(config) - models = builder.build_fragment_classifier() - builder.compile_model(models, train_branch="classifier") - - # Restore checkpoints and load weights - if backup_dir: - shutil.move(backup_dir, checkpoint_dir) - - if checkpoint_path and Path(checkpoint_path).exists(): - print(f" Loading weights from: {checkpoint_path}") - models["jaeger_model"].load_weights(checkpoint_path) - return models["jaeger_model"] - - -def create_dataset(csv_path, crop_size, max_samples=5000, batch_size=256): - """Create tf.data dataset with proper preprocessing.""" - df = pd.read_csv(csv_path, header=None, names=["label", "seq"]) - if max_samples and len(df) > max_samples: - df = df.sample(n=max_samples, random_state=42) - - # Create CSV strings - csv_strings = [f"{label},{seq}" for label, seq in zip(df["label"], df["seq"])] - labels = df["label"].values.astype(np.int32) - - # Create processor - processor = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - seq_onehot=False, - crop_size=crop_size, - timesteps=False, - mutate=False, - mutation_rate=0.0, - shuffle=False, - masking=False, - class_label_onehot=False, - ngram_width=3, - ) - - def _parse(csv_str): - out, _ = processor(csv_str) - return out["translated"] - - dataset = tf.data.Dataset.from_tensor_slices(csv_strings) - dataset = dataset.map(_parse, num_parallel_calls=tf.data.AUTOTUNE) - dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) - - return dataset, labels - - -def evaluate_model(model, dataset, true_labels, length, model_name): - print(f"\n{'=' * 60}") - print(f"Evaluating {model_name} on {length}bp sequences") - print(f"{'=' * 60}") - - preds = [] - for batch in dataset: - out = model(batch, training=False) - logits = out["prediction"].numpy() if isinstance(out, dict) else out.numpy() - preds.extend(np.argmax(logits, axis=-1)) - - preds = np.array(preds) - acc = accuracy_score(true_labels, preds) - print(f"Accuracy: {acc:.4f}") - print("\nClassification Report:") - print( - classification_report( - true_labels, preds, target_names=CLASS_NAMES, digits=4, zero_division=0 - ) - ) - print("Confusion Matrix:") - print(confusion_matrix(true_labels, preds)) - - return {"model": model_name, "length": length, "accuracy": acc} - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--config", required=True) - parser.add_argument("--checkpoint", default=None) - parser.add_argument("--test-csv", required=True) - parser.add_argument("--model-name", required=True) - parser.add_argument("--lengths", type=int, nargs="+", default=[500, 1000, 2000]) - parser.add_argument("--max-samples", type=int, default=5000) - parser.add_argument("--batch-size", type=int, default=256) - args = parser.parse_args() - - print(f"Loading model: {args.model_name}") - model = load_model(args.config, args.checkpoint) - - results = [] - for length in args.lengths: - dataset, labels = create_dataset( - args.test_csv, length, args.max_samples, args.batch_size - ) - result = evaluate_model(model, dataset, labels, length, args.model_name) - results.append(result) - - print(f"\n{'=' * 60}") - print("SUMMARY") - print(f"{'=' * 60}") - print(f"{'Model':<25} {'Length':<10} {'Accuracy':<10}") - print("-" * 60) - for r in results: - print(f"{r['model']:<25} {r['length']:<10} {r['accuracy']:<10.4f}") - - output_dir = Path("evaluation_results") - output_dir.mkdir(exist_ok=True) - np.savez( - output_dir / f"{args.model_name}_results.npz", - lengths=np.array([r["length"] for r in results]), - accuracies=np.array([r["accuracy"] for r in results]), - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/eval_simple.py b/scripts/eval_simple.py deleted file mode 100644 index 6830ad8..0000000 --- a/scripts/eval_simple.py +++ /dev/null @@ -1,138 +0,0 @@ -#!/usr/bin/env python3 -"""Simple evaluation using DynamicModelBuilder with checkpoint protection.""" - -import sys - -sys.path.insert(0, "src") -import os - -os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" -os.environ["GRPC_VERBOSITY"] = "ERROR" - -import tensorflow as tf -import numpy as np -import pandas as pd -from pathlib import Path -from sklearn.metrics import accuracy_score, classification_report -import shutil - -from jaeger.utils.misc import load_model_config -from jaeger.commands.train import DynamicModelBuilder -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - -CLASS_NAMES = ["chromosome", "virus", "plasmid"] - - -def load_model(config_path, checkpoint_path): - config = load_model_config(Path(config_path)) - config["from_last_checkpoint"] = False - config["force"] = False - - # Temporarily move checkpoint dir to avoid the warning - checkpoint_dir = str(Path(checkpoint_path).parent) - backup_dir = checkpoint_dir + "_backup" - - if Path(checkpoint_dir).exists(): - if Path(backup_dir).exists(): - shutil.rmtree(backup_dir) - shutil.move(checkpoint_dir, backup_dir) - - try: - builder = DynamicModelBuilder(config) - models = builder.build_fragment_classifier() - builder.compile_model(models, train_branch="classifier") - finally: - # Always restore checkpoint dir - if Path(backup_dir).exists(): - shutil.move(backup_dir, checkpoint_dir) - - print(f" Loading weights from: {checkpoint_path}") - models["jaeger_model"].load_weights(checkpoint_path) - return models["jaeger_model"] - - -def create_dataset(csv_path, crop_size, max_samples=5000, batch_size=256): - df = pd.read_csv(csv_path, header=None, names=["label", "seq"]) - if max_samples and len(df) > max_samples: - df = df.sample(n=max_samples, random_state=42) - - csv_strings = [f"{label},{seq}" for label, seq in zip(df["label"], df["seq"])] - labels = df["label"].values.astype(np.int32) - - processor = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - seq_onehot=False, - crop_size=crop_size, - timesteps=False, - mutate=False, - mutation_rate=0.0, - shuffle=False, - masking=False, - class_label_onehot=False, - ngram_width=3, - ) - - def _parse(csv_str): - out, _ = processor(csv_str) - return out["translated"] - - dataset = tf.data.Dataset.from_tensor_slices(csv_strings) - dataset = dataset.map(_parse, num_parallel_calls=tf.data.AUTOTUNE) - dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) - - return dataset, labels - - -def evaluate(model, dataset, true_labels, length, model_name): - print(f"\nEvaluating {model_name} on {length}bp") - preds = [] - for batch in dataset: - out = model(batch, training=False) - logits = out["prediction"].numpy() if isinstance(out, dict) else out.numpy() - preds.extend(np.argmax(logits, axis=-1)) - preds = np.array(preds) - acc = accuracy_score(true_labels, preds) - print(f" Accuracy: {acc:.4f}") - print( - f" Report:\n{classification_report(true_labels, preds, target_names=CLASS_NAMES, digits=4, zero_division=0)}" - ) - return {"model": model_name, "length": length, "accuracy": acc} - - -def main(): - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--config", required=True) - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--test-csv", required=True) - parser.add_argument("--model-name", required=True) - parser.add_argument("--lengths", type=int, nargs="+", default=[500, 1000, 2000]) - parser.add_argument("--max-samples", type=int, default=5000) - parser.add_argument("--batch-size", type=int, default=256) - args = parser.parse_args() - - print(f"Loading model: {args.model_name}") - model = load_model(args.config, args.checkpoint) - - results = [] - for length in args.lengths: - dataset, labels = create_dataset( - args.test_csv, length, args.max_samples, args.batch_size - ) - result = evaluate(model, dataset, labels, length, args.model_name) - results.append(result) - - print(f"\n{'=' * 50}") - print("SUMMARY") - print(f"{'=' * 50}") - print(f"{'Model':<15} {'500bp':<10} {'1000bp':<10} {'2000bp':<10}") - print("-" * 50) - accs = [r["accuracy"] for r in results] - print(f"{args.model_name:<15} {accs[0]:<10.4f} {accs[1]:<10.4f} {accs[2]:<10.4f}") - - -if __name__ == "__main__": - main() diff --git a/scripts/train_and_eval_length_generalization.py b/scripts/train_and_eval_length_generalization.py deleted file mode 100644 index eef6e3c..0000000 --- a/scripts/train_and_eval_length_generalization.py +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env python3 -""" -Train models on 500bp and evaluate length generalization. -Trains baseline, cross-frame, and axial models, then evaluates on 500/1000/2000bp. -""" - -import sys - -sys.path.insert(0, "src") - -import os - -os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" -os.environ["GRPC_VERBOSITY"] = "ERROR" - -import tensorflow as tf -import numpy as np -import pandas as pd -from pathlib import Path -from sklearn.metrics import accuracy_score - -from jaeger.utils.misc import load_model_config -from jaeger.commands.train import DynamicModelBuilder -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - -CLASS_NAMES = ["chromosome", "virus", "plasmid"] - - -def train_model(config_path): - """Train model and return the trained model + config.""" - config = load_model_config(Path(config_path)) - config["from_last_checkpoint"] = False - config["force"] = True - - builder = DynamicModelBuilder(config) - models = builder.build_fragment_classifier() - builder.compile_model(models, train_branch="classifier") - - # Get data - string_processor_config = builder._get_string_processor_config() - _train_data = builder._get_fragment_paths() - - train_data = {"train": None, "validation": None} - for k, v in _train_data.items(): - from jaeger.commands.train import check_files - - paths = check_files(v.get("paths")) - _data = tf.data.TextLineDataset( - paths, num_parallel_reads=len(paths), buffer_size=200 - ) - _buffer_size = string_processor_config.get("buffer_size") - - padded_shape = {"translated": [6, None]} - train_data[k] = ( - _data.map( - process_string_train( - codons=string_processor_config.get("codon"), - codon_num=string_processor_config.get("codon_id"), - codon_depth=string_processor_config.get("codon_depth"), - label_original=string_processor_config.get( - "classifier_labels", None - ), - label_alternative=string_processor_config.get( - "classifier_labels_map", None - ), - ngram_width=string_processor_config.get("ngram_width"), - seq_onehot=string_processor_config.get("seq_onehot"), - crop_size=string_processor_config.get("crop_size"), - input_type=string_processor_config.get("input_type"), - masking=string_processor_config.get("masking"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle=string_processor_config.get("shuffle"), - shuffle_frames=string_processor_config.get("shuffle_frames"), - ), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .shuffle(_buffer_size, reshuffle_each_iteration=True) - .padded_batch( - builder.train_cfg.get("batch_size", 32), - padded_shapes=( - padded_shape, - {"classifier": builder.classifier_out_dim}, - ), - ) - .prefetch(tf.data.AUTOTUNE) - ) - - # Train - train_args = { - "validation_data": train_data.get("validation").take( - builder.train_cfg.get("classifier_validation_steps") - ), - "epochs": builder.train_cfg.get("classifier_epochs"), - "callbacks": builder.get_callbacks(branch="classifier"), - } - - models.get("jaeger_classifier").fit( - train_data.get("train").take(builder.train_cfg.get("classifier_train_steps")), - class_weight=builder.train_cfg.get("classifier_class_weights"), - **train_args, - ) - - return models["jaeger_model"], builder - - -def create_eval_dataset(csv_path, crop_size, max_samples=5000, batch_size=256): - """Create eval dataset.""" - df = pd.read_csv(csv_path, header=None, names=["label", "seq"]) - if max_samples and len(df) > max_samples: - df = df.sample(n=max_samples, random_state=42) - - csv_strings = [f"{label},{seq}" for label, seq in zip(df["label"], df["seq"])] - labels = df["label"].values.astype(np.int32) - - processor = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - seq_onehot=False, - crop_size=crop_size, - timesteps=False, - mutate=False, - mutation_rate=0.0, - shuffle=False, - masking=False, - class_label_onehot=False, - ngram_width=3, - ) - - def _parse(csv_str): - out, _ = processor(csv_str) - return out["translated"] - - dataset = tf.data.Dataset.from_tensor_slices(csv_strings) - dataset = dataset.map(_parse, num_parallel_calls=tf.data.AUTOTUNE) - dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) - - return dataset, labels - - -def evaluate(model, dataset, true_labels, length, model_name): - print(f"\nEvaluating {model_name} on {length}bp") - preds = [] - for batch in dataset: - out = model(batch, training=False) - logits = out["prediction"].numpy() if isinstance(out, dict) else out.numpy() - preds.extend(np.argmax(logits, axis=-1)) - preds = np.array(preds) - acc = accuracy_score(true_labels, preds) - print(f" Accuracy: {acc:.4f}") - return {"model": model_name, "length": length, "accuracy": acc} - - -def main(): - configs = [ - ("train_config/nn_config_500bp_baseline.yaml", "baseline"), - ("train_config/nn_config_500bp_crossframe.yaml", "crossframe"), - ("train_config/nn_config_500bp_axial.yaml", "axial"), - ] - - test_csv = "/home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/training/genomad_train_data/jaeger_format/val.csv" - lengths = [500, 1000, 2000] - max_samples = 5000 - - all_results = [] - - for config_path, model_name in configs: - print(f"\n{'#' * 60}") - print(f"# Training {model_name}") - print(f"{'#' * 60}") - - model, builder = train_model(config_path) - - for length in lengths: - dataset, labels = create_eval_dataset(test_csv, length, max_samples) - result = evaluate(model, dataset, labels, length, model_name) - all_results.append(result) - - # Summary - print(f"\n{'=' * 60}") - print("FINAL SUMMARY") - print(f"{'=' * 60}") - print(f"{'Model':<15} {'500bp':<10} {'1000bp':<10} {'2000bp':<10}") - print("-" * 60) - - for model_name in ["baseline", "crossframe", "axial"]: - accs = [r["accuracy"] for r in all_results if r["model"] == model_name] - print(f"{model_name:<15} {accs[0]:<10.4f} {accs[1]:<10.4f} {accs[2]:<10.4f}") - - # Save - output_dir = Path("evaluation_results") - output_dir.mkdir(exist_ok=True) - np.savez(output_dir / "length_generalization_results.npz", results=all_results) - - -if __name__ == "__main__": - main() diff --git a/src/jaeger/__init__.py b/src/jaeger/__init__.py index f23c366..e69de29 100644 --- a/src/jaeger/__init__.py +++ b/src/jaeger/__init__.py @@ -1,15 +0,0 @@ -import os - -# Suppress TensorFlow and C++ logging noise before any TF-related imports. -# These environment variables are read when TensorFlow's native libraries load, -# so they must be set at package import time. -defaults = { - "TF_CPP_MIN_LOG_LEVEL": "3", # 0=DEBUG, 1=INFO, 2=WARN, 3=ERROR - "ABSL_MIN_LOG_LEVEL": "2", # Suppress absl INFO/WARNING before init - "TF_ENABLE_ONEDNN_OPTS": "0", # Disable oneDNN custom ops warning - "GRPC_VERBOSITY": "ERROR", # Suppress gRPC logs - "GLOG_minloglevel": "2", # Suppress glog INFO messages -} -for key, value in defaults.items(): - if not os.environ.get(key): - os.environ[key] = value diff --git a/src/jaeger/cli.py b/src/jaeger/cli.py index 3bce569..3ddd8a6 100644 --- a/src/jaeger/cli.py +++ b/src/jaeger/cli.py @@ -10,67 +10,6 @@ both phages and prophages in metagenomic assemblies. """ -import os -import sys - -# Suppress TensorFlow and C++ warnings before any TF imports. -# These must be set before the first tensorflow import. -os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" # 0=DEBUG, 1=INFO, 2=WARN, 3=ERROR -os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" # Disable oneDNN custom ops warning -os.environ["GRPC_VERBOSITY"] = "ERROR" # Suppress gRPC logs -os.environ["GLOG_minloglevel"] = "2" # Suppress glog INFO messages -os.environ["WRAPT_DISABLE_EXTENSIONS"] = "true" - -if sys.platform.startswith("linux"): - os.environ["XLA_FLAGS"] = "--xla_gpu_cuda_data_dir=/usr/lib/cuda" - - -# Patch TensorFlow 2.18 + Python 3.12 compatibility bug before any TF imports. -# TF's internal _DictWrapper objects fail inspect.getattr_static during -# isinstance checks with typing generics, causing SavedModel loading to crash. -# This bug is fixed in TensorFlow 2.21+. -# See: https://github.com/tensorflow/tensorflow/issues/78784 -# We also temporarily silence native stderr during this first TF import to -# suppress harmless but noisy startup messages (oneDNN, CUDA plugin checks, -# absl pre-init warnings, etc.). -def _silence_tf_import(): - import os as _os - - stderr_fd = _os.dup(2) - devnull = _os.open(_os.devnull, _os.O_WRONLY) - try: - _os.dup2(devnull, 2) - import tensorflow as tf # noqa: F401 - from tensorflow.python.framework import tensor_util - - _original_is_tf_type = tensor_util.is_tf_type - - def _patched_is_tf_type(x): - try: - return _original_is_tf_type(x) - except TypeError: - return False - - # Detect the bug by trying the actual _DictWrapper class - from tensorflow.python.trackable.data_structures import _DictWrapper - - try: - _original_is_tf_type(_DictWrapper()) - except TypeError: - tensor_util.is_tf_type = _patched_is_tf_type - finally: - _os.dup2(stderr_fd, 2) - _os.close(devnull) - _os.close(stderr_fd) - - -try: - _silence_tf_import() -except Exception: - pass # TF not installed, incompatible version, or bug already fixed -finally: - del _silence_tf_import - import click import logging from pathlib import Path @@ -78,14 +17,9 @@ def _patched_is_tf_type(x): from importlib.resources import files from jaeger.utils.misc import json_to_dict, add_data_to_json, AvailableModels import warnings -from jaeger.commands.quantize import quantize_model -from jaeger.commands.convert_graph import convert_graph warnings.filterwarnings("ignore") -if sys.platform.startswith("linux"): - os.environ["XLA_FLAGS"] = "--xla_gpu_cuda_data_dir=/usr/lib/cuda" - USER_MODEL_PATHS = json_to_dict(files("jaeger.data") / "config.json").get("model_paths") DEFAULT_MODEL_PATH = [files("jaeger.data")] @@ -93,10 +27,7 @@ def _patched_is_tf_type(x): i for i in AvailableModels(path=DEFAULT_MODEL_PATH + USER_MODEL_PATHS).info.keys() if i != "jaeger_fragment" -] + ["default"] - -# Models that still use the legacy prediction workflow. -LEGACY_PREDICT_MODELS = {"default", "experimental_1", "experimental_2"} +] logger = logging.getLogger("Jaeger") @@ -144,22 +75,25 @@ def health(**kwargs): @click.option( "-m", "--model", - default="default", - help=( - f"Select a deep-learning model to use. " - f"Available choices (when --config is not set): {', '.join(AVAIL_MODELS)}" - ), + default=None, + help=("Model name (deprecated; use --model_path or --config)."), ) @click.option( "--model_path", default=None, - help=("Give the path to a model. overrides --model"), + help=("Path to a PyTorch model directory containing config.yaml and model.pt."), ) @click.option( "--config", type=click.Path(exists=True), help="Path to Jaeger config file (e.g., when using Apptainer or Docker)", ) +@click.option( + "--checkpoint", + type=click.Path(exists=True), + default=None, + help="Path to a PyTorch checkpoint (model.pt) to use with --config.", +) @click.option( "-p", "--prophage", @@ -244,72 +178,19 @@ def health(**kwargs): default=1, ) @click.option("-f", "--overwrite", is_flag=True, help="Overwrite existing files") -@click.option( - "--quantized", - type=click.Choice(["dynamic", "float16", "full_int8"]), - help="Use quantized TFLite model for inference (dynamic|float16|full_int8)", -) -@click.option( - "--precision", - type=click.Choice(["fp32", "fp16", "bf16"]), - default="fp32", - help="GPU inference precision: fp32 (default), fp16, or bf16. fp16/bf16 reduce memory and may speed up inference on compatible GPUs.", -) -@click.option( - "--xla", - is_flag=True, - help="Enable XLA JIT compilation for inference. May provide 2-3x speedup on GPU after initial compilation overhead.", -) -@click.option( - "--onnx", - is_flag=True, - help="Use ONNX Runtime for inference. Requires converting the model first with 'jaeger utils convert-graph --mode onnx'.", -) -@click.option( - "--int8", - is_flag=True, - help="Use INT8 quantized ONNX model (use with --onnx). Requires 'jaeger utils convert-graph --mode onnx --int8'.", -) def predict(**kwargs): - """ - Runs Jaeger on a dataset - """ - model = kwargs.get("model") + """Runs Jaeger on a dataset using the PyTorch inference path.""" model_path = kwargs.get("model_path") config = kwargs.get("config") - if model_path: - model = "from_path" - else: - if model is None: - model = "default" - elif config is None and model not in AVAIL_MODELS: - raise click.BadParameter( - f"Model '{model}' is not one of the available options: {', '.join(AVAIL_MODELS)}" - ) - - """Run jaeger inference pipeline.""" - if model in LEGACY_PREDICT_MODELS: - import click as _click - - _click.echo( - _click.style( - f"Warning: model '{model}' uses the legacy prediction workflow and is deprecated. " - "It will be removed in a future release. Please migrate to a modern model " - "(e.g., 'jaeger_38341_1.4M_fragment').", - fg="yellow", - ), - err=True, + if not model_path and not config: + raise click.BadParameter( + "PyTorch inference requires either --model_path or --config." ) - if model == "default": - from jaeger.commands.predict_legacy import run_core - - run_core(**kwargs) - else: - from jaeger.commands.predict import run_core + from jaeger.commands.predict import run_core - run_core(**kwargs) + run_core(**kwargs) @click.command(context_settings={"show_default": True}) @@ -954,8 +835,7 @@ def convert(**kwargs): jaeger utils optimize-data -i train.csv -o train.npz --format numpy_full Supported formats: - tfrecord - TFRecord with preprocessed tensors - numpy_raw - int8 sequences + TF preprocessing at train time + numpy_raw - int8 sequences + runtime preprocessing at train time numpy_full - fully preprocessed, fastest loading numpy_raw_variable - variable-length int8 sequences """, @@ -977,7 +857,7 @@ def convert(**kwargs): @click.option( "--format", type=click.Choice( - ["tfrecord", "numpy_raw", "numpy_full", "numpy_raw_variable"], + ["numpy_raw", "numpy_full", "numpy_raw_variable"], case_sensitive=False, ), required=True, @@ -1007,7 +887,7 @@ def convert(**kwargs): "--use-embedding-layer/--no-embedding-layer", default=True, show_default=True, - help="Use embedding layer (int indices) vs one-hot (tfrecord only)", + help="Use embedding layer (int indices) vs one-hot (numpy formats)", ) @click.option( "--max-length", @@ -1033,45 +913,6 @@ def optimize_data(**kwargs): pass -@utils.command( - context_settings=dict(ignore_unknown_options=True, show_default=True), - help="""Combine multiple model graphs into an ensemble model. Outputs can be summarized with majority voting, - sum or mean of logits of each model - - usage\n - ----- - - jaeger utils combine_models -i path/to/model1 -i path/to/model2 -i path/to/model3 -c mv - """, -) -@click.option( - "-i", - "--input", - type=click.Path(exists=True), - required=True, - help="Path to saved model", - multiple=True, -) -@click.option( - "-o", "--output", type=str, required=True, help="Path to save the ensemble model" -) -@click.option( - "-c", - "--comb", - type=click.Choice(["MV", "SUM", "MEAN", "NONE"], case_sensitive=False), - required=True, - help="how to summarize the model outputs " - "MV: majority voting" - "SUM: sum of logits" - "MEAN: Mean of logits" - "NONE: Do not aggregate", -) -def combine_models(**kwargs): - from jaeger.commands.utils_models import combine_models_core - - combine_models_core(**kwargs) - - @utils.command( context_settings=dict(ignore_unknown_options=True, show_default=True), help=""" @@ -1100,385 +941,12 @@ def stats(**kwargs): pass -@utils.command( - context_settings=dict(ignore_unknown_options=True, show_default=True), - help=""" - Quantize a Jaeger model for faster inference - - usage - ----- - - jaeger utils quantize -m default -o ./quantized_models - jaeger utils quantize -m jaeger_57341_1.5M_fragment -o ./quantized --mode float16 - """, -) -@click.option( - "-m", - "--model", - type=str, - required=True, - help="Model name to quantize", -) -@click.option( - "-o", - "--output", - type=click.Path(path_type=Path), - required=True, - help="Output directory for the quantized model", -) -@click.option( - "--mode", - type=click.Choice(["dynamic", "float16", "full_int8"]), - default="dynamic", - help="Quantization mode", -) -@click.option( - "-v", - "--verbose", - count=True, - help="Verbosity level", - default=1, -) -def quantize_cmd(**kwargs): - """Quantize a Jaeger model""" - quantize_model(**kwargs) - - -@utils.command( - context_settings=dict(ignore_unknown_options=True, show_default=True), - help=""" - Convert a Jaeger SavedModel to an optimized inference graph. - - Supports XLA, TFLite, ONNX, and TensorRT backends. - - usage - ----- - - jaeger utils convert-graph -m default -o ./optimized --mode xla - jaeger utils convert-graph -m default -o ./optimized --mode onnx - jaeger utils convert-graph -m default -o ./optimized --mode onnx --int8 - """, -) -@click.option( - "-m", - "--model", - type=str, - required=True, - help="Model name to convert", -) -@click.option( - "-o", - "--output", - type=click.Path(path_type=Path), - required=True, - help="Output directory for the converted model", -) -@click.option( - "--mode", - type=click.Choice(["xla", "tflite", "onnx", "tensorrt"]), - default="xla", - help="Conversion mode", -) -@click.option( - "--int8", - is_flag=True, - help="Apply static INT8 quantization (ONNX mode only)", -) -@click.option( - "-v", - "--verbose", - count=True, - help="Verbosity level", - default=1, -) -def convert_graph_cmd(model, output, mode, int8, verbose): - """Convert a Jaeger SavedModel to an optimized inference graph.""" - convert_graph(model, output, mode, verbose, int8=int8) - - -@click.group(context_settings=dict(help_option_names=["-h", "--help"])) -@click.pass_context -def taxonomy(obj): - """ - exprimental taxonomy prediction pipeline - """ - pass - - -@taxonomy.command( - context_settings=dict(ignore_unknown_options=True), - help=""" - Build a taxonomy database - - usage - ----- - - jaeger taxonomy build [OPTIONS] -i contigs.fasta -o taxonomy_db - - """, -) -@click.option( - "-i", - "--input", - type=click.Path(exists=True), - required=True, - help="Path to input file", -) -@click.option( - "-t", - "--tax", - type=click.Path(exists=True), - required=True, - help="Path to taxdump", -) -@click.option( - "-a", - "--acc2tax", - type=click.Path(exists=True), - required=True, - help="Path to accession2taxid file", -) -@click.option( - "-o", "--output", type=str, required=True, help="Path to output directory" -) -@click.option( - "--fsize", - type=int, - default=2048, - help="Length of the sliding window", -) -@click.option( - "--stride", - type=int, - default=2048, - help="The length of the gap between two sliding windows (stride==fsize)", -) -@click.option( - "-m", - "--model", - default="default", - help=( - f"Select a deep-learning model to use. " - f"Available choices (when --config is not set): {', '.join(AVAIL_MODELS)}" - ), -) -@click.option( - "--model_path", - default=None, - help=("Give the path to a model. overrides --model"), -) -@click.option( - "--config", - type=click.Path(exists=True), - help="Path to Jaeger config file (useful when using Apptainer or Docker)", -) -@click.option( - "--rc", - type=float, - default=0.1, - help="Minimum reliability score required to accept predictions", -) -@click.option( - "--batch", - type=int, - default=96, - help="Parallel batch size, lower if GPU runs out of memory", -) -@click.option("--workers", type=int, default=4, help="Number of threads to use") -@click.option( - "--cpu", - is_flag=True, - help="Ignore available GPUs and explicitly run on CPU", -) -@click.option( - "--physicalid", - type=int, - default=0, - help="Set default GPU device ID for multi-GPU systems", -) -@click.option( - "--mem", - type=int, - default=4, - help="GPU memory limit", -) -@click.option( - "--precision", - type=click.Choice(["fp32", "fp16", "bf16"]), - default="fp32", - help="GPU inference precision: fp32 (default), fp16, or bf16. fp16/bf16 reduce memory and may speed up inference on compatible GPUs.", -) -@click.option( - "--xla", - is_flag=True, - help="Enable XLA JIT compilation for inference. May provide 2-3x speedup on GPU after initial compilation overhead.", -) -@click.option( - "-v", - "--verbose", - count=True, - help="Verbosity level: -vv debug, -v info", - default=1, -) -@click.option("-f", "--overwrite", is_flag=True, help="Overwrite existing database") -def build(**kwargs): - model = kwargs.get("model") - if kwargs.get("config") is None: - if model is None: - model = "default" - elif model not in AVAIL_MODELS: - raise click.BadParameter( - f"Model '{model}' is not one of the available options: {', '.join(AVAIL_MODELS)}" - ) - else: - if model is None: - model = "default" - - """Run jaeger taxonomy database generation pipeline""" - if model == "default": - raise click.BadParameter(f"Model '{model}' is not supported") - - else: - from jaeger.commands.taxonomy import build_taxdb - - build_taxdb(**kwargs) - - -@taxonomy.command( - "predict", - context_settings=dict(ignore_unknown_options=True), - help=""" - Use exeperimental taxonomy prediction workflow - - usage - ----- - - jaeger taxonomy predict [OPTIONS] -i contigs.fasta -d taxonomy_db -o output_dir - - """, -) -@click.option( - "-i", - "--input", - type=click.Path(exists=True), - required=True, - help="Path to input file", -) -@click.option( - "-d", - "--db", - type=click.Path(exists=True), - required=True, - help="Path to taxonomy database", -) -@click.option( - "-o", "--output", type=str, required=True, help="Path to output directory" -) -@click.option( - "--fsize", - type=int, - default=2048, - help="Length of the sliding window", -) -@click.option( - "--stride", - type=int, - default=2048, - help="The gap between two sliding of the sliding window (fsize = stride)", -) -@click.option( - "-m", - "--model", - default="default", - help=( - f"Select a deep-learning model to use. " - f"Available models: (when --config is not set): {', '.join(AVAIL_MODELS)}" - ), -) -@click.option( - "--config", - type=click.Path(exists=True), - help="Path to Jaeger config file (e.g., when using Apptainer or Docker)", -) -@click.option( - "--rc", - type=float, - default=0.1, - help="Minimum reliability score required to accept predictions", -) -@click.option( - "--batch", - type=int, - default=96, - help="Parallel batch size, lower if GPU runs out of memory", -) -@click.option("--workers", type=int, default=4, help="Number of threads to use") -@click.option( - "--cpu", - is_flag=True, - help="Ignore available GPUs and explicitly run on CPU", -) -@click.option( - "--physicalid", - type=int, - default=0, - help="Set default GPU device ID for multi-GPU systems", -) -@click.option( - "--mem", - type=int, - default=4, - help="GPU memory limit", -) -@click.option( - "--precision", - type=click.Choice(["fp32", "fp16", "bf16"]), - default="fp32", - help="GPU inference precision: fp32 (default), fp16, or bf16. fp16/bf16 reduce memory and may speed up inference on compatible GPUs.", -) -@click.option( - "--xla", - is_flag=True, - help="Enable XLA JIT compilation for inference. May provide 2-3x speedup on GPU after initial compilation overhead.", -) -@click.option( - "-v", - "--verbose", - count=True, - help="Verbosity level: -vv debug, -v info", - default=1, -) -@click.option("-f", "--overwrite", is_flag=True, help="Overwrite existing database") -def predict_tax(**kwargs): # noqa: F811 - model = kwargs.get("model") - if kwargs.get("config") is None: - if model is None: - model = "default" - elif model not in AVAIL_MODELS: - raise click.BadParameter( - f"Model '{model}' is not one of the available options: {', '.join(AVAIL_MODELS)}" - ) - else: - if model is None: - model = "default" - - """Run jaeger taxonomy database generation pipeline""" - if model == "default": - raise click.BadParameter(f"Model '{model}' is not supported") - - else: - from jaeger.commands.taxonomy import predict_taxonomy - - predict_taxonomy(**kwargs) - - main.add_command(health) main.add_command(predict) main.add_command(train) main.add_command(register_models) main.add_command(download) main.add_command(utils) -main.add_command(taxonomy) if __name__ == "__main__": diff --git a/src/jaeger/commands/convert_graph.py b/src/jaeger/commands/convert_graph.py deleted file mode 100644 index 4b775bf..0000000 --- a/src/jaeger/commands/convert_graph.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Graph conversion CLI for Jaeger. - -Thin Click wrapper around `jaeger.nnlib.conversion`. -""" - -from __future__ import annotations - -from pathlib import Path - -import click - -from jaeger.nnlib.conversion import convert_graph - - -@click.command() -@click.option( - "-m", - "--model", - type=str, - required=True, - help="Model name to convert (e.g., jaeger_57341_1.5M_fragment)", -) -@click.option( - "-o", - "--output", - type=click.Path(path_type=Path), - required=True, - help="Output directory for the converted model", -) -@click.option( - "--mode", - type=click.Choice(["xla", "tflite", "onnx", "tensorrt"]), - default="xla", - help="Conversion mode: xla (default), tflite, onnx, or tensorrt", -) -@click.option( - "--int8", - is_flag=True, - help=( - "Apply static INT8 quantization (ONNX mode only). " - "Produces a smaller model that can use TensorRT INT8 tensor cores." - ), -) -@click.option( - "-v", - "--verbose", - count=True, - help="Verbosity level", - default=1, -) -def convert_graph_cmd(model: str, output: Path, mode: str, int8: bool, verbose: int): - """Convert a Jaeger SavedModel to an optimized inference graph. - - Supports four optimization backends: - - \b - xla (default): - JIT-compiles the graph for GPU. Provides 2-3x speedup after - initial compilation. Best for large datasets with repeated shapes. - - \b - tflite: - Converts to TensorFlow Lite for mobile/edge deployment. - Produces a smaller model (~3.5x size reduction). - - \b - onnx: - Converts to ONNX format for cross-platform deployment. - Recommended for TensorRT acceleration with pip-installed TF. - Use --int8 to additionally quantize to INT8 for smaller models - and INT8 tensor-core inference. - - \b - tensorrt: - NVIDIA TensorRT optimization for maximum GPU performance. - Requires TensorFlow built with TensorRT support (not available - in standard pip packages). Use NVIDIA NGC containers. - - Examples: - jaeger utils convert-graph -m default -o ./optimized --mode xla - jaeger utils convert-graph -m default -o ./optimized --mode onnx - jaeger utils convert-graph -m default -o ./optimized --mode onnx --int8 - """ - convert_graph(model, output, mode, verbose, int8=int8) diff --git a/src/jaeger/commands/health.py b/src/jaeger/commands/health.py index 7aaa83e..54a5647 100644 --- a/src/jaeger/commands/health.py +++ b/src/jaeger/commands/health.py @@ -1,61 +1,18 @@ -import os - -os.environ["WRAPT_DISABLE_EXTENSIONS"] = "true" -import json import sys import platform from pathlib import Path from importlib.resources import files from importlib.metadata import version, PackageNotFoundError -import tensorflow as tf -from jaeger.nnlib.inference import JaegerModel -from jaeger.nnlib.v1.layers import WRes_model_embeddings -from jaeger.preprocess.v1.convert import process_string -from jaeger.seqops.io import fragment_generator +import warnings + from jaeger.utils.fs import validate_fasta_entries from jaeger.utils.logging import get_logger -from jaeger.utils.test import test_tf +from jaeger.utils.test import test_torch from jaeger.utils.misc import AvailableModels, json_to_dict -import warnings -from typing import Generator, Any -from jaeger.utils.misc import track_ms as track warnings.filterwarnings("ignore") -class InferModel: - """ - loads a graph given a dict with model graph location and class map - consumnes batched iterators and returns logits per iterator element - """ - - def __init__(self, graph_path): - self.loaded_model = tf.saved_model.load(graph_path) - self.inference_fn = self.loaded_model.signatures["serving_default"] - - @tf.function - def _predict_step(self, batch): - # Unpack the data - x, y = batch[0], batch[1:] - # set model to inference mode - y_logits = self.inference_fn( - inputs=x["forward_1"], - inputs_1=x["forward_2"], - inputs_2=x["forward_3"], - inputs_3=x["reverse_1"], - inputs_4=x["reverse_2"], - inputs_5=x["reverse_3"], - ) - - return {"y_hat": y_logits, "meta": y} - - def predict(self, x) -> Generator[Any, Any, Any]: - accum = [] - for batch in track(x, description="[cyan]Crunching data..."): - accum.append(self._predict_step(batch)) - return accum - - def _get_package_version(pkg: str) -> str: try: return version(pkg) @@ -65,7 +22,6 @@ def _get_package_version(pkg: str) -> str: def _get_jaeger_version() -> str: """Get jaeger-bio version, preferring local pyproject.toml over system package metadata.""" - # Prefer local editable install's pyproject.toml try: import jaeger @@ -77,7 +33,6 @@ def _get_jaeger_version() -> str: return line.split("=")[-1].strip().strip('"') except Exception: pass - # Fallback to importlib.metadata try: return version("jaeger-bio") except PackageNotFoundError: @@ -90,19 +45,16 @@ def _print_diagnostics(logger) -> None: logger.info("Jaeger Health Diagnostics") logger.info("=" * 60) - # Python environment logger.info(f"Python version: {platform.python_version()}") logger.info(f"Python executable: {sys.executable}") logger.info(f"Platform: {platform.platform()}") logger.info(f"Machine: {platform.machine()}") logger.info(f"Processor: {platform.processor()}") - # Core dependency versions logger.info("-" * 40) logger.info("Core Dependencies:") logger.info(f" jaeger-bio: {_get_jaeger_version()}") - logger.info(f" tensorflow: {_get_package_version('tensorflow')}") - logger.info(f" keras: {_get_package_version('keras')}") + logger.info(f" torch: {_get_package_version('torch')}") logger.info(f" numpy: {_get_package_version('numpy')}") logger.info(f" click: {_get_package_version('click')}") logger.info(f" parasail: {_get_package_version('parasail')}") @@ -116,18 +68,18 @@ def _print_diagnostics(logger) -> None: logger.info(f" pycirclize: {_get_package_version('pycirclize')}") logger.info(f" biopython: {_get_package_version('biopython')}") - # TensorFlow hardware info logger.info("-" * 40) - logger.info("TensorFlow Hardware:") - cpus = tf.config.list_physical_devices("CPU") - gpus = tf.config.list_physical_devices("GPU") - logger.info(f" CPUs detected: {len(cpus)}") - logger.info(f" GPUs detected: {len(gpus)}") - for i, gpu in enumerate(gpus): - logger.info(f" GPU {i}: {gpu.name}") - logger.info(f" TF built with CUDA: {tf.test.is_built_with_cuda()}") + logger.info("PyTorch Hardware:") + try: + import torch + + logger.info(" CPUs detected: 1") + logger.info(f" GPUs detected: {torch.cuda.device_count()}") + for i in range(torch.cuda.device_count()): + logger.info(f" GPU {i}: {torch.cuda.get_device_name(i)}") + except Exception as e: + logger.info(f" Could not inspect PyTorch devices: {e}") - # Installed models logger.info("-" * 40) logger.info("Installed Models:") config_path = files("jaeger.data").joinpath("config.json") @@ -150,7 +102,6 @@ def _print_diagnostics(logger) -> None: else: logger.info(" (no models found)") - # Bundled config entries logger.info("-" * 40) logger.info("Bundled Model Configs:") for key in sorted(k for k in config.keys() if k != "model_paths"): @@ -171,9 +122,6 @@ def health_core(**kwargs) -> None: output_path = Path.cwd() / "test_log" output_path.mkdir(parents=True, exist_ok=True) - fsize = 2048 - stride = 2048 - batch = 64 fnames = ["test_short.fasta", "test_empty.fasta", "test_contigs.fasta"] log_file = "test_jaeger.log" @@ -181,15 +129,14 @@ def health_core(**kwargs) -> None: log_path=Path(output_path), log_file=log_file, level=kwargs.get("verbose") ) - # Print diagnostics first _print_diagnostics(logger) - # Test 1-3 + # Test 1-3: validate bundled FASTA files for i, f in enumerate(fnames): input_file = str(files("jaeger.data.test").joinpath(f)) logger.info(input_file) try: - num = validate_fasta_entries(input_file) + validate_fasta_entries(input_file) passed += 1 logger.info(f"{i} test passed {f}!") except Exception as e: @@ -200,63 +147,13 @@ def health_core(**kwargs) -> None: logger.info(f"{i} test passed {f}!") logger.debug(e) - # Test 4 - result = test_tf() + # Test 4: PyTorch smoke test + result = test_torch() if isinstance(result, Exception): - logger.error("4 tensorflow test failed!") + logger.error("4 pytorch test failed!") logger.debug(result) else: passed += 1 - logger.info("4 tensorflow test passed!") + logger.info("4 pytorch test passed!") - # Test 5 - try: - tf.config.set_soft_device_placement(True) - config_path = files("jaeger.data").joinpath("config.json") - config = json.loads(config_path.read_text()) - weights_path = files("jaeger.data.models.default").joinpath( - config["default"]["weights"] - ) - - input_dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - input_file, - fragsize=fsize, - stride=stride, - num=num, - ), - output_signature=(tf.TensorSpec(shape=(), dtype=tf.string)), - ) - - logger.info("loading the dataset") - idataset = input_dataset.map( - process_string(crop_size=fsize), - ).batch(batch) - - inputs, outputs = WRes_model_embeddings( - input_shape=(None,), dropout_active=False - ) - - logger.info("creating the model") - model = JaegerModel(inputs=inputs, outputs=outputs) - model.load_weights(filepath=weights_path) - model.summary() - logger.info(files("jaeger.data.models.test").joinpath("jaeger_fragment_graph")) - tf.saved_model.save( - model, - files("jaeger.data.models.test").joinpath("jaeger_fragment_graph"), - ) - - logger.info("loading the model") - model = InferModel( - files("jaeger.data.models.test").joinpath("jaeger_fragment_graph") - ) - logger.info("starting model inference") - _ = model.predict(idataset) - logger.info("5 test model passed!") - passed += 1 - except Exception as e: - logger.exception("5 test model failed!") - logger.debug(e) - finally: - logger.info(f"{passed}/5 tests passed!") + logger.info(f"{passed}/4 tests passed!") diff --git a/src/jaeger/commands/predict_legacy.py b/src/jaeger/commands/predict_legacy.py deleted file mode 100644 index adc1579..0000000 --- a/src/jaeger/commands/predict_legacy.py +++ /dev/null @@ -1,342 +0,0 @@ -import os -import traceback -import sys -import psutil -import time -import numpy as np -from importlib.resources import files -from importlib.metadata import version -import json -import h5py -import joblib -from pathlib import Path -import tensorflow as tf - -from jaeger.nnlib.inference import JaegerModel -from jaeger.nnlib.v1.layers import WRes_model_embeddings, create_jaeger_model -from jaeger.postprocess.collect import write_fasta_from_results -from jaeger.postprocess.prophages import ( - logits_to_df, - plot_scores, - plot_scores_linear, - prophage_report, - segment, -) -from jaeger.seqops.io import fragment_generator -from jaeger.utils.gpu import get_device_name -from jaeger.utils.termini import scan_for_terminal_repeats -from jaeger.utils.fs import validate_fasta_entries -from jaeger.utils.logging import description, get_logger - -GB_BYTES = 1024**3 - - -def run_core(**kwargs): - current_process = psutil.Process() - MODEL = kwargs.get("model") - DATA_PATH = files("jaeger.data") - MODEL_PATH = DATA_PATH / "models" / MODEL - CONFIG_PATH = DATA_PATH / "config.json" - config = json.loads(CONFIG_PATH.read_text()).get(MODEL) - config["model"] = MODEL - input_file_path = Path(kwargs.get("input")) - input_file = input_file_path.name - file_base = input_file_path.stem - OUTPUT_DIR = Path(kwargs.get("output")) / MODEL - OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - - log_file = Path(f"{file_base}_jaeger.log") - logger = get_logger(OUTPUT_DIR, log_file, level=kwargs.get("verbose")) - - logger.info( - description(version("jaeger-bio")) + "\n{:-^80}".format("validating parameters") - ) - logger.debug(DATA_PATH) - - try: - num = validate_fasta_entries(str(input_file_path), min_len=kwargs.get("fsize")) - except Exception as e: - logger.error(e) - logger.debug(traceback.format_exc()) - sys.exit(1) - - output_table_path = OUTPUT_DIR / f"{file_base}_jaeger.tsv" - output_phage_table_path = OUTPUT_DIR / f"{file_base}_phages_jaeger.tsv" - - if output_table_path.exists() and not kwargs.get("overwrite"): - logger.error( - "output file exists. enable --overwrite option to overwrite the output file." - ) - sys.exit(1) - - weights_path = MODEL_PATH / config["weights"] - num_class = config["num_classes"] - ood_params = None - - if not weights_path.exists(): - logger.error("could not find model weights. please check the data dir") - sys.exit(1) - - if config["ood"]: - ood_weights_path = MODEL_PATH / config["ood"] - - if ood_weights_path.suffix == ".h5": - with h5py.File(ood_weights_path, "r") as hf: - ood_params = { - "type": "params", - "coeff": np.array(hf["coeff"]), - "intercept": np.array(hf["intercept"]), - "batch_mean": np.array(hf["mean_batch"]), - } - elif ood_weights_path.suffix == ".pkl": - mean_file = MODEL_PATH / "batch_means.npy" - std_file = MODEL_PATH / "batch_std.npy" - batch_mean = np.load(mean_file) - batch_std = np.load(std_file) - ood_params = { - "type": "sklearn", - "model": joblib.load(ood_weights_path), - "batch_mean": batch_mean, - "batch_std": batch_std, - } - - gpus = tf.config.list_physical_devices("GPU") - mode = None - - if kwargs.get("cpu"): - mode = "CPU" - tf.config.set_visible_devices([], "GPU") - logger.info("CPU only mode selected") - elif gpus: - mode = "GPU" - tf.config.set_visible_devices([gpus[kwargs.get("physicalid")]], "GPU") - try: - tf.config.set_logical_device_configuration( - gpus[kwargs.get("physicalid")], - [ - tf.config.LogicalDeviceConfiguration( - memory_limit=4096, experimental_device_ordinal=10 - ) - ], - ) - except Exception as e: - logger.error(f"an error {e} occurred during virtual device initialization ") - logger.debug(traceback.format_exc()) - else: - mode = "CPU" - logger.warn( - "could not find a GPU on the system. For optimal performance run Jaeger on a GPU." - ) - - logger.info(f"tensorflow: {version('tensorflow')}") - logger.info(f"input file: {input_file}") - logger.info(f"log file: {log_file.name}") - logger.info(f"outpath: {OUTPUT_DIR.resolve()}") - logger.info(f"fragment size: {kwargs.get('fsize')}") - logger.info(f"stride: {kwargs.get('stride')}") - logger.info(f"batch size: {kwargs.get('batch')}") - logger.info(f"mode: {mode}") - logger.info(f"avail mem: {psutil.virtual_memory().available / (GB_BYTES):.2f}GB") - logger.info(f"avail cpus: {psutil.cpu_count()}") - logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") - logger.info( - f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB ({current_process.memory_percent():.2f}%)" - ) - - term_repeats = scan_for_terminal_repeats( - file_path=str(input_file_path), - num=num, - workers=kwargs.get("workers"), - fsize=kwargs.get("fsize"), - ) - - device = tf.config.list_logical_devices(mode) - device_names = [get_device_name(i) for i in device] - logger.debug(f"{device}, {device_names}") - if len(device) > 1: - logger.info(f"Using MirroredStrategy {device_names}") - strategy = tf.distribute.MirroredStrategy(device_names) - else: - logger.info(f"Using OneDeviceStrategy {device_names}") - strategy = tf.distribute.OneDeviceStrategy(device_names[0]) - - tf.config.set_soft_device_placement(True) - with strategy.scope(): - input_dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - str(input_file_path), - no_progress=False, - fragsize=kwargs.get("fsize"), - stride=kwargs.get("stride"), - num=num, - ), - output_signature=(tf.TensorSpec(shape=(), dtype=tf.string)), - ) - if kwargs.get("model") == "default": - from jaeger.preprocess.v1.convert import process_string - - idataset = ( - input_dataset.map( - process_string(crop_size=kwargs.get("fsize")), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .batch(kwargs.get("batch"), num_parallel_calls=tf.data.AUTOTUNE) - .prefetch(25) - ) - inputs, outputs = WRes_model_embeddings( - input_shape=(None,), dropout_active=False - ) - else: - from jaeger.preprocess.v2.convert import process_string - - insize = (int(kwargs.get("fsize")) // 3) - 1 - idataset = ( - input_dataset.map( - process_string(crop_size=kwargs.get("fsize")), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .batch(kwargs.get("batch"), num_parallel_calls=tf.data.AUTOTUNE) - .prefetch(50) - ) - - inputs, outputs = create_jaeger_model( - input_shape=(6, insize, 11), out_shape=num_class - ) - model = JaegerModel(inputs=inputs, outputs=outputs) - - logger.info("loading model to memory") - model.load_weights(filepath=weights_path) - logger.info( - f"avail mem : {psutil.virtual_memory().available / (GB_BYTES): .2f}GB\n{'-' * 80}" - ) - - if mode == "GPU": - for device_number, d in enumerate(device_names): - gpu_mem = tf.config.experimental.get_memory_info(d) - logger.info( - f"GPU {device_number} current : {gpu_mem['current'] / (GB_BYTES): .2f}GB peak : {gpu_mem['peak'] / (GB_BYTES): .2f}GB" - ) - - try: - logger.info("starting model inference") - y_pred = model.predict(idataset, verbose=0) - except Exception as e: - logger.debug(traceback.format_exc()) - logger.error( - f"an error {e} occured during inference on {'|'.join(device_names)}! check {log_file} for traceback." - ) - sys.exit(1) - - from jaeger.postprocess.collect import pred_to_dict_legacy, write_output_legacy - - if kwargs.get("getalllabels"): - config["labels"] = [v for k, v in config["all_labels"].items()] - else: - config["labels"] = [v for k, v in config["default_labels"].items()] - - data, data_full = pred_to_dict_legacy( - config, - y_pred, - model=kwargs.get("model"), - fsize=kwargs.get("fsize"), - ood_params=ood_params, - term_repeats=term_repeats, - ) - - write_output_legacy( - config, - data, - output_table_path=output_table_path, - output_phage_table_path=output_phage_table_path, - reliability_cutoff=kwargs.get("rc", 0.5), - phage_score=kwargs.get("pc", 3), - ) - - logger.info(f"processed {data.get('headers').shape[0]}/{num} sequences") - - if kwargs.get("getsequences"): - output_fasta_file = f"{file_base}_{config['suffix']}_phages_jaeger.fasta" - output_fasta_file_path = OUTPUT_DIR / output_fasta_file - write_fasta_from_results( - input_fasta=input_file_path, - output_tsv=output_phage_table_path, - output_fasta=output_fasta_file_path, - ) - logger.info(f"{output_fasta_file} created") - - if kwargs.get("window_scores"): - output_scores = f"{file_base}_{config['suffix']}_window_scores.npz" - output_scores_path = os.path.join(OUTPUT_DIR, output_scores) - logger.info( - f"writing window-wise scores and metadata to {output_scores_path}" - ) - np.savez( - output_scores_path, - headers=data_full["headers"], - lengths=data_full["lengths"], - predictions=np.array(data_full["predictions"], dtype=object), - gc_skews=np.array(data_full["gc_skews"], dtype=object), - gcs=np.array(data_full["gcs"], dtype=object), - ) - logger.info(f"{output_scores} created") - - if kwargs.get("prophage"): - # still experimental - needs more testing!!!! - try: - if logits_df := logits_to_df( - config, cmdline_kwargs=kwargs, **data_full - ): - logger.info("identifying prophages") - pro_dir = OUTPUT_DIR / f"{file_base}_{config['suffix']}_prophages" - plots_dir = pro_dir / "plots" - for dir in [pro_dir, plots_dir]: - dir.mkdir(parents=True, exist_ok=True) - - phage_cord = segment( - logits_df, - outdir=plots_dir, - cutoff_length=kwargs.get("lc"), - sensitivity=kwargs.get("sensitivity"), - identifier="phage", - ) - plot_type = kwargs.get("plot_type", "circular") - if plot_type in ("circular", "both"): - plot_scores( - logits_df, - config=config, - model=MODEL, - fsize=kwargs.get("fsize"), - infile_base=file_base, - outdir=plots_dir, - phage_cordinates=phage_cord, - ) - if plot_type in ("linear", "both"): - plot_scores_linear( - logits_df, - config=config, - model=MODEL, - fsize=kwargs.get("fsize"), - infile_base=file_base, - outdir=plots_dir, - phage_cordinates=phage_cord, - ) - prophage_report( - fsize=kwargs.get("fsize"), - filehandle=str(input_file_path), - prophage_cordinates=phage_cord, - outdir=pro_dir, - ) - else: - logger.info("no prophage regions found") - except Exception as e: - logger.error( - f"an error {e} occured during the prophage prediction step" - ) - logger.debug(traceback.format_exc()) - - logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") - logger.info( - f"memory usage : {current_process.memory_full_info().rss / 1024**3:.2f}Gb ({current_process.memory_percent():.2f}%)" - ) diff --git a/src/jaeger/commands/quantize.py b/src/jaeger/commands/quantize.py deleted file mode 100644 index 34d2fd5..0000000 --- a/src/jaeger/commands/quantize.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Model quantization utilities for Jaeger. - -Supports: -- Dynamic range quantization (INT8 weights, FP32 activations) -- Full integer quantization (INT8 weights and activations) -- Float16 quantization (FP16 weights, FP32 activations) - -Note: Quantization is performed via TensorFlow Lite conversion. -The quantized model is saved as a TFLite model alongside metadata. - -Important: - TFLite models produced by this tool are primarily intended for - edge/mobile deployment where model size matters. On desktop GPUs, - TFLite inference may not be faster than the original SavedModel - due to fallback TF ops and interpreter overhead. -""" - -from __future__ import annotations - -import logging -import shutil -import sys -from pathlib import Path - -import click -import numpy as np -import tensorflow as tf -import yaml -from tensorflow.python.framework.convert_to_constants import ( - convert_variables_to_constants_v2, -) - -from jaeger.utils.misc import AvailableModels - -logger = logging.getLogger("Jaeger") - - -def quantize_model(model: str, output: Path, mode: str, verbose: int): - """Core quantization logic (non-CLI).""" - log = logging.getLogger("Jaeger") - log.info(f"Quantizing model '{model}' with mode '{mode}'") - - # Resolve model path - model_info = _resolve_model(model) - if model_info is None: - log.error(f"Model '{model}' not found") - sys.exit(1) - - graph_dir = Path(model_info["graph"]) - output = Path(output) - output.mkdir(parents=True, exist_ok=True) - - # Create output subdirectory - quantized_dir = output / f"{model}_{mode}" - if quantized_dir.exists(): - shutil.rmtree(quantized_dir) - quantized_dir.mkdir(parents=True) - - # Copy metadata files - for key in ("classes", "project", "weights"): - if key in model_info and model_info[key]: - src = Path(model_info[key]) - dst = quantized_dir / src.name - shutil.copy2(src, dst) - log.info(f"Copied {key}: {src.name}") - - # Perform quantization via frozen graph -> TFLite - # Frozen graph eliminates resource variables that TFLite cannot handle. - log.info(f"Loading SavedModel from {graph_dir}") - loaded = tf.saved_model.load(str(graph_dir)) - infer = loaded.signatures["serving_default"] - - log.info("Freezing graph (converting variables to constants)...") - frozen_func = convert_variables_to_constants_v2(infer) - - log.info(f"Converting to TFLite (mode={mode})...") - converter = tf.lite.TFLiteConverter.from_concrete_functions([frozen_func]) - # Jaeger models use ops (e.g., Gelu, BatchMatMul) that are not in the - # baseline TFLite builtin set; allow TensorFlow fallback ops. - converter.target_spec.supported_ops = [ - tf.lite.OpsSet.TFLITE_BUILTINS, - tf.lite.OpsSet.SELECT_TF_OPS, - ] - - if mode == "dynamic": - converter.optimizations = [tf.lite.Optimize.DEFAULT] - elif mode == "float16": - converter.optimizations = [tf.lite.Optimize.DEFAULT] - converter.target_spec.supported_types = [tf.float16] - elif mode == "full_int8": - log.warning( - "full_int8 mode is experimental and may reduce accuracy. " - "Use a representative dataset from your target domain for best results." - ) - converter.optimizations = [tf.lite.Optimize.DEFAULT] - converter.representative_dataset = _make_rep_dataset() - converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] - # Keep input/output as FLOAT32 so the TFLite runtime handles - # quantization/dequantization automatically. This avoids needing - # to manually quantize inputs in the inference pipeline. - # converter.inference_input_type = tf.int8 - # converter.inference_output_type = tf.int8 - - try: - tflite_model = converter.convert() - except Exception as e: - log.error(f"Quantization failed: {e}") - sys.exit(1) - - # Save TFLite model - tflite_path = quantized_dir / f"{model}_{mode}.tflite" - tflite_path.write_bytes(tflite_model) - log.info( - f"Saved quantized model: {tflite_path} ({len(tflite_model) / 1024:.1f} KB)" - ) - - # Save quantization info - info = { - "original_graph": str(graph_dir), - "quantization_mode": mode, - "tflite_model": str(tflite_path.name), - "model_format": "tflite", - "note": ( - "TFLite model for edge deployment. " - "Resize interpreter input tensor at runtime for variable sequence lengths." - ), - } - info_path = quantized_dir / "quantization_info.yaml" - info_path.write_text(yaml.safe_dump(info)) - - log.info(f"Quantization complete. Output: {quantized_dir}") - - -@click.command() -@click.option( - "-m", - "--model", - type=str, - required=True, - help="Model name to quantize (e.g., jaeger_57341_1.5M_fragment)", -) -@click.option( - "-o", - "--output", - type=click.Path(path_type=Path), - required=True, - help="Output directory for the quantized model", -) -@click.option( - "--mode", - type=click.Choice(["dynamic", "float16", "full_int8"]), - default="dynamic", - help=( - "Quantization mode. dynamic (default) is recommended. " - "full_int8 is experimental and may reduce accuracy." - ), -) -@click.option( - "-v", - "--verbose", - count=True, - help="Verbosity level", - default=1, -) -def quantize(model: str, output: Path, mode: str, verbose: int): - """Quantize a Jaeger model for edge deployment. - - Converts a SavedModel to a quantized TensorFlow Lite model. The primary - benefit is reduced model size (e.g., ~6 MB -> ~1.6 MB with dynamic - quantization). Note that TFLite inference speed depends heavily on the - target hardware and may not exceed the original SavedModel on desktop GPUs. - - Examples: - jaeger utils quantize -m default -o ./quantized_models - jaeger utils quantize -m jaeger_57341_1.5M_fragment -o ./quantized --mode float16 - """ - quantize_model(model, output, mode, verbose) - - -def _resolve_model(model_name: str) -> dict | None: - """Resolve model name to path dict.""" - from importlib.resources import files - - from jaeger.utils.misc import json_to_dict - - CONFIG_PATH = files("jaeger.data") / "config.json" - model_paths = json_to_dict(CONFIG_PATH).get("model_paths", []) - models = AvailableModels(path=model_paths) - return models.info.get(model_name) - - -def _make_rep_dataset(): - """Create a representative dataset generator for INT8 quantization.""" - - def dataset_gen(): - for _ in range(100): - yield [np.random.randn(1, 6, 100, 64).astype(np.float32)] - - return dataset_gen diff --git a/src/jaeger/commands/taxonomy.py b/src/jaeger/commands/taxonomy.py deleted file mode 100644 index 4c3fb8a..0000000 --- a/src/jaeger/commands/taxonomy.py +++ /dev/null @@ -1,673 +0,0 @@ -# experimental taxonomy predicton workflow -import os - -os.environ["GRPC_VERBOSITY"] = "ERROR" -os.environ["GLOG_minloglevel"] = "2" -import traceback -import sys -import psutil -import shutil -import time - -try: - import taxopy - import faiss -except ModuleNotFoundError as exc: - raise ImportError( - "The taxonomy workflow requires the 'taxopy' and 'faiss' packages. " - "Install them with e.g.: pip install taxopy faiss-cpu" - ) from exc - -import numpy as np -import tensorflow as tf -from typing import Union -from pathlib import Path -from importlib.resources import files -from importlib.metadata import version -from typing import Any, List -from jaeger.nnlib.inference import InferModel -from jaeger.seqops.io import fragment_generator -from jaeger.utils.gpu import get_device_name -from jaeger.utils.fs import validate_fasta_entries -from jaeger.utils.misc import json_to_dict, AvailableModels, get_model_id # noqa: F401 -from jaeger.utils.logging import description, get_logger - - -# Standard abbreviated rank prefixes for the full-lineage string. -_RANK_PREFIXES = { - "superkingdom": "d", - "kingdom": "k", - "subkingdom": "sk", - "superphylum": "sp", - "phylum": "p", - "subphylum": "sph", - "superclass": "sc", - "class": "c", - "subclass": "ssc", - "infraclass": "ic", - "superorder": "so", - "order": "o", - "suborder": "sor", - "infraorder": "io", - "parvorder": "po", - "superfamily": "sf", - "family": "f", - "subfamily": "sfa", - "tribe": "t", - "subtribe": "st", - "genus": "g", - "subgenus": "sg", - "species group": "sg", - "species subgroup": "ssg", - "species": "s", - "subspecies": "ss", - "forma": "fo", - "varietas": "v", - "biotype": "b", - "realm": "r", - "subrealm": "sr", -} - - -def _format_ranked_lineage(lca) -> str: - """Return a `r__Name;k__Name;...` lineage string from a taxopy Taxon.""" - parts = [] - for rank, name in getattr(lca, "rank_name_dictionary", {}).items(): - prefix = _RANK_PREFIXES.get(rank, rank[0] if rank else "?") - parts.append(f"{prefix}__{name}") - return ";".join(parts) - - -class TaxonomyModel: - def __init__( - self, - faiss_path: Union[str, Path], - index2taxid_path: Union[str, Path], - taxdump_path: Union[str, Path], - ): - self.index2taxid = self.load_index2taxid(path=index2taxid_path) - self.faiss_index = self.load_faiss_index(path=faiss_path) - self.taxdb = self.load_taxdb(path=taxdump_path) - # Cache Taxon objects because constructing them is expensive. - self._taxon_cache: dict[int, Any] = {} - - def load_taxdb(self, path: Union[str, Path]): - return taxopy.TaxDb( - nodes_dmp=path / "nodes.dmp", - names_dmp=path / "names.dmp", - merged_dmp=path / "merged.dmp", - ) - - def load_index2taxid(self, path): - data = np.load(path) - return data["taxids"] - - def load_faiss_index(self, path): - return faiss.read_index(str(path)) - - def _get_taxon(self, taxid: int): - taxon = self._taxon_cache.get(taxid) - if taxon is None: - taxon = taxopy.Taxon(taxid, self.taxdb) - self._taxon_cache[taxid] = taxon - return taxon - - def get_lca(self, taxids: List[int], fraction=0.6): - taxons = [self._get_taxon(i) for i in taxids] - if len(taxons) > 1: - return taxopy.find_majority_vote( - taxons, taxdb=self.taxdb, fraction=fraction - ) - return taxons[0] - - def predict(self, embeddings: List, headers: List, k: int = 1): - """Batch all contig windows into a single FAISS search. - - ``embeddings`` is a list of 2-D arrays, one per contig. The result is an - LCA call per contig using the ``k`` nearest neighbors of every window. - """ - if not embeddings: - return [] - - # Stack all windows from all contigs into one contiguous float32 array. - all_embeddings = np.vstack(embeddings).astype(np.float32, copy=False) - all_embeddings = np.ascontiguousarray(all_embeddings) - faiss.normalize_L2(all_embeddings) - - _, all_indices = self.faiss_index.search(all_embeddings, k) - all_taxids = self.index2taxid[all_indices.flatten()] - - # Map each window back to its contig. - offsets = np.cumsum([0] + [e.shape[0] for e in embeddings]) * k - - tmp_l = [] - for idx, (embedding, header) in enumerate(zip(embeddings, headers)): - start = offsets[idx] - end = offsets[idx + 1] - taxids = all_taxids[start:end] - tmp_l.append({"header": header, "lca": self.get_lca(taxids)}) - return tmp_l - - -def copy_tax_files(in_tax, output_tax): - src = Path(in_tax) - dst = Path(output_tax) - - dst.mkdir(parents=True, exist_ok=True) # Ensure destination exists - - for fname in ["names.dmp", "nodes.dmp", "merged.dmp"]: - shutil.copy(src / fname, dst / fname) - - -def unravel(y_pred: dict, **kwargs) -> tuple[Any, Any, Any]: - """ - Processes model predictions and associated metadata into structured dictionaries. - num_classes - fsize - - Returns: - data (dict): Core prediction metrics and metadata - data_full (dict): Full output for auxiliary processing - """ - - # -- Step 1: Determine split points - # split_flags = np.array(y_pred["meta_2"], dtype=np.int32) - # split_indices = np.where(split_flags == 1)[0] + 1 - - # if y_pred["prediction"].shape[0] == split_indices[-1]: - # split_indices = split_indices[:-1] - - # -- Step 2: Split predictions and embeddings - # _ = np.split(y_pred["prediction"], split_indices, axis=0) - # _ = np.split(y_pred["reliability"], split_indices, axis=0) - # embeddings = np.split(y_pred["embedding"], split_indices, axis=0) - - # -- headers - # headers = np.split(y_pred["meta_0"], split_indices, axis=0) - # return (embeddings, headers) - - taxids = [kwargs.get("acc2taxid").get(i.decode()) for i in y_pred["meta_0"]] - - return (y_pred["embedding"], taxids, y_pred["meta_0"]) - - -def unravel_inference(y_pred: dict) -> tuple[Any, Any]: - """ - Processes model predictions and associated metadata into tuple. - - Returns: - embeddings - headers - """ - - # -- Step 1: Determine split points - split_flags = np.array(y_pred["meta_2"], dtype=np.int32) - split_indices = np.where(split_flags == 1)[0] + 1 - - if y_pred["prediction"].shape[0] == split_indices[-1]: - split_indices = split_indices[:-1] - - # -- Step 2: Split predictions and embeddings - _ = np.split(y_pred["prediction"], split_indices, axis=0) - _ = np.split(y_pred["reliability"], split_indices, axis=0) - embeddings = np.split(y_pred["embedding"], split_indices, axis=0) - - # -- headers - headers = [i[0] for i in np.split(y_pred["meta_0"], split_indices, axis=0)] - return (embeddings, headers) - - -def create_cosine_index(vectors: np.ndarray, save_path: str) -> None: - """ - Creates a cosine similarity FAISS index and saves it to disk. - - Args: - vectors (np.ndarray): 2D array of shape (n_samples, dim). Will be cast - to float32 if necessary. - save_path (str): Path to save the FAISS index (e.g., 'index.faiss'). - """ - assert vectors.ndim == 2, "Input must be a 2D array" - vectors = vectors.astype(np.float32, copy=False) - - faiss.normalize_L2(vectors) # Normalize vectors for cosine similarity - - dim = vectors.shape[1] - index = faiss.IndexFlatIP( - dim - ) # Inner product = cosine similarity (after normalization) - - index.add(vectors) # Add vectors to index - faiss.write_index(index, str(save_path)) - - -def get_acc2taxid(path: str): - map_ = {} - with open(path, "r") as fh: - for index, line in enumerate(fh): - line = line.strip() - if index == 0 or not line: - continue - fields = line.split("\t") - if len(fields) < 3: - continue - map_[fields[1]] = int(fields[2]) - return map_ - - -def build_taxdb(**kwargs): - # acc2tax, tax, input, out - current_process = psutil.Process() - GB_BYTES = 1024**3 - MODEL = kwargs.get("model") - DATA_PATH = files("jaeger.data") - if kwargs.get("config") is None: - CONFIG_PATH = DATA_PATH / "config.json" - else: - CONFIG_PATH = kwargs.get("config") - - USER_MODEL_PATHS = json_to_dict(CONFIG_PATH).get("model_paths") - MODEL_INFO = AvailableModels(path=USER_MODEL_PATHS).info[MODEL] - MEMORY_LIMIT = 1024 * kwargs.get("mem", 4) - THREADS = kwargs.get("workers") - input_file_path = Path(kwargs.get("input")) - input_file = input_file_path.name - file_base = input_file_path.stem - - OUTPUT_DIR = Path(kwargs.get("output")) - OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - acc2taxid = get_acc2taxid(kwargs.get("acc2tax")) - log_file = Path(f"{file_base}_jaeger.log") - logger = get_logger(OUTPUT_DIR, log_file, level=kwargs.get("verbose")) - logger.info( - description(version("jaeger-bio")) + "\n{:-^80}".format("validating parameters") - ) - logger.debug(DATA_PATH) - logger.debug(AvailableModels(path=USER_MODEL_PATHS).info) - logger.debug(MODEL_INFO) - - try: - num = validate_fasta_entries(str(input_file_path), min_len=kwargs.get("fsize")) - except Exception as e: - logger.error(e) - logger.debug(traceback.format_exc()) - sys.exit(1) - - output_faissdb = OUTPUT_DIR / "genomes.faiss" - output_tax = OUTPUT_DIR / "tax" - output_tax.mkdir(exist_ok=True) - output_index2tax = OUTPUT_DIR / "index2tax.npz" - - if output_faissdb.exists() and not kwargs.get("overwrite"): - logger.error( - "a faiss index exists in the output path. enable --overwrite option to overwrite the output file." - ) - sys.exit(1) - - if not MODEL_INFO["graph"].exists(): - logger.error(f"could not find model graph. please check {USER_MODEL_PATHS}") - sys.exit(1) - tf.config.threading.set_inter_op_parallelism_threads(THREADS) - tf.config.threading.set_intra_op_parallelism_threads(THREADS) - tf.config.set_soft_device_placement(True) - gpus = tf.config.list_physical_devices("GPU") - mode = None - - if kwargs.get("cpu"): - mode = "CPU" - tf.config.set_visible_devices([], "GPU") - logger.info("CPU only mode selected") - elif gpus: - mode = "GPU" - tf.config.set_visible_devices([gpus[kwargs.get("physicalid")]], "GPU") - - # Set mixed precision policy if requested - precision = kwargs.get("precision", "fp32") - if precision == "fp16": - tf.keras.mixed_precision.set_global_policy("mixed_float16") - logger.info("GPU precision: mixed_float16 (FP16 compute, FP32 variables)") - elif precision == "bf16": - tf.keras.mixed_precision.set_global_policy("mixed_bfloat16") - logger.info("GPU precision: mixed_bfloat16 (BF16 compute, FP32 variables)") - else: - logger.info("GPU precision: float32 (FP32)") - - try: - tf.config.set_logical_device_configuration( - gpus[kwargs.get("physicalid")], - [ - tf.config.LogicalDeviceConfiguration( - memory_limit=MEMORY_LIMIT, experimental_device_ordinal=10 - ) - ], - ) - except Exception as e: - logger.error(f"an error {e} occurred during virtual device initialization ") - logger.debug(traceback.format_exc()) - else: - mode = "CPU" - logger.warning( - "could not find a GPU on the system. For optimal performance run Jaeger on a GPU." - ) - - logger.info(f"tensorflow: {version('tensorflow')}") - logger.info(f"input file: {input_file}") - logger.info(f"log file: {log_file.name}") - logger.info(f"outpath: {OUTPUT_DIR.resolve()}") - logger.info(f"fragment size: {kwargs.get('fsize')}") - logger.info(f"stride: {kwargs.get('stride')}") - logger.info(f"batch size: {kwargs.get('batch')}") - logger.info(f"mode: {mode}") - logger.info(f"avail mem: {psutil.virtual_memory().available / (GB_BYTES):.2f}GB") - logger.info( - f"intra threads: {tf.config.threading.get_intra_op_parallelism_threads()}" - ) - logger.info( - f"inter threads: {tf.config.threading.get_inter_op_parallelism_threads()}" - ) - logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") - logger.info( - f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB ({current_process.memory_percent():.2f}%)" - ) - - device = tf.config.list_logical_devices(mode) - device_names = [get_device_name(i) for i in device] - logger.debug(f"{device}, {device_names}") - if len(device) > 1: - logger.info(f"Using MirroredStrategy {device_names}") - strategy = tf.distribute.MirroredStrategy(device_names) - else: - logger.info(f"Using OneDeviceStrategy {device_names}") - strategy = tf.distribute.OneDeviceStrategy(device_names[0]) - - use_xla = kwargs.get("xla", False) - if use_xla: - logger.info( - "Using XLA-compiled inference (first batch may be slow due to compilation)" - ) - try: - logger.info("loading model") - model = InferModel(MODEL_INFO, use_xla=use_xla, return_embedding=True) - except ValueError as e: - logger.error(str(e)) - sys.exit(1) - string_processor_config = model.string_processor_config - logger.info("generating sequence fragments") - input_dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - str(input_file_path), - no_progress=False, - fragsize=kwargs.get("fsize"), - stride=kwargs.get("stride"), - num=num, - ), - output_signature=(tf.TensorSpec(shape=(), dtype=tf.string)), - ) - - from jaeger.seqops.encode import process_string_inference - - logger.info("preprocessing fragments") - idataset = ( - input_dataset.map( - process_string_inference( - crop_size=kwargs.get("fsize"), - codons=string_processor_config.get("codon"), - codon_num=string_processor_config.get("codon_id"), - codon_depth=string_processor_config.get("codon_depth"), - ngram_width=string_processor_config.get("ngram_width"), - seq_onehot=string_processor_config.get("seq_onehot"), - input_type=string_processor_config.get("input_type"), - masking=string_processor_config.get("masking"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle=string_processor_config.get("shuffle"), - ), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .batch(kwargs.get("batch"), num_parallel_calls=tf.data.AUTOTUNE) - .prefetch(25) - ) - - with strategy.scope(): - try: - logger.info("running embedding inference") - y_pred = model.predict(idataset, no_progress=True) - except Exception as e: - logger.debug(traceback.format_exc()) - logger.error( - f"an error {e} occured during generating embeddings on {'|'.join(device_names)}! check {log_file} for traceback." - ) - sys.exit(1) - - # save embeddings and headers - logger.info("saving embeddings") - embeddings, taxids, _ = unravel(y_pred, acc2taxid=acc2taxid) - np.savez(OUTPUT_DIR / "embeddings.npz", embeddings=embeddings, taxids=taxids) - np.savez(output_index2tax, taxids=taxids) - # build faiss db - logger.info("building FAISS index") - create_cosine_index(embeddings, save_path=output_faissdb) - # cp taxdump - logger.info("copying taxdump files") - copy_tax_files(in_tax=kwargs.get("tax"), output_tax=output_tax) - - -def predict_taxonomy(**kwargs): - # acc2tax, tax, input, out - current_process = psutil.Process() - GB_BYTES = 1024**3 - MODEL = kwargs.get("model") - DATA_PATH = files("jaeger.data") - if kwargs.get("config") is None: - CONFIG_PATH = DATA_PATH / "config.json" - else: - CONFIG_PATH = kwargs.get("config") - - USER_MODEL_PATHS = json_to_dict(CONFIG_PATH).get("model_paths") - MODEL_INFO = AvailableModels(path=USER_MODEL_PATHS).info[MODEL] - MEMORY_LIMIT = 1024 * kwargs.get("mem", 4) - THREADS = kwargs.get("workers") - input_file_path = Path(kwargs.get("input")) - input_file = input_file_path.name - file_base = input_file_path.stem - - OUTPUT_DIR = Path(kwargs.get("output")) - OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - DATABASE_DIR = Path(kwargs.get("db")) - - log_file = Path(f"{file_base}_jaeger.log") - logger = get_logger(OUTPUT_DIR, log_file, level=kwargs.get("verbose")) - logger.info( - description(version("jaeger-bio")) + "\n{:-^80}".format("validating parameters") - ) - logger.debug(DATA_PATH) - logger.debug(AvailableModels(path=USER_MODEL_PATHS).info) - logger.debug(MODEL_INFO) - - # load taxdb - logger.info("loading taxonomy database") - tax_model = TaxonomyModel( - faiss_path=DATABASE_DIR / "genomes.faiss", - taxdump_path=DATABASE_DIR / "tax", - index2taxid_path=DATABASE_DIR / "index2tax.npz", - ) - - try: - num = validate_fasta_entries(str(input_file_path), min_len=kwargs.get("fsize")) - except Exception as e: - logger.error(e) - logger.debug(traceback.format_exc()) - sys.exit(1) - - output_taxonomy = OUTPUT_DIR / f"{file_base}_taxonomy.tsv" - - if output_taxonomy.exists() and not kwargs.get("overwrite"): - logger.error( - "output path is not empty. enable --overwrite option to overwrite the output file." - ) - sys.exit(1) - - if not MODEL_INFO["graph"].exists(): - logger.error(f"could not find model graph. please check {USER_MODEL_PATHS}") - sys.exit(1) - tf.config.threading.set_inter_op_parallelism_threads(THREADS) - tf.config.threading.set_intra_op_parallelism_threads(THREADS) - tf.config.set_soft_device_placement(True) - gpus = tf.config.list_physical_devices("GPU") - mode = None - - if kwargs.get("cpu"): - mode = "CPU" - tf.config.set_visible_devices([], "GPU") - logger.info("CPU only mode selected") - elif gpus: - mode = "GPU" - tf.config.set_visible_devices([gpus[kwargs.get("physicalid")]], "GPU") - - # Set mixed precision policy if requested - precision = kwargs.get("precision", "fp32") - if precision == "fp16": - tf.keras.mixed_precision.set_global_policy("mixed_float16") - logger.info("GPU precision: mixed_float16 (FP16 compute, FP32 variables)") - elif precision == "bf16": - tf.keras.mixed_precision.set_global_policy("mixed_bfloat16") - logger.info("GPU precision: mixed_bfloat16 (BF16 compute, FP32 variables)") - else: - logger.info("GPU precision: float32 (FP32)") - - try: - tf.config.set_logical_device_configuration( - gpus[kwargs.get("physicalid")], - [ - tf.config.LogicalDeviceConfiguration( - memory_limit=MEMORY_LIMIT, experimental_device_ordinal=10 - ) - ], - ) - except Exception as e: - logger.error(f"an error {e} occurred during virtual device initialization ") - logger.debug(traceback.format_exc()) - else: - mode = "CPU" - logger.warning( - "could not find a GPU on the system. For optimal performance run Jaeger on a GPU." - ) - - logger.info(f"tensorflow: {version('tensorflow')}") - logger.info(f"input file: {input_file}") - logger.info(f"log file: {log_file.name}") - logger.info(f"outpath: {OUTPUT_DIR.resolve()}") - logger.info(f"fragment size: {kwargs.get('fsize')}") - logger.info(f"stride: {kwargs.get('stride')}") - logger.info(f"batch size: {kwargs.get('batch')}") - logger.info(f"mode: {mode}") - logger.info(f"avail mem: {psutil.virtual_memory().available / (GB_BYTES):.2f}GB") - logger.info( - f"intra threads: {tf.config.threading.get_intra_op_parallelism_threads()}" - ) - logger.info( - f"inter threads: {tf.config.threading.get_inter_op_parallelism_threads()}" - ) - logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") - logger.info( - f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB ({current_process.memory_percent():.2f}%)" - ) - - device = tf.config.list_logical_devices(mode) - device_names = [get_device_name(i) for i in device] - logger.debug(f"{device}, {device_names}") - if len(device) > 1: - logger.info(f"Using MirroredStrategy {device_names}") - strategy = tf.distribute.MirroredStrategy(device_names) - else: - logger.info(f"Using OneDeviceStrategy {device_names}") - strategy = tf.distribute.OneDeviceStrategy(device_names[0]) - - use_xla = kwargs.get("xla", False) - if use_xla: - logger.info( - "Using XLA-compiled inference (first batch may be slow due to compilation)" - ) - try: - logger.info("loading model") - model = InferModel(MODEL_INFO, use_xla=use_xla, return_embedding=True) - except ValueError as e: - logger.error(str(e)) - sys.exit(1) - string_processor_config = model.string_processor_config - logger.info("generating sequence fragments") - input_dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - str(input_file_path), - no_progress=False, - fragsize=kwargs.get("fsize"), - stride=kwargs.get("stride"), - num=num, - ), - output_signature=(tf.TensorSpec(shape=(), dtype=tf.string)), - ) - - from jaeger.seqops.encode import process_string_inference - - logger.info("preprocessing fragments") - idataset = ( - input_dataset.map( - process_string_inference( - crop_size=kwargs.get("fsize"), - codons=string_processor_config.get("codon"), - codon_num=string_processor_config.get("codon_id"), - codon_depth=string_processor_config.get("codon_depth"), - ngram_width=string_processor_config.get("ngram_width"), - seq_onehot=string_processor_config.get("seq_onehot"), - input_type=string_processor_config.get("input_type"), - masking=string_processor_config.get("masking"), - mutate=string_processor_config.get("mutate"), - mutation_rate=string_processor_config.get("mutation_rate"), - shuffle=string_processor_config.get("shuffle"), - ), - num_parallel_calls=tf.data.AUTOTUNE, - ) - .batch(kwargs.get("batch"), num_parallel_calls=tf.data.AUTOTUNE) - .prefetch(25) - ) - - with strategy.scope(): - try: - logger.info("running embedding inference") - y_pred = model.predict(idataset, no_progress=True) - except Exception as e: - logger.debug(traceback.format_exc()) - logger.error( - f"an error {e} occured during generating embeddings on {'|'.join(device_names)}! check {log_file} for traceback." - ) - sys.exit(1) - - # unravel embedding vectors - logger.info("running taxonomy assignment") - embeddings, headers = unravel_inference(y_pred) - - predictions = tax_model.predict(embeddings, headers) - - # Write predictions to TSV using pandas for consistency with the rest of the package. - logger.info("writing taxonomy predictions") - import pandas as pd - - rows = [] - for pred in predictions: - header = pred["header"] - if isinstance(header, bytes): - header = header.decode("utf-8") - lca = pred["lca"] - rows.append( - { - "header": header, - "taxid": getattr(lca, "taxid", ""), - "name": getattr(lca, "name", ""), - "rank": getattr(lca, "rank", ""), - "lineage": _format_ranked_lineage(lca), - } - ) - - pd.DataFrame(rows).to_csv(output_taxonomy, sep="\t", index=False) - logger.info(f"taxonomy predictions written to {output_taxonomy}") diff --git a/src/jaeger/commands/test.py b/src/jaeger/commands/test.py deleted file mode 100644 index c3db810..0000000 --- a/src/jaeger/commands/test.py +++ /dev/null @@ -1,146 +0,0 @@ -import os - -os.environ["WRAPT_DISABLE_EXTENSIONS"] = "true" -import json -from pathlib import Path -from importlib.resources import files -import tensorflow as tf -from jaeger.nnlib.inference import JaegerModel -from jaeger.nnlib.v1.layers import WRes_model_embeddings -from jaeger.preprocess.v1.convert import process_string -from jaeger.seqops.io import fragment_generator -from jaeger.utils.fs import validate_fasta_entries -from jaeger.utils.logging import get_logger -from jaeger.utils.test import test_tf -import warnings -from typing import Generator, Any -from jaeger.utils.misc import track_ms as track - -warnings.filterwarnings("ignore") - - -class InferModel: - """ - loads a graph given a dict with model graph location and class map - consumnes batched iterators and returns logits per iterator element - """ - - def __init__(self, graph_path): - self.loaded_model = tf.saved_model.load(graph_path) - self.inference_fn = self.loaded_model.signatures["serving_default"] - - @tf.function - def _predict_step(self, batch): - # Unpack the data - x, y = batch[0], batch[1:] - # set model to inference mode - y_logits = self.inference_fn( - inputs=x["forward_1"], - inputs_1=x["forward_2"], - inputs_2=x["forward_3"], - inputs_3=x["reverse_1"], - inputs_4=x["reverse_2"], - inputs_5=x["reverse_3"], - ) - - return {"y_hat": y_logits, "meta": y} - - def predict(self, x) -> Generator[Any, Any, Any]: - accum = [] - for batch in track(x, description="[cyan]Crunching data..."): - accum.append(self._predict_step(batch)) - return accum - - -def test_core(**kwargs) -> None: - """Run tests to check installation""" - passed = 0 - - output_path = Path.cwd() / "test_log" - output_path.mkdir(parents=True, exist_ok=True) - fsize = 2048 - stride = 2048 - batch = 64 - - fnames = ["test_short.fasta", "test_empty.fasta", "test_contigs.fasta"] - log_file = "test_jaeger.log" - logger = get_logger( - log_path=Path(output_path), log_file=log_file, level=kwargs.get("verbose") - ) - - # Test 1-3 - for i, f in enumerate(fnames): - input_file = str(files("jaeger.data.test").joinpath(f)) - logger.info(input_file) - try: - num = validate_fasta_entries(input_file) - passed += 1 - logger.info(f"{i} test passed {f}!") - except Exception as e: - if i > 1: - logger.error(f"{i} test failed {f}!") - else: - passed += 1 - logger.info(f"{i} test passed {f}!") - logger.debug(e) - - # Test 4 - result = test_tf() - if isinstance(result, Exception): - logger.error("4 tensorflow test failed!") - logger.debug(result) - else: - passed += 1 - logger.info("4 tensorflow test passed!") - - # Test 5 - try: - tf.config.set_soft_device_placement(True) - config_path = files("jaeger.data").joinpath("config.json") - config = json.loads(config_path.read_text()) - weights_path = files("jaeger.data.models.default").joinpath( - config["default"]["weights"] - ) - - input_dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - input_file, - fragsize=fsize, - stride=stride, - num=num, - ), - output_signature=(tf.TensorSpec(shape=(), dtype=tf.string)), - ) - - logger.info("loading the dataset") - idataset = input_dataset.map( - process_string(crop_size=fsize), - ).batch(batch) - - inputs, outputs = WRes_model_embeddings( - input_shape=(None,), dropout_active=False - ) - - logger.info("creating the model") - model = JaegerModel(inputs=inputs, outputs=outputs) - model.load_weights(filepath=weights_path) - model.summary() - logger.info(files("jaeger.data.models.test").joinpath("jaeger_fragment_graph")) - tf.saved_model.save( - model, - files("jaeger.data.models.test").joinpath("jaeger_fragment_graph"), - ) - - logger.info("loading the model") - model = InferModel( - files("jaeger.data.models.test").joinpath("jaeger_fragment_graph") - ) - logger.info("starting model inference") - _ = model.predict(idataset) - logger.info("5 test model passed!") - passed += 1 - except Exception as e: - logger.exception("5 test model failed!") - logger.debug(e) - finally: - logger.info(f"{passed}/5 tests passed!") diff --git a/src/jaeger/commands/utils.py b/src/jaeger/commands/utils.py index f5060f0..51e2a74 100644 --- a/src/jaeger/commands/utils.py +++ b/src/jaeger/commands/utils.py @@ -561,9 +561,8 @@ def optimize_data_core( ): """Convert Jaeger CSV training data to an optimized format. - Supports four output formats: - - tfrecord: Preprocessed tensors in TFRecord format - - numpy_raw: int8 sequences + TF preprocessing at train time + Supports three output formats: + - numpy_raw: int8 sequences + runtime preprocessing at train time - numpy_full: Fully preprocessed, fastest loading - numpy_raw_variable: Variable-length int8 sequences @@ -574,7 +573,7 @@ def optimize_data_core( output_path : str Path to output file. format : str - One of: tfrecord, numpy_raw, numpy_full, numpy_raw_variable. + One of: numpy_raw, numpy_full, numpy_raw_variable. crop_size : int Sequence crop size (default: 500). num_classes : int @@ -583,13 +582,12 @@ def optimize_data_core( Number of parallel workers for CPU-bound formats. Defaults to all CPUs. use_embedding_layer : bool - For tfrecord: use embedding layer (int indices) vs one-hot. + Unused (kept for API compatibility). max_length : int For numpy_raw_variable: maximum sequence length. """ format = format.lower() valid_formats = [ - "tfrecord", "numpy_raw", "numpy_full", "numpy_raw_variable", diff --git a/src/jaeger/commands/utils_models.py b/src/jaeger/commands/utils_models.py deleted file mode 100644 index 4886e7f..0000000 --- a/src/jaeger/commands/utils_models.py +++ /dev/null @@ -1,166 +0,0 @@ -import sys -import tensorflow as tf -import logging -from pathlib import Path -import shutil -from jaeger.utils.misc import AvailableModels - -logger = logging.getLogger("Jaeger") - - -class EnsembleModel(tf.Module): - def __init__(self, model_paths, method="mean"): - super().__init__() - self.method = method - self.models = [tf.saved_model.load(p) for p in model_paths] - self._signatures = [m.signatures["serving_default"] for m in self.models] - - @tf.function - def __call__(self, inputs): - """ - Combine outputs from all sub-models. - - Supported methods: - - mean: average logits/other tensors. - - sum: sum logits/other tensors. - - mv: majority vote over ``prediction`` logits; other tensors are - averaged over models that voted for the winning class. - - none: fallback to mean (keeps a valid SavedModel output). - """ - outputs = [sig(inputs) for sig in self._signatures] - - # Use only keys common to all model outputs. - common_keys = set(outputs[0].keys()) - for out in outputs[1:]: - common_keys &= set(out.keys()) - - if self.method in ("mean", "sum", "none"): - combined = {} - for key in common_keys: - stacked = tf.stack([out[key] for out in outputs], axis=0) - if self.method == "sum": - combined[key] = tf.reduce_sum(stacked, axis=0) - else: - combined[key] = tf.reduce_mean(stacked, axis=0) - return combined - - # Majority vote over the prediction head. - if "prediction" not in common_keys: - raise ValueError("Majority voting requires a 'prediction' output key") - - preds = tf.stack([out["prediction"] for out in outputs], axis=0) - argmaxes = tf.argmax(preds, axis=-1, output_type=tf.int32) - n_classes = tf.shape(preds)[-1] - - votes = tf.one_hot(argmaxes, depth=n_classes) - vote_counts = tf.reduce_sum(votes, axis=0) - majority_class = tf.argmax(vote_counts, axis=-1, output_type=tf.int32) - - majority_mask = tf.one_hot(majority_class, depth=n_classes) - majority_mask = tf.expand_dims(majority_mask, axis=0) - masked = preds * tf.cast(majority_mask, preds.dtype) - - sums = tf.reduce_sum(masked, axis=0) - counts = tf.reduce_sum(tf.cast(tf.not_equal(masked, 0.0), preds.dtype), axis=0) - majority_means = sums / tf.maximum(counts, 1.0) - - combined = {} - for key in common_keys: - stacked = tf.stack([out[key] for out in outputs], axis=0) - if key == "prediction": - combined[key] = majority_means - else: - combined[key] = tf.reduce_mean(stacked, axis=0) - return combined - - -def _resolve_model(path): - """Resolve a user-supplied path to a model dict with graph/project/classes.""" - p = Path(path) - - # Direct path to a SavedModel graph directory. - if p.is_dir() and p.name.endswith("_graph") and (p / "saved_model.pb").exists(): - graph_dir = p - base_name = graph_dir.name.removesuffix("_graph") - project = None - classes = None - for parent in [graph_dir.parent] + list(graph_dir.parents): - if project is None: - candidates = list(parent.glob(f"{base_name}_project.yaml")) - if candidates: - project = candidates[0] - if classes is None: - candidates = list(parent.glob(f"{base_name}_classes.yaml")) - if candidates: - classes = candidates[0] - if project and classes: - break - return {"graph": graph_dir, "project": project, "classes": classes} - - # Otherwise expect a Jaeger experiment/model directory containing a - # ``model`` sub-directory scanned by AvailableModels. - info = AvailableModels(p).info - if not info: - return None - # ``info`` is a dict mapping model name -> dict of artifacts. - return next(iter(info.values())) - - -def combine_models_core(**kwargs): - inputs = kwargs.get("input") - method = kwargs.get("comb", "mean") - - models = [] - for raw_path in inputs: - model = _resolve_model(raw_path) - if model is None: - logger.error(f"No Jaeger model found at: {raw_path}") - logger.error( - "Provide either the experiment directory that contains a " - "'model' sub-directory, or the direct path to a *_graph " - "SavedModel directory." - ) - sys.exit(1) - missing = [k for k in ("graph", "project", "classes") if not model.get(k)] - if missing: - logger.error( - f"Model at {raw_path} is missing required artifacts: {', '.join(missing)}" - ) - sys.exit(1) - models.append(model) - - out_dir = Path(kwargs.get("output")) - out_dir.mkdir(parents=True, exist_ok=True) - out_model_path = out_dir / "model" - out_model_path.mkdir(parents=True, exist_ok=True) - - graph_dir = models[0]["graph"] - model_name = graph_dir.name.removesuffix("_graph") - out_graph_path = out_model_path / f"{model_name}_{len(models)}_ensemble_graph" - out_project_path = ( - out_model_path / f"{model_name}_{len(models)}_ensemble_project.yaml" - ) - out_class_path = ( - out_model_path / f"{model_name}_{len(models)}_ensemble_classes.yaml" - ) - - model_paths = [m["graph"] for m in models] - project_path = models[0]["project"] - class_path = models[0]["classes"] - - ensemble = EnsembleModel(model_paths, method=method) - - # Get the input signature dynamically from the first sub-model. - input_spec = list(ensemble._signatures[0].structured_input_signature[1].values())[0] - - tf.saved_model.save( - ensemble, - export_dir=out_graph_path, - signatures={ - "serving_default": ensemble.__call__.get_concrete_function(input_spec) - }, - ) - shutil.copy(project_path, out_project_path) - shutil.copy(class_path, out_class_path) - - print(f"✅ Ensemble SavedModel successfully created at: {out_graph_path}") diff --git a/src/jaeger/data/__init__.py b/src/jaeger/data/__init__.py index 230f09f..d8dd661 100644 --- a/src/jaeger/data/__init__.py +++ b/src/jaeger/data/__init__.py @@ -37,16 +37,6 @@ # FASTA fragment generation from jaeger.seqops.io import fragment_generator, fragment_generator_lib -# Preprocessing functions -from jaeger.seqops.encode import process_string_train, process_string_inference - -# Dataset loaders -from jaeger.data.loaders import ( - _load_numpy_full_dataset as load_numpy_full, - _load_numpy_raw_dataset as load_numpy_raw, - _load_numpy_raw_variable_dataset as load_numpy_raw_variable, -) - # Public converter API from jaeger.dataops.convert import convert_dataset @@ -75,13 +65,6 @@ # FASTA "fragment_generator", "fragment_generator_lib", - # Preprocessing - "process_string_train", - "process_string_inference", - # Loaders - "load_numpy_full", - "load_numpy_raw", - "load_numpy_raw_variable", # Converters "convert_dataset", ] diff --git a/src/jaeger/data/loaders.py b/src/jaeger/data/loaders.py deleted file mode 100644 index 2b14555..0000000 --- a/src/jaeger/data/loaders.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Dataset loaders for NumPy-based training formats. - -Provides `tf.data.Dataset` builders for: -- ``numpy_raw`` — int8 DNA sequences processed on-the-fly to 6-frame codon IDs. -- ``numpy_raw_variable`` — variable-length int8 sequences. -- ``numpy_full`` — fully preprocessed arrays (fastest, no runtime augmentations). -""" - -from __future__ import annotations - -import numpy as np -import tensorflow as tf - -from jaeger.seqops.encode import ( - _make_process_raw_sequence_fn, - _make_process_variable_sequence_fn, -) - - -def _load_numpy_raw_dataset( - path: str, - crop_size: int = 500, - ngram_width: int = 3, - num_classes: int = 3, - shuffle: bool = False, - mutate: bool = False, - mutation_rate: float = 0.1, - shuffle_frames: bool = False, -): - """Loads a raw NumPy ``.npz`` file (int8 sequences) and returns a ``tf.data.Dataset``. - - The ``.npz`` must contain: - - ``sequences``: int8 array of shape ``(N, crop_size)`` - - ``labels``: float32 array of shape ``(N, num_classes)`` - - Each sequence is converted to 6-frame codon IDs via - :func:`jaeger.data.raw_processors._make_process_raw_sequence_fn`. - """ - data = np.load(path, allow_pickle=False) - sequences = data["sequences"] - labels = data["labels"] - - dataset = tf.data.Dataset.from_tensor_slices((sequences, labels)) - - process_fn = _make_process_raw_sequence_fn( - crop_size=crop_size, - ngram_width=ngram_width, - num_classes=num_classes, - shuffle=shuffle, - mutate=mutate, - mutation_rate=mutation_rate, - shuffle_frames=shuffle_frames, - ) - - dataset = dataset.map(process_fn, num_parallel_calls=tf.data.AUTOTUNE) - return dataset - - -def _load_numpy_raw_variable_dataset( - path: str, - ngram_width: int = 3, - num_classes: int = 3, - shuffle: bool = False, - mutate: bool = False, - mutation_rate: float = 0.1, - shuffle_frames: bool = False, -): - """Loads a variable-length raw NumPy ``.npz`` file and returns a ``tf.data.Dataset``. - - The ``.npz`` must contain: - - ``sequences``: int8 array of shape ``(N, max_length)`` (padded) - - ``lengths``: int array of shape ``(N,)`` with actual lengths - - ``labels``: float32 array of shape ``(N, num_classes)`` - """ - data = np.load(path, allow_pickle=False) - sequences = data["sequences"] - lengths = data["lengths"] - labels = data["labels"] - - dataset = tf.data.Dataset.from_tensor_slices((sequences, lengths, labels)) - - process_fn = _make_process_variable_sequence_fn( - ngram_width=ngram_width, - num_classes=num_classes, - shuffle=shuffle, - mutate=mutate, - mutation_rate=mutation_rate, - shuffle_frames=shuffle_frames, - ) - - dataset = dataset.map(process_fn, num_parallel_calls=tf.data.AUTOTUNE) - return dataset - - -def _load_numpy_full_dataset( - path: str, - input_type: str = "translated", - use_embedding_layer: bool = True, - codon_depth: int = 21, - crop_size: int = 500, -): - """Loads a fully-preprocessed NumPy ``.npz`` file and returns a ``tf.data.Dataset``. - - Supports both structured arrays ``(N, 6, seq_len)`` and flattened arrays - ``(N, flat)`` for backward compatibility. Also supports nucleotide input. - - The ``.npz`` file must contain: - - ``translated`` or ``nucleotide``: preprocessed sequences - - ``label``: float32 array of shape ``(N, num_classes)`` - - No further preprocessing is applied — the data is ready for direct training. - This gives the fastest possible data loading at the cost of no runtime - augmentations (shuffle, mutate, frame_shuffle). - """ - data = np.load(path, allow_pickle=False) - - if input_type == "translated": - if use_embedding_layer: - seq_shape = [6, crop_size // 3 - 1] - else: - seq_shape = [6, crop_size // 3 - 1, codon_depth] - seq_key = "translated" - elif input_type == "nucleotide": - seq_shape = [2, crop_size, 4] - seq_key = "nucleotide" - else: - raise ValueError(f"Unsupported input_type: {input_type}") - - seqs = data[seq_key] - labels = data["label"] - - # Ensure correct shapes if needed (flattened -> structured) - expected_flat = np.prod(seq_shape) - if seqs.ndim == 2 and seqs.shape[1] == expected_flat: - seqs = seqs.reshape([-1] + list(seq_shape)) - - dataset = tf.data.Dataset.from_tensor_slices(({seq_key: seqs}, labels)) - return dataset diff --git a/src/jaeger/data/tfrecord.py b/src/jaeger/data/tfrecord.py deleted file mode 100644 index 067df65..0000000 --- a/src/jaeger/data/tfrecord.py +++ /dev/null @@ -1,82 +0,0 @@ -"""TFRecord parsing utilities for preprocessed training data. - -Provides feature-description builders and parsing functions for -TensorFlow TFRecord datasets created by the conversion pipeline. -""" - -from __future__ import annotations - -import tensorflow as tf - - -def _get_tfrecord_feature_description( - input_type: str, - use_embedding_layer: bool, - codon_depth: int, - crop_size: int, - num_classes: int, -) -> dict: - """Returns TFRecord feature description for parsing preprocessed data.""" - if input_type == "translated": - if use_embedding_layer: - feature_description = { - "translated": tf.io.FixedLenFeature( - [6 * (crop_size // 3 - 1)], tf.int64 - ), - "label": tf.io.FixedLenFeature([num_classes], tf.float32), - } - else: - feature_description = { - "translated": tf.io.FixedLenFeature( - [6 * (crop_size // 3 - 1) * codon_depth], tf.float32 - ), - "label": tf.io.FixedLenFeature([num_classes], tf.float32), - } - elif input_type == "nucleotide": - feature_description = { - "nucleotide": tf.io.FixedLenFeature([2 * crop_size * 4], tf.float32), - "label": tf.io.FixedLenFeature([num_classes], tf.float32), - } - else: - raise ValueError(f"Unsupported input_type: {input_type}") - return feature_description - - -def _make_parse_tfrecord_fn( - input_type: str, - use_embedding_layer: bool, - codon_depth: int, - crop_size: int, - num_classes: int, -): - """Creates a TFRecord parsing function for the given config.""" - feature_description = _get_tfrecord_feature_description( - input_type, use_embedding_layer, codon_depth, crop_size, num_classes - ) - - if input_type == "translated": - if use_embedding_layer: - seq_shape = [6, crop_size // 3 - 1] - else: - seq_shape = [6, crop_size // 3 - 1, codon_depth] - elif input_type == "nucleotide": - seq_shape = [2, crop_size, 4] - - @tf.function - def _parse_tfrecord(example_proto): - parsed = tf.io.parse_single_example(example_proto, feature_description) - if input_type == "translated": - if use_embedding_layer: - seq = tf.cast(parsed["translated"], tf.int32) - else: - seq = tf.cast(parsed["translated"], tf.float32) - seq = tf.reshape(seq, seq_shape) - features = {"translated": seq} - elif input_type == "nucleotide": - seq = tf.cast(parsed["nucleotide"], tf.float32) - seq = tf.reshape(seq, seq_shape) - features = {"nucleotide": seq} - label = parsed["label"] - return features, label - - return _parse_tfrecord diff --git a/src/jaeger/dataops/convert.py b/src/jaeger/dataops/convert.py index 111219d..ec1a8f6 100644 --- a/src/jaeger/dataops/convert.py +++ b/src/jaeger/dataops/convert.py @@ -1,7 +1,6 @@ """Dataset format converters. Converts CSV training data to optimized formats: -- ``tfrecord`` — TensorFlow TFRecord with preprocessed tensors. - ``numpy_raw`` — int8 DNA sequences (fast loading, runtime preprocessing). - ``numpy_full`` — fully preprocessed tensors (fastest loading, no augmentations). - ``numpy_raw_variable`` — variable-length int8 sequences. @@ -14,8 +13,6 @@ import time from functools import partial from multiprocessing import Pool, cpu_count -from pathlib import Path - import numpy as np # --------------------------------------------------------------------------- @@ -37,117 +34,6 @@ def wrapper(func): return wrapper -# --------------------------------------------------------------------------- -# TFRecord helpers -# --------------------------------------------------------------------------- -def _int64_feature(value): - """Returns an int64_list from a bool / enum / int / uint.""" - import tensorflow as tf - - return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) - - -def _float_feature(value): - """Returns a float_list from a float / double.""" - import tensorflow as tf - - return tf.train.Feature(float_list=tf.train.FloatList(value=value)) - - -def _serialize_tfrecord_embedding(translated, label): - """Serialize a single example for embedding layer (int32 indices).""" - import tensorflow as tf - - translated_flat = tf.reshape(translated, [-1]) - feature = { - "translated": _int64_feature(translated_flat.numpy().tolist()), - "label": _float_feature(label.numpy().flatten().tolist()), - } - return tf.train.Example( - features=tf.train.Features(feature=feature) - ).SerializeToString() - - -def _serialize_tfrecord_onehot(translated, label): - """Serialize a single example for one-hot encoded input (float32).""" - import tensorflow as tf - - translated_flat = tf.reshape(translated, [-1]) - feature = { - "translated": _float_feature(translated_flat.numpy().flatten().tolist()), - "label": _float_feature(label.numpy().flatten().tolist()), - } - return tf.train.Example( - features=tf.train.Features(feature=feature) - ).SerializeToString() - - -def _convert_to_tfrecord( - csv_path, output_path, total, preprocess_fn, use_embedding_layer -): - """Convert CSV to TFRecord with preprocessed tensors.""" - import tensorflow as tf - - serialize_fn = ( - _serialize_tfrecord_embedding - if use_embedding_layer - else _serialize_tfrecord_onehot - ) - - Path(output_path).parent.mkdir(parents=True, exist_ok=True) - - with tf.io.TFRecordWriter(output_path) as writer: - start = time.time() - with open(csv_path) as f: - for i, line in enumerate(f): - outputs, label = preprocess_fn(line.strip().encode()) - translated = outputs["translated"] - example = serialize_fn(translated, label) - writer.write(example) - if (i + 1) % 5000 == 0: - elapsed = time.time() - start - rate = (i + 1) / elapsed - print( - f" {i + 1}/{total} ({100 * (i + 1) / total:.1f}%) - {rate:.1f} samples/sec" - ) - elapsed = time.time() - start - print( - f"Done! {total} samples in {elapsed:.1f}s ({total / elapsed:.1f} samples/sec)" - ) - - -def _convert_with_tf( - csv_path: str, - output_path: str, - format: str, - crop_size: int, - num_classes: int, - use_embedding_layer: bool, -): - """Convert CSV using TensorFlow preprocessing (tfrecord only).""" - from jaeger.seqops.encode import process_string_train - from jaeger.seqops.maps import CODONS, CODON_ID - - total = sum(1 for _ in open(csv_path)) - print(f"Total samples: {total}") - - preprocess_fn = process_string_train( - crop_size=crop_size, - seq_onehot=False, - input_type="translated", - class_label_onehot=True, - num_classes=num_classes, - shuffle=False, - ngram_width=3, - codons=CODONS, - codon_num=CODON_ID, - ) - - _convert_to_tfrecord( - csv_path, output_path, total, preprocess_fn, use_embedding_layer - ) - - # --------------------------------------------------------------------------- # numpy_raw (int8 sequences) # --------------------------------------------------------------------------- @@ -727,15 +613,15 @@ def convert_dataset( Args: input_path: Path to input CSV file (label,sequence format). output_path: Path to output file. - format: Target format — ``tfrecord``, ``numpy_raw``, ``numpy_full``, + format: Target format — ``numpy_raw``, ``numpy_full``, or ``numpy_raw_variable``. crop_size: Sequence crop size (for fixed-length formats). num_classes: Number of output classes. - use_embedding_layer: Whether model uses embedding layer (TFRecord only). + use_embedding_layer: Unused (kept for API compatibility). max_length: Maximum sequence length (variable-length only). num_workers: Number of parallel workers (``None`` = auto). """ - valid_formats = ["tfrecord", "numpy_raw", "numpy_full", "numpy_raw_variable"] + valid_formats = ["numpy_raw", "numpy_full", "numpy_raw_variable"] if format not in valid_formats: raise ValueError( f"Invalid format: {format}. Choose from: {', '.join(valid_formats)}" @@ -744,11 +630,7 @@ def convert_dataset( print(f"Converting {input_path} -> {output_path}") print(f"Format: {format}, Crop size: {crop_size}, Num classes: {num_classes}") - if format == "tfrecord": - _convert_with_tf( - input_path, output_path, format, crop_size, num_classes, use_embedding_layer - ) - elif format == "numpy_raw": + if format == "numpy_raw": _convert_to_numpy_raw( input_path, output_path, crop_size, num_classes, num_workers ) diff --git a/src/jaeger/nnlib/builder.py b/src/jaeger/nnlib/builder.py deleted file mode 100644 index eb4b849..0000000 --- a/src/jaeger/nnlib/builder.py +++ /dev/null @@ -1,969 +0,0 @@ -"""Model builders for Jaeger neural networks. - -This module contains the `DynamicModelBuilder` class which constructs -Keras models from YAML/JSON configuration dictionaries. It is used by -the training pipeline and can be imported independently of the CLI layer. -""" - -from __future__ import annotations - -import json -import math -import os -import random -import re -import shutil -from pathlib import Path -from typing import Any, Optional -from uuid import uuid4 - -import numpy as np -import tensorflow as tf -import yaml - -from jaeger.seqops.maps import ( - AA_ID, - CODONS, - CODON_ID, - DICODONS, - DICODON_ID, - MURPHY10_ID, - PC5_ID, -) -from jaeger.nnlib.metrics import ( - PrecisionForClass, - RecallForClass, - SpecificityForClass, -) -from jaeger.nnlib.v2.layers import ( - AxialAttention, - CrossFrameAttention, - GatedFrameGlobalMaxPooling, - MaskedBatchNorm, - MaskedConv1D, - MetricModel, - ResidualBlock_wrapper, - TransformerEncoder, -) -from jaeger.nnlib.v2.losses import ArcFaceLoss, HierarchicalLoss -from jaeger.utils.logging import get_logger -from jaeger.utils.misc import clear_directory - -logger = get_logger(log_file=None, log_path=None, level=3) - - -def set_global_seed(seed: int = 42) -> None: - """Set deterministic seeds for Python, NumPy and TensorFlow.""" - random.seed(seed) - np.random.seed(seed) - tf.random.set_seed(seed) - os.environ["TF_DETERMINISTIC_OPS"] = "1" - os.environ["PYTHONHASHSEED"] = str(seed) - - -def check_files(files: list[str]) -> list[str]: - """Return subset of *files* that exist on disk.""" - return [f for f in files if Path(f).is_file()] - - -GRAPH_RE = re.compile( - r"^jaeger_(?P[0-9a-fA-F]+)_(?P\d+(?:\.\d+)?[A-Z])_fragment_graph$" -) - - -def find_existing_graph_id(path: Path) -> Optional[str]: - """Search *direct children* of *path* for jaeger__1.2M_fragment_graph. - - Returns the extracted or ``None`` if not found. - """ - if not path.exists() or not path.is_dir(): - return None - for p in path.iterdir(): - m = GRAPH_RE.match(p.name) - if m: - return m.group("id") - return None - - -class DynamicModelBuilder: - """Builds Keras models from a configuration dictionary. - - The builder supports: - - Embedding layers (translated codons or nucleotide one-hot) - - Representation learners (residual blocks, transformers, attention) - - Projection heads (for self-supervised pre-training with ArcFace loss) - - Classification heads - - Reliability heads (binary/multi-class confidence estimation) - - Combined ``jaeger_model`` that exposes all outputs - - Parameters - ---------- - config: - Full training configuration dict. Must contain ``model`` and - ``training`` top-level keys. - """ - - def __init__(self, config: dict[str, Any]) -> None: - self.uid: str = uuid4().hex[:8] - self.cfg: dict[str, Any] = config - self.model_cfg: dict[str, Any] = config.get("model", {}) or {} - self.train_cfg: dict[str, Any] = config.get("training", {}) or {} - - self.inputs: Optional[tf.Tensor] = None - self.outputs: list[tf.Tensor] = [] - - embedding_cfg = self.model_cfg.get("embedding", {}) or {} - self.input_shape = embedding_cfg.get("input_shape") - - self._fragment_paths = self._get_fragment_paths() - self._contig_paths = self._get_contig_paths() - self._reliability_fragment_paths = self._get_reliability_fragment_paths() - self._reliability_contig_paths = self._get_reliability_contig_paths() - - self._saving_config = self._get_model_saving_configuration() - self._from_last_checkpoint: bool = bool( - config.get("from_last_checkpoint", False) - ) - self._force: bool = bool(config.get("force", False)) - self.use_xla: bool = bool(config.get("use_xla", False)) - self._checkpoints: dict[str, Path] = {} - - self.classifier_out_dim: int = int(self.model_cfg.get("classifier_out_dim", 0)) - self.reliability_out_dim: int = int( - self.model_cfg.get("reliability_out_dim", 0) - ) - - self.loss_classifier_name = self.train_cfg.get( - "loss_classifier", "categorical_crossentropy" - ).lower() - self.loss_reliability_name = self.train_cfg.get( - "loss_reliability", "binary_crossentropy" - ).lower() - self.loss_classifier = None - self.loss_reliability = None - self.metrics_classifier: list[Any] = [] - self.metrics_reliability: list[Any] = [] - - self._regularizer = { - "l2": tf.keras.regularizers.L2, - "l1": tf.keras.regularizers.L1, - } - - self._layers = { - "masked_conv1d": MaskedConv1D, - "conv1d": tf.keras.layers.Conv1D, - "masked_batchnorm": MaskedBatchNorm, - "batchnorm": tf.keras.layers.BatchNormalization, - "transformer_encoder": TransformerEncoder, - "cross_frame_attention": CrossFrameAttention, - "axial_attention": AxialAttention, - "residual_block": ResidualBlock_wrapper, - "dense": tf.keras.layers.Dense, - "activation": tf.keras.layers.Activation, - "dropout": tf.keras.layers.Dropout, - "crop": tf.keras.layers.Cropping2D, - } - - self._load_training_params() - self._prepare_checkpoint_dirs() - - # ------------------------------------------------------------------ - # Checkpoint helpers - # ------------------------------------------------------------------ - - def _prepare_checkpoint_dirs(self, config: dict | None = None) -> None: - if not config: - config = self.cfg - callbacks_cfg = config.get("training", {}).get("callbacks", {}) - directories = callbacks_cfg.get("directories", []) or [] - - should_exit = False - for dir_str in directories: - path = Path(dir_str) - if path.exists(): - if self._from_last_checkpoint: - self._checkpoints[path.name] = self.get_latest_h5_with_metadata( - path - ) - elif self._force: - shutil.rmtree(path) - else: - logger.warning( - "Checkpoint(s) exist at %s. " - "Use --force to delete the existing checkpoints and continue!", - path, - ) - logger.info( - "Or set --from_last_checkpoint to continue training " - "from the last checkpoint." - ) - should_exit = True - continue - path.mkdir(parents=True, exist_ok=True) - if should_exit: - exit(1) - - def get_latest_h5_with_metadata( - self, - path: str | Path, - check_convergence: str = "classifier", - pattern: str = r"epoch:(\d+)-loss:(\d+\.\d+)", - ) -> dict: - """Scan *path* for the most recent ``.h5`` checkpoint matching *pattern*.""" - path = Path(path) - h5_files = sorted( - path.glob("*.h5"), key=lambda f: f.stat().st_mtime, reverse=True - ) - for file in h5_files: - match = re.search(pattern, file.name) - if match: - epoch, loss = match.groups() - return { - "path": file, - "epoch": int(epoch), - "loss": float(loss), - "is_converged": False if path.name == check_convergence else False, - } - return {"path": None, "epoch": 0, "loss": None, "is_converged": False} - - # ------------------------------------------------------------------ - # Model construction - # ------------------------------------------------------------------ - - def build_fragment_classifier(self) -> dict[str, tf.keras.Model]: - """Build the full fragment-level model graph. - - Returns a dict with keys such as ``rep_model``, ``classification_head``, - ``reliability_head``, ``jaeger_classifier``, ``jaeger_reliability``, - ``jaeger_model``, etc. - """ - models: dict[str, tf.keras.Model] = {} - - # === 1. EMBEDDING === - if "embedding" in self.model_cfg: - inputs, x = self._build_embedding(self._get_string_processor_config()) - self.inputs = inputs - else: - raise ValueError("Missing 'embedding' section in config") - - # === 2. REPRESENTATION LEARNER === - if "representation_learner" in self.model_cfg: - rep_out = self._build_block( - x, self.model_cfg["representation_learner"], prefix="rep" - ) - models["rep_model"] = tf.keras.Model( - inputs=self.inputs, outputs=rep_out, name="rep_model" - ) - models["rep_model"].summary() - - # === 3. PRETRAINING (projection head) === - if "projection" in self.model_cfg: - input_shape = (self.model_cfg["projection"].get("input_shape"),) - inputs = tf.keras.Input(shape=input_shape, name="projection_input") - x_projection = self._build_block( - inputs, self.model_cfg["projection"], prefix="projection" - ) - models["projection_head"] = tf.keras.Model( - inputs=inputs, outputs=x_projection, name="projection_head" - ) - projection_dim = self.model_cfg["projection"]["hidden_layers"][-2][ - "config" - ].get("units") - - x = models["rep_model"].output[0] - x = models["projection_head"](x) - models["jaeger_projection"] = MetricModel( - inputs=models["rep_model"].input, outputs=x, name="Jaeger_projection" - ) - labels = tf.keras.Input(shape=(self.classifier_out_dim,), name="labels") - embeddings = tf.keras.Input(shape=(projection_dim,), name="embedding") - loss = ArcFaceLoss( - num_classes=self.classifier_out_dim, - embedding_dim=projection_dim, - margin=self.model_cfg["projection"]["margin"], - scale=self.model_cfg["projection"]["scale"], - onehot=True, - )(labels, embeddings) - models["arcface_loss"] = tf.keras.Model( - inputs=[labels, embeddings], outputs=loss, name="Arcface" - ) - if self._checkpoints.get("projection", {}).get("path", False): - models["jaeger_projection"].load_weights( - self._checkpoints.get("projection").get("path"), - skip_mismatch=True, - ) - logger.info( - f"Loaded projection model weights from " - f"{self._checkpoints.get('projection').get('path')}" - ) - - # === 3. CLASSIFIER === - if "classifier" in self.model_cfg: - input_shape = (self.model_cfg["classifier"].get("input_shape"),) - inputs = tf.keras.Input(shape=input_shape, name="classifier_input") - x_classifier = self._build_block( - inputs, self.model_cfg["classifier"], prefix="classifier" - ) - models["classification_head"] = tf.keras.Model( - inputs=inputs, outputs=x_classifier, name="classification_head" - ) - rep_out = models["rep_model"].output - if isinstance(rep_out, tuple): - x_rep = rep_out[0] - else: - x_rep = rep_out - x = models["classification_head"](x_rep) - models["jaeger_classifier"] = tf.keras.Model( - inputs=models["rep_model"].input, outputs=x, name="Jaeger_classifier" - ) - if self._checkpoints.get("classifier", {}).get("path", False): - models["jaeger_classifier"].load_weights( - self._checkpoints.get("classifier").get("path"), - skip_mismatch=True, - ) - logger.info( - f"Loaded classification model weights from " - f"{self._checkpoints.get('classifier').get('path')}" - ) - - # === 4. RELIABILITY === - if "reliability_model" in self.model_cfg: - input_shape = (self.model_cfg["reliability_model"].get("input_shape"),) - inputs = tf.keras.Input(shape=input_shape, name="reliability_input") - x_reliability = self._build_block( - inputs, self.model_cfg["reliability_model"], prefix="reliability" - ) - models["reliability_head"] = tf.keras.Model( - inputs=inputs, outputs=x_reliability, name="reliability_head" - ) - x = models["rep_model"].output[1] - x = models["reliability_head"](x) - models["jaeger_reliability"] = tf.keras.Model( - inputs=models["rep_model"].input, outputs=x, name="Jaeger_reliability" - ) - try: - if self._checkpoints.get("reliability", {}).get("path", False): - models["jaeger_reliability"].load_weights( - self._checkpoints.get("reliability").get("path"), - skip_mismatch=True, - ) - logger.info( - f"Loaded reliability model weights from " - f"{self._checkpoints.get('reliability').get('path')}" - ) - except Exception: - logger.warning( - "could not load the weights to reliability model from checkpoint. " - "trying to load weights partially" - ) - models["jaeger_reliability"].load_weights( - self._checkpoints.get("reliability").get("path"), - skip_mismatch=True, - ) - self._checkpoints["reliability"] = { - "path": None, - "epoch": 0, - "loss": None, - "is_converged": False, - } - - # === 5. COMBINED MODEL === - rep_out = models["rep_model"].output - has_reliability = "reliability_head" in models - if isinstance(rep_out, tuple): - if len(rep_out) == 2: - x1, x2 = rep_out - class_ = models["classification_head"](x1) - outputs: dict[str, Any] = { - "prediction": class_, - "embedding": x1, - "nmd": x2, - } - if has_reliability: - outputs["reliability"] = models["reliability_head"](x2) - models["jaeger_model"] = tf.keras.Model( - inputs=models["rep_model"].input, - outputs=outputs, - name="Jaeger_model", - ) - elif len(rep_out) == 3: - x1, x2, g = rep_out - class_ = models["classification_head"](x1) - outputs = { - "prediction": class_, - "embedding": x1, - "nmd": x2, - "gate": g, - } - if has_reliability: - outputs["reliability"] = models["reliability_head"](x2) - models["jaeger_model"] = tf.keras.Model( - inputs=models["rep_model"].input, - outputs=outputs, - name="Jaeger_model", - ) - else: - class_ = models["classification_head"](rep_out) - models["jaeger_model"] = tf.keras.Model( - inputs=models["rep_model"].input, - outputs={"prediction": class_, "embedding": rep_out}, - name="Jaeger_model", - ) - - return models - - def build_contig_classifier(self) -> None: - """Placeholder for future contig-level consensus model.""" - pass - - def _build_embedding(self, cfg: dict[str, Any]) -> tuple[tf.Tensor, tf.Tensor]: - """Create the embedding layer from *cfg*.""" - input_shape = cfg.get("input_shape", (6, None, 64)) - embedding_size = cfg.get("embedding_size", 4) - - inputs = tf.keras.Input(shape=input_shape, name=cfg.get("input_type")) - masked_inputs = tf.keras.layers.Masking(name="input_mask", mask_value=0.0)( - inputs - ) - - match cfg.get("input_type"): - case "translated": - if embedding_size > 0: - if cfg.get("use_embedding_layer", False): - x = tf.keras.layers.Embedding( - input_dim=cfg.get("vocab_size"), - output_dim=embedding_size, - mask_zero=True, - embeddings_initializer=tf.keras.initializers.Orthogonal(), - embeddings_regularizer=self._regularizer.get( - cfg.get("embedding_regularizer"), - )(cfg.get("embedding_regularizer_w")), - )(inputs) - else: - x = tf.keras.layers.Dense( - embedding_size, - name=f"{cfg.get('input_type')}_embedding", - use_bias=False, - kernel_initializer=tf.keras.initializers.Orthogonal(), - kernel_regularizer=self._regularizer.get( - cfg.get("embedding_regularizer"), - )(cfg.get("embedding_regularizer_w")), - )(masked_inputs) - else: - x = masked_inputs - case "nucleotide": - x = masked_inputs - case _: - raise ValueError(f"{cfg.get('input_type')} is invalid") - - if cfg.get("use_positional_embeddings", False): - from jaeger.nnlib.v2.layers import SinusoidalPositionEmbedding - - positional = SinusoidalPositionEmbedding( - max_wavelength=cfg.get("positional_embedding_length") - )(x) - x = tf.keras.layers.Add()([x, positional]) - - return inputs, x - - def _get_bias(self, data_path: str, kind: str, label_map: list) -> np.ndarray: - """Compute class-frequency bias for the final layer.""" - import polars as pl - - def _sigmoid(f: dict): - n, p = f.values() - t = p / (p + n) - return np.log(t / (1 - t)) - - def _softmax(f: dict): - f = np.array(list(f.values())) - return np.log(f / np.sum(f)) - - def _correct_label_map(f: dict, label_map: list): - _tmp = {i: 0 for i in range(max(label_map) + 1)} - for k, v in f.items(): - _tmp[label_map[k]] += v - return _tmp - - df = pl.read_csv(data_path, columns=[0], has_header=False) - counts_dict = df["column_1"].value_counts().to_dict(as_series=False) - counts_dict = dict(zip(counts_dict["column_1"], counts_dict["count"])) - counts_dict = {k: counts_dict[k] for k in sorted(list(counts_dict.keys()))} - if len(label_map) > 0: - counts_dict = _correct_label_map(counts_dict, label_map) - match kind: - case "softmax": - return _softmax(counts_dict) - case "sigmoid": - return _sigmoid(counts_dict) - - def _build_block(self, x, cfg: dict[str, Any], prefix: str): - """Build a stack of layers from *cfg* and return the output tensor(s).""" - nmd = [] - nmd_con = tf.keras.layers.Concatenate(axis=-1) - previous_channels = None - for i, layer_cfg in enumerate(cfg.get("hidden_layers", [])): - layer_name = layer_cfg.get("name", "").lower() - cfg_layer = dict(layer_cfg.get("config", {})) - cfg_layer["name"] = f"{prefix}_{layer_name}_{i}" - - layer_class = self._layers.get(layer_name) - if layer_class is None: - raise ValueError(f"Unknown layer type: {layer_name}") - - if "kernel_regularizer" in cfg_layer: - reg_name = cfg_layer.pop("kernel_regularizer") - reg_w = cfg_layer.pop("kernel_regularizer_w", None) - cfg_layer["kernel_regularizer"] = self._regularizer[reg_name](reg_w) - - if "kernel_initializer" in cfg_layer: - init_name = cfg_layer.pop("kernel_initializer") - cfg_layer["kernel_initializer"] = tf.keras.initializers.get(init_name) - - if "bias_initializer" in cfg_layer: - if "calculate_from" in cfg_layer.get("bias_initializer"): - classifier_label_map = self.model_cfg.get( - "string_processor", {} - ).get("classifier_labels_map", []) - reliability_label_map = self.model_cfg.get( - "string_processor", {} - ).get("reliability_labels_map", []) - if "relia" in prefix: - path = ( - self._get_reliability_fragment_paths() - .get("train", {}) - .get("paths", [])[-1] - ) - cfg_layer["bias_initializer"] = tf.keras.initializers.Constant( - self._get_bias( - path, - kind="sigmoid" - if "binary" in self.loss_reliability_name - else "softmax", - label_map=reliability_label_map, - ) - ) - elif "classi" in prefix: - path = ( - self._get_fragment_paths() - .get("train", {}) - .get("paths", [])[-1] - ) - cfg_layer["bias_initializer"] = tf.keras.initializers.Constant( - self._get_bias( - path, - kind="sigmoid" - if "binary" in self.loss_classifier_name - else "softmax", - label_map=classifier_label_map, - ) - ) - - if "block_size" in cfg_layer: - block_size = cfg_layer.pop("block_size") - shape = (self.input_shape[0], None, previous_channels) - if block_size > 0: - if cfg_layer.get("return_nmd", False): - x, nmd_ = layer_class(block_size, shape, **cfg_layer)(x) - nmd.append(nmd_) - else: - x = layer_class(block_size, shape, **cfg_layer)(x) - if "filters" in cfg_layer: - previous_channels = cfg_layer.get("filters") - elif "units" in cfg_layer: - previous_channels = cfg_layer.get("units") - continue - - if "return_nmd" in cfg_layer: - if cfg_layer.get("return_nmd"): - x, nmd_ = layer_class(**cfg_layer)(x) - nmd.append(nmd_) - else: - x = layer_class(**cfg_layer)(x) - if "filters" in cfg_layer: - previous_channels = cfg_layer.get("filters") - elif "units" in cfg_layer: - previous_channels = cfg_layer.get("units") - continue - - x = layer_class(**cfg_layer)(x) - - if "filters" in cfg_layer: - previous_channels = cfg_layer.get("filters") - elif "units" in cfg_layer: - previous_channels = cfg_layer.get("units") - elif "embed_dim" in cfg_layer: - previous_channels = cfg_layer.get("embed_dim") - - # Aggregation - if "pooling" in cfg: - pooling = cfg.get("pooling", "average").lower() - pooler = self._get_pooler(pooling) - has_nmd = len(nmd) > 0 - if len(nmd) > 1: - nmd = nmd_con(nmd) - elif len(nmd) == 1: - nmd = nmd[0] - if "gated" not in pooling: - x = pooler(name=f"global_{pooling}pool")(x) - return (x, nmd) if has_nmd else x - else: - x, g = pooler(return_gate=True, name=f"global_{pooling}pool")(x) - return (x, nmd, g) if has_nmd else (x, g) - return x - - # ------------------------------------------------------------------ - # Metrics / loss / optimizer factories - # ------------------------------------------------------------------ - - def _get_metrics(self, config: list[dict]) -> list[Any]: - _metrics = { - "categorical_accuracy": tf.keras.metrics.CategoricalAccuracy, - "binary_accuracy": tf.keras.metrics.BinaryAccuracy, - "sparse_categorical_accuracy": tf.keras.metrics.SparseCategoricalAccuracy, - "categorical_crossentropy": tf.keras.metrics.CategoricalCrossentropy, - "binary_crossentropy": tf.keras.metrics.BinaryCrossentropy, - "sparse_categorical_crossentropy": tf.keras.metrics.SparseCategoricalCrossentropy, - "auc": tf.keras.metrics.AUC, - "precision": tf.keras.metrics.Precision, - "recall": tf.keras.metrics.Recall, - "per_class_precision": PrecisionForClass, - "per_class_recall": RecallForClass, - "per_class_specificity": SpecificityForClass, - } - metrics = [] - for c in config: - if "per_class_" not in c.get("name"): - if c.get("params") is not None: - metrics.append(_metrics.get(c.get("name"))(**c.get("params"))) - else: - metrics.append(_metrics.get(c.get("name"))()) - else: - for cls in range(self.classifier_out_dim): - metrics.append(_metrics.get(c.get("name"))(class_id=cls)) - return metrics - - def _load_training_params(self) -> None: - opt_name = self.train_cfg.get("optimizer", "adam").lower() - opt_params = self.train_cfg.get("optimizer_params", {}) - self.optimizer = self._get_optimizer(opt_name, opt_params) - loss_classifier_name = self.train_cfg.get( - "loss_classifier", "categorical_crossentropy" - ).lower() - loss_classifier_params = self.train_cfg.get("loss_params_classifier", {}) - self.loss_classifier = self._get_loss( - loss_classifier_name, loss_classifier_params - ) - - loss_reliability_name = self.train_cfg.get( - "loss_reliability", "binary_crossentropy" - ).lower() - loss_reliability_params = self.train_cfg.get("loss_params_reliability", {}) - self.loss_reliability = self._get_loss( - loss_reliability_name, loss_reliability_params - ) - - if len(self.train_cfg.get("metrics_classifier", [])) > 0: - self.metrics_classifier = self._get_metrics( - self.train_cfg.get("metrics_classifier") - ) - else: - self.metrics_classifier = self._get_default_metrics(branch="classifier") - - if len(self.train_cfg.get("metrics_reliability", [])) > 0: - self.metrics_reliability = self._get_metrics( - self.train_cfg.get("metrics_reliability") - ) - else: - self.metrics_reliability = self._get_default_metrics(branch="reliability") - - def _get_default_metrics(self, branch: str) -> list[Any]: - match branch: - case "classifier": - return [tf.keras.metrics.CategoricalAccuracy(name="acc")] - case "reliability": - return [ - tf.keras.metrics.AUC(name="auc", from_logits=True), - tf.keras.metrics.BinaryAccuracy(name="acc"), - ] - - def compile_model(self, model: dict, train_branch: str = "classifier") -> None: - """Compile a specific branch of the model graph.""" - opt_name = self.train_cfg.get("optimizer", "adam").lower() - opt_params = self.train_cfg.get("optimizer_params", {}) - self.optimizer = self._get_optimizer(opt_name, opt_params) - jit_compile = self.use_xla - if train_branch == "pretrain": - model.get("rep_model").trainable = True - model.get("jaeger_projection").compile( - optimizer=self.optimizer, - loss_fn=model.get("arcface_loss"), - run_eagerly=False, - jit_compile=jit_compile, - ) - logger.info(f"model compiled for {train_branch} (xla={jit_compile})") - elif train_branch == "classifier": - model.get("rep_model").trainable = True - model.get("jaeger_classifier").compile( - optimizer=self.optimizer, - loss=self.loss_classifier, - metrics=self.metrics_classifier, - jit_compile=jit_compile, - ) - logger.info(f"model compiled for {train_branch} (xla={jit_compile})") - elif train_branch == "reliability": - if model.get("jaeger_reliability") is None: - logger.warning( - "jaeger_reliability not built — skipping reliability compilation" - ) - return - model.get("rep_model").trainable = False - model.get("jaeger_reliability").compile( - optimizer=self.optimizer, - loss=self.loss_reliability, - metrics=self.metrics_reliability, - jit_compile=jit_compile, - ) - else: - raise ValueError( - "train_branch must be 'pretrain', 'classifier' or 'reliability'" - ) - - # ------------------------------------------------------------------ - # Saving / callbacks - # ------------------------------------------------------------------ - - def _prepare_save_path( - self, - num_params: str | None = None, - suffix: str | None = None, - metadata: str | None = None, - clear: bool = True, - ) -> tuple[Path, str]: - """Create the save directory and build the model name. - - Returns the target path and the computed model name. - """ - path = Path(self._saving_config.get("path")) - path.mkdir(parents=True, exist_ok=True) - - if metadata: - metadata = Path(metadata).resolve() - metadata.write_text( - json.dumps( - {"model_path": str(path), "experiment_path": str(path.parent)}, - indent=2, - ) - ) - graph_id = find_existing_graph_id(path) - - if clear and any(path.iterdir()): - logger.warning(f"{path} is not empty. deleting existing files!") - if graph_id is not None: - logger.warning("if a graph is found, it's id will be reused") - logger.warning(f"found existing graph id: {graph_id}") - clear_directory(path=path) - - model_name = self.model_cfg.get("name") - model_name += ( - f"{'_' + graph_id if graph_id else '_' + self.uid}" - f"{'_' + num_params if num_params else ''}" - f"{'_' + suffix if suffix else ''}" - ) - return path, model_name - - def save_model( - self, - model: tf.keras.Model, - num_params: str | None = None, - suffix: str | None = None, - metadata: str | None = None, - ) -> None: - """Save model weights, SavedModel graph, class map and project config.""" - path, model_name = self._prepare_save_path( - num_params=num_params, suffix=suffix, metadata=metadata, clear=True - ) - - if self._saving_config.get("save_weights"): - model.save_weights(path / f"{model_name}.weights.h5") - logger.info( - f"model weights are written to {path / f'{model_name}.weights.h5'}" - ) - - if self._saving_config.get("save_exec_graph"): - tf.saved_model.save(model, path / f"{model_name}_graph") - logger.info( - f"model computational graph is written to " - f"{path / f'{model_name}_graph'}" - ) - - with open(path / f"{model_name}_classes.yaml", "w") as yaml_file: - logger.info("writing class_labels_map") - yaml.dump( - dict(classes=self.model_cfg.get("class_label_map")), - yaml_file, - default_flow_style=False, - ) - - shutil.copy(self.cfg.get("config_path"), path / f"{model_name}_project.yaml") - - def save_embedding_model( - self, - models: dict[str, tf.keras.Model], - num_params: str | None = None, - suffix: str | None = None, - metadata: str | None = None, - ) -> None: - """Save a dedicated embedding-only SavedModel graph. - - The embedding graph returns the representation vector just before the - classifier head. This is used by downstream tasks such as the taxonomy - workflow, and is kept separate so it can be loaded without the larger - classification / reliability heads. - """ - if not self._saving_config.get("save_embedding_graph", True): - return - - rep_model = models.get("rep_model") - if rep_model is None: - logger.warning( - "No 'rep_model' found; skipping embedding-only graph export." - ) - return - - path, model_name = self._prepare_save_path( - num_params=num_params, suffix=suffix, metadata=metadata, clear=False - ) - - rep_output = rep_model.output - if isinstance(rep_output, (list, tuple)): - rep_output = rep_output[0] - - # Wrap the representation in an Identity layer and export it under the - # key "embedding" so the SavedModel signature is easy to consume. - # Use a distinct operation name to avoid collisions with an existing - # "embedding" layer in the representation model. - embedding_output = tf.keras.layers.Lambda(lambda x: x, name="embedding_output")( - rep_output - ) - - embedding_model = tf.keras.Model( - inputs=rep_model.input, - outputs={"embedding": embedding_output}, - name=f"{model_name}_embedding", - ) - - graph_path = path / f"{model_name}_embedding_graph" - tf.saved_model.save(embedding_model, graph_path) - logger.info(f"embedding-only computational graph is written to {graph_path}") - - def get_callbacks(self, branch: str = "classifier") -> list[Any]: - """Build Keras callback list from config for a given *branch*.""" - cb_list = self.train_cfg.get("callbacks", dict()) - callbacks = [] - for cb in cb_list.get(branch, []): - name = cb.get("name") - params = cb.get("params", {}) - try: - cb_class = getattr(tf.keras.callbacks, name) - callbacks.append(cb_class(**params)) - except AttributeError: - raise ValueError(f"Unsupported callback: {name}") - return callbacks - - # ------------------------------------------------------------------ - # Config helpers - # ------------------------------------------------------------------ - - def _get_string_processor_config(self) -> dict[str, Any]: - _map = { - "CODON": CODONS, - "CODON_ID": CODON_ID, - "AA_ID": AA_ID, - "MURPHY10_ID": MURPHY10_ID, - "PC5_ID": PC5_ID, - "DICODON": DICODONS, - "DICODON_ID": DICODON_ID, - } - _config = dict(self.model_cfg.get("embedding", {})) - _config.update(self.model_cfg.get("string_processor", {})) - _config["seq_onehot"] = _config.get("seq_onehot", False) - if _config["input_type"] == "translated": - _config["codon"] = _map.get(_config.get("codon")) - _config["codon_id"] = _map.get(_config.get("codon_id")) - _config["codon_depth"] = max(_config.get("codon_id", [])) + 1 - _config["vocab_size"] = len(_config.get("codon_id", [])) + 1 - _config["ngram_width"] = int(math.log(len(_config.get("codon", [])), 4)) - if _config["seq_onehot"] is False: - _config["codon_depth"] = 1 - else: - _config["codon"] = None - _config["codon_id"] = None - _config["codon_depth"] = None - _config["vocab_size"] = 4 - _config["ngram_width"] = None - return _config - - def _get_optimizer(self, name: str, kwargs: dict) -> Any: - optimizers = { - "adam": tf.keras.optimizers.Adam, - "sgd": tf.keras.optimizers.SGD, - "rmsprop": tf.keras.optimizers.RMSprop, - "adagrad": tf.keras.optimizers.Adagrad, - } - return optimizers[name](**kwargs) - - def _get_pooler(self, name: str): - poolers = { - "max": tf.keras.layers.GlobalMaxPooling2D, - "average": tf.keras.layers.GlobalAveragePooling2D, - "gatedframe": GatedFrameGlobalMaxPooling, - } - return poolers[name] - - def _get_loss(self, name: str, kwargs: dict) -> Any: - losses = { - "categorical_crossentropy": tf.keras.losses.CategoricalCrossentropy, - "sparse_categorical_crossentropy": tf.keras.losses.SparseCategoricalCrossentropy, - "binary_crossentropy": tf.keras.losses.BinaryCrossentropy, - "mse": tf.keras.losses.MeanSquaredError, - "hierachical_loss": HierarchicalLoss, - } - return losses[name](**kwargs) - - def _get_paths(self, key: str) -> dict[str, Any]: - fcd_dict = self.train_cfg.get(key, {}) - paths: dict[str, Any] = {} - for fcd_k, fcd_v in fcd_dict.items(): - tmp_paths = [] - tmp_class = [] - tmp_label = [] - for i in fcd_v: - tmp_class.extend(i.get("class", [])) - tmp_label.extend(i.get("label", [])) - tmp_paths.extend(i.get("path", [])) - paths[fcd_k] = { - "paths": tmp_paths, - "class": tmp_class, - "label": tmp_label, - } - return paths - - def _get_last_checkpoint(self) -> None: - """Placeholder for checkpoint resumption logic.""" - pass - - def _get_model_saving_configuration(self) -> dict: - return self.train_cfg.get("model_saving", {}) - - def _get_fragment_paths(self) -> dict[str, Any]: - return self._get_paths("fragment_classifier_data") - - def _get_contig_paths(self) -> dict[str, Any]: - return self._get_paths("contig_classifier_data") - - def _get_reliability_fragment_paths(self) -> dict[str, Any]: - return self._get_paths("fragment_reliability_data") - - def _get_reliability_contig_paths(self) -> dict[str, Any]: - return self._get_paths("contig_reliability_data") diff --git a/src/jaeger/nnlib/conversion.py b/src/jaeger/nnlib/conversion.py deleted file mode 100644 index 3efd060..0000000 --- a/src/jaeger/nnlib/conversion.py +++ /dev/null @@ -1,610 +0,0 @@ -"""Model graph conversion utilities for Jaeger. - -Supports optimizing SavedModel graphs for inference via: -- XLA (JIT compilation) -- TensorFlow Lite (mobile/edge deployment) -- ONNX (cross-platform deployment with TensorRT/CUDA support) -- TensorRT (NVIDIA GPU acceleration) - -This module is independent of the CLI layer and can be imported directly: - - from jaeger.nnlib.conversion import convert_graph - convert_graph("default", Path("./optimized"), mode="xla", verbose=1) -""" - -from __future__ import annotations - -import logging -import shutil -import sys -from pathlib import Path - -import numpy as np -import tensorflow as tf -import yaml -from tensorflow.python.framework.convert_to_constants import ( - convert_variables_to_constants_v2, -) - -from jaeger.utils.misc import AvailableModels - -logger = logging.getLogger("Jaeger") - - -def convert_graph( - model: str, - output: Path, - mode: str, - verbose: int = 0, - int8: bool = False, -): - """Core graph conversion logic (non-CLI). - - Parameters - ---------- - model: - Model name or identifier (resolved via ``AvailableModels``). - output: - Output directory for the converted model. - mode: - One of ``"xla"``, ``"tflite"``, ``"onnx"``, ``"tensorrt"``. - verbose: - Verbosity level (currently unused but kept for API consistency). - int8: - If ``True`` and *mode* is ``"onnx"``, apply static INT8 quantization. - """ - log = logging.getLogger("Jaeger") - log.info(f"Converting model '{model}' with mode '{mode}'") - - model_info = _resolve_model(model) - if model_info is None: - raise ValueError(f"Model '{model}' not found") - - graph_dir = Path(model_info["graph"]) - output = Path(output) - output.mkdir(parents=True, exist_ok=True) - - suffix = "_int8" if (mode == "onnx" and int8) else "" - converted_dir = output / f"{model}_{mode}{suffix}" - if converted_dir.exists(): - shutil.rmtree(converted_dir) - converted_dir.mkdir(parents=True) - - for key in ("classes", "project", "weights"): - if key in model_info and model_info[key]: - src = Path(model_info[key]) - dst = converted_dir / src.name - shutil.copy2(src, dst) - log.info(f"Copied {key}: {src.name}") - - if mode == "xla": - _convert_xla(graph_dir, converted_dir, model, log) - elif mode == "tflite": - _convert_tflite(graph_dir, converted_dir, model, log) - elif mode == "onnx": - _convert_onnx(graph_dir, converted_dir, model, log, int8=int8) - elif mode == "tensorrt": - _convert_tensorrt(graph_dir, converted_dir, model, log) - else: - raise ValueError(f"Unknown conversion mode: {mode}") - - info = { - "original_graph": str(graph_dir), - "conversion_mode": mode, - "int8_quantization": bool(int8) if mode == "onnx" else None, - "output_dir": str(converted_dir), - } - info_path = converted_dir / "conversion_info.yaml" - info_path.write_text(yaml.safe_dump(info)) - - log.info(f"Conversion complete. Output: {converted_dir}") - - -def _resolve_model(model_name: str) -> dict | None: - """Resolve model name to a path dict via ``AvailableModels``.""" - from importlib.resources import files - - from jaeger.utils.misc import json_to_dict - - CONFIG_PATH = files("jaeger.data") / "config.json" - model_paths = json_to_dict(CONFIG_PATH).get("model_paths", []) - models = AvailableModels(path=model_paths) - return models.info.get(model_name) - - -# ------------------------------------------------------------------ -# XLA -# ------------------------------------------------------------------ - - -def _get_input_spec(infer): - """Return a tf.TensorSpec for the first real input of a serving signature. - - Resource tensors (SavedModel variables) are ignored. - """ - input_tensor = next( - tensor for tensor in infer.inputs if tensor.dtype != tf.resource - ) - return tf.TensorSpec( - shape=input_tensor.shape, - dtype=input_tensor.dtype, - name=input_tensor.name.split(":")[0], - ) - - -def _get_rep_input(infer): - """Build a representative input tensor from a SavedModel serving signature. - - Unknown dimensions are replaced with representative sizes (1 for batch, - 500 for variable-length axes). - """ - input_spec = _get_input_spec(infer) - shape = [] - for dim in input_spec.shape.as_list(): - if dim is None: - # Use 1 for the batch dimension and 500 for variable-length axes. - shape.append(1 if not shape else 500) - else: - shape.append(dim) - dtype = input_spec.dtype - if dtype.is_integer: - return tf.constant( - np.random.randint(0, 64, size=shape).astype(dtype.as_numpy_dtype) - ) - return tf.constant(np.random.randn(*shape).astype(dtype.as_numpy_dtype)) - - -def _convert_xla( - graph_dir: Path, output_dir: Path, model_name: str, log: logging.Logger -): - """Convert SavedModel to XLA-optimized SavedModel. - - XLA (Accelerated Linear Algebra) uses JIT compilation to fuse operations - and optimize memory access patterns. This can provide 2-3x speedup on GPU - for compute-bound models after initial compilation overhead. - """ - log.info("Loading SavedModel for XLA optimization...") - loaded = tf.saved_model.load(str(graph_dir)) - infer = loaded.signatures["serving_default"] - - log.info("Wrapping with XLA JIT compilation...") - - input_spec = _get_input_spec(infer) - - class XLAModel(tf.Module): - def __init__(self, loaded_model): - super().__init__() - # Track the original model so its variables are saved with us. - self.loaded_model = loaded_model - - @tf.function(input_signature=[input_spec], jit_compile=True) - def serving_default(self, inputs): - return self.loaded_model.signatures["serving_default"](inputs=inputs) - - xla_module = XLAModel(loaded) - - log.info("Tracing XLA graph (this may take a while)...") - rep_input = _get_rep_input(infer) - _ = xla_module.serving_default(rep_input) - - log.info(f"Saving XLA-optimized model to {output_dir}...") - tf.saved_model.save(xla_module, str(output_dir)) - log.info("XLA model saved") - - -# ------------------------------------------------------------------ -# TFLite -# ------------------------------------------------------------------ - - -def _convert_tflite( - graph_dir: Path, output_dir: Path, model_name: str, log: logging.Logger -): - """Convert SavedModel to TFLite (same as quantize but without quantization).""" - log.info("Loading SavedModel for TFLite conversion...") - loaded = tf.saved_model.load(str(graph_dir)) - infer = loaded.signatures["serving_default"] - - log.info("Converting to frozen graph...") - frozen_func = convert_variables_to_constants_v2(infer) - - log.info("Converting to TFLite...") - converter = tf.lite.TFLiteConverter.from_concrete_functions([frozen_func]) - converter.optimizations = [tf.lite.Optimize.DEFAULT] - # Jaeger models use ops (e.g., Gelu, BatchMatMul) that are not in the - # baseline TFLite builtin set; allow TensorFlow fallback ops. - converter.target_spec.supported_ops = [ - tf.lite.OpsSet.TFLITE_BUILTINS, - tf.lite.OpsSet.SELECT_TF_OPS, - ] - - tflite_model = converter.convert() - - tflite_path = output_dir / f"{model_name}.tflite" - tflite_path.write_bytes(tflite_model) - log.info(f"Saved TFLite model: {tflite_path} ({len(tflite_model) / 1024:.1f} KB)") - - -# ------------------------------------------------------------------ -# ONNX -# ------------------------------------------------------------------ - - -def _convert_onnx( - graph_dir: Path, - output_dir: Path, - model_name: str, - log: logging.Logger, - int8: bool = False, -): - """Convert SavedModel to ONNX format. - - ONNX models can be run with ONNX Runtime, which supports multiple - execution providers (TensorRT, CUDA, CPU) without requiring - TensorFlow to be built with those backends. - - This is the recommended approach for TensorRT acceleration when - using pip-installed TensorFlow. - - Args: - int8: If True, apply static INT8 quantization using ONNX Runtime's - quantization tools. This produces a smaller model that can run - on TensorRT's INT8 tensor cores. Requires a calibration step. - """ - try: - import tf2onnx # noqa: F401 - except ImportError: - log.error( - "tf2onnx is not installed. Install it with:\n" - " pip install tf2onnx onnxruntime-gpu" - ) - sys.exit(1) - - log.info("Converting SavedModel to ONNX...") - onnx_path = output_dir / f"{model_name}.onnx" - - import subprocess - - cmd = [ - sys.executable, - "-m", - "tf2onnx.convert", - "--saved-model", - str(graph_dir), - "--output", - str(onnx_path), - "--opset", - "13", - "--signature", - "serving_default", - "--tag", - "serve", - ] - - result = subprocess.run(cmd, capture_output=True, text=True) - if result.returncode != 0: - log.error(f"ONNX conversion failed:\n{result.stderr}") - sys.exit(1) - - log.info( - f"Saved ONNX model: {onnx_path} ({onnx_path.stat().st_size / 1024 / 1024:.1f} MB)" - ) - - try: - import onnx - - onnx_model = onnx.load(str(onnx_path)) - onnx.checker.check_model(onnx_model) - log.info("ONNX model validation passed") - except Exception as e: - log.warning(f"ONNX validation warning: {e}") - - if int8: - _quantize_onnx_int8(onnx_path, graph_dir, output_dir, model_name, log) - - -def _quantize_onnx_int8( - onnx_path: Path, - graph_dir: Path, - output_dir: Path, - model_name: str, - log: logging.Logger, -): - """Apply static INT8 quantization to an ONNX model. - - Uses ONNX Runtime's quantization tools with a calibration dataset - derived from real DNA sequences. The quantized model is saved in the - output directory as ``{model_name}_int8.onnx``. - """ - try: - from onnxruntime.quantization import ( - CalibrationDataReader, - QuantFormat, - QuantType, - quantize_static, - ) - except ImportError as e: - log.error( - "ONNX Runtime quantization tools are not available. " - f"Install with: pip install onnxruntime-gpu sympy\n{e}" - ) - sys.exit(1) - - import onnx - - log.info("Preparing ONNX model for INT8 quantization...") - - try: - inferred_path = onnx_path.with_suffix(".inferred.onnx") - onnx_model = onnx.load(str(onnx_path)) - inferred_model = onnx.shape_inference.infer_shapes(onnx_model) - onnx.save(inferred_model, str(inferred_path)) - input_model_path = str(inferred_path) - preprocessed_path = inferred_path - except Exception as e: - log.warning(f"ONNX shape inference failed, using original model: {e}") - input_model_path = str(onnx_path) - preprocessed_path = None - - onnx_model = onnx.load(input_model_path) - input_tensor = onnx_model.graph.input[0] - input_shape = [dim.dim_value for dim in input_tensor.type.tensor_type.shape.dim] - batch = input_shape[0] if input_shape[0] > 0 else 1 - frames = input_shape[1] if input_shape[1] > 0 else 6 - seq_len = input_shape[2] if input_shape[2] > 0 else 500 - depth = input_shape[3] if len(input_shape) > 3 and input_shape[3] > 0 else 64 - input_name = input_tensor.name - - log.info( - f"Calibration input shape: ({batch}, {frames}, {seq_len}, {depth}); " - f"name={input_name}" - ) - - calibration_inputs = _build_calibration_inputs( - graph_dir=graph_dir, - target_shape=(batch, frames, seq_len, depth), - num_samples=100, - log=log, - ) - - class _Calibrator(CalibrationDataReader): - def __init__(self, input_name: str, inputs: list[np.ndarray]): - self.input_name = input_name - self.inputs = inputs - self._idx = 0 - - def get_next(self): - if self._idx >= len(self.inputs): - return None - item = {self.input_name: self.inputs[self._idx]} - self._idx += 1 - return item - - calibrator = _Calibrator(input_name, calibration_inputs) - - quantized_path = output_dir / f"{model_name}_int8.onnx" - - log.info("Running INT8 static quantization (this may take a while)...") - try: - quantize_static( - model_input=input_model_path, - model_output=str(quantized_path), - calibration_data_reader=calibrator, - quant_format=QuantFormat.QDQ, - activation_type=QuantType.QInt8, - weight_type=QuantType.QInt8, - extra_options={ - "ActivationSymmetric": True, - "WeightSymmetric": True, - }, - ) - except Exception as e: - log.error(f"INT8 quantization failed: {e}") - if preprocessed_path and preprocessed_path.exists(): - preprocessed_path.unlink() - sys.exit(1) - - if preprocessed_path and preprocessed_path.exists(): - preprocessed_path.unlink() - - log.info( - f"Saved INT8 ONNX model: {quantized_path} " - f"({quantized_path.stat().st_size / 1024 / 1024:.1f} MB)" - ) - - -def _build_calibration_inputs( - graph_dir: Path, - target_shape: tuple[int, int, int, int], - num_samples: int, - log: logging.Logger, -) -> list[np.ndarray]: - """Generate calibration inputs by running real DNA through the SavedModel. - - Falls back to Gaussian noise if the real-data path fails. - """ - from importlib.resources import files - - batch, frames, seq_len, depth = target_shape - inputs: list[np.ndarray] = [] - - test_fasta = files("jaeger.data") / "test" / "test_contigs.fasta" - if not test_fasta.exists(): - log.warning( - "No bundled test FASTA found; falling back to synthetic one-hot calibration data. " - "INT8 accuracy may be degraded." - ) - return _synthetic_one_hot_samples(target_shape, num_samples) - - try: - log.info(f"Building calibration dataset from {test_fasta}") - - from jaeger.seqops.io import fragment_generator - from jaeger.seqops.encode import process_string_inference - from jaeger.nnlib.inference import InferModel - - _tmp_info = {"graph": graph_dir, "classes": None, "project": None} - _tmp_infer = InferModel.__new__(InferModel) - _tmp_infer.class_map = {} - _tmp_infer.string_processor_config = _tmp_infer._load_string_processor_config( - None - ) - - project_yaml = graph_dir.parent / f"{graph_dir.parent.name}_project.yaml" - if not project_yaml.exists(): - candidates = list(graph_dir.parent.glob("*_project.yaml")) - if candidates: - project_yaml = candidates[0] - if project_yaml.exists(): - import yaml - - with project_yaml.open() as fh: - _tmp_infer.string_processor_config = yaml.safe_load(fh) - - spc = _tmp_infer.string_processor_config - - dataset = tf.data.Dataset.from_generator( - lambda: fragment_generator( - file_path=str(test_fasta), - no_progress=True, - fragsize=spc.get("fragsize", 200), - stride=spc.get("stride", spc.get("fragsize", 200) // 2), - ), - output_signature=tf.TensorSpec(shape=(), dtype=tf.string), - ) - - process_fn = process_string_inference( - codons=spc.get("codon"), - codon_num=spc.get("codon_id"), - codon_depth=spc.get("codon_depth", 1), - ngram_width=spc.get("ngram_width", 3), - seq_onehot=spc.get("seq_onehot"), - crop_size=spc.get("crop_size"), - input_type=spc.get("input_type", "translated"), - masking=spc.get("masking"), - mutate=spc.get("mutate"), - mutation_rate=spc.get("mutation_rate"), - shuffle=spc.get("shuffle"), - ) - - dataset = dataset.map(process_fn, num_parallel_calls=tf.data.AUTOTUNE) - dataset = dataset.batch(1).prefetch(tf.data.AUTOTUNE) - - input_type = spc.get("input_type", "translated") - - for inputs_dict in dataset.take(num_samples): - x = inputs_dict.get(input_type) - if isinstance(x, tf.Tensor): - x = x.numpy() - if x.shape[2] < seq_len: - pad = seq_len - x.shape[2] - x = np.pad(x, ((0, 0), (0, 0), (0, pad), (0, 0)), mode="constant") - elif x.shape[2] > seq_len: - x = x[:, :, :seq_len, :] - if x.shape[3] != depth: - if x.shape[3] < depth: - pad_d = depth - x.shape[3] - x = np.pad(x, ((0, 0), (0, 0), (0, 0), (0, pad_d)), mode="constant") - else: - x = x[:, :, :, :depth] - inputs.append(x.astype(np.float32)) - - if not inputs: - raise RuntimeError("No calibration inputs generated from FASTA") - - log.info(f"Generated {len(inputs)} real calibration samples") - return inputs - - except Exception as e: - log.warning( - f"Failed to build real calibration dataset ({e}); " - "falling back to synthetic one-hot data. INT8 accuracy may be degraded." - ) - return _synthetic_one_hot_samples(target_shape, num_samples) - - -def _synthetic_one_hot_samples( - target_shape: tuple[int, int, int, int], num_samples: int -) -> list[np.ndarray]: - """Generate random one-hot tensors matching codon-encoded DNA inputs. - - Real Jaeger inputs are one-hot codon tensors (shape: B, 6, L, 64). Using - random one-hot samples for calibration is closer to the true distribution - than Gaussian noise and usually yields better INT8 accuracy. - """ - batch, frames, seq_len, depth = target_shape - samples = [] - for _ in range(num_samples): - indices = np.random.randint(0, depth, size=(batch, frames, seq_len)) - one_hot = np.eye(depth)[indices].astype(np.float32) - samples.append(one_hot) - return samples - - -# ------------------------------------------------------------------ -# TensorRT -# ------------------------------------------------------------------ - - -def _convert_tensorrt( - graph_dir: Path, output_dir: Path, model_name: str, log: logging.Logger -): - """Convert SavedModel to TensorRT-optimized SavedModel. - - Requires TensorFlow built with TensorRT support. Most pip-installed - TensorFlow packages do NOT include TensorRT. Use NVIDIA's NGC containers - or build from source with --config=cuda and TensorRT headers. - - Alternative: Use --mode onnx to create an ONNX model, then run with - ONNX Runtime's TensorRT execution provider. - - When available, TensorRT can provide 2-5x speedup on NVIDIA GPUs by: - - Fusing layers (conv+bias+relu) - - Selecting optimal CUDA kernels - - Using FP16/INT8 tensor cores - """ - try: - from tensorflow.python.compiler.tensorrt import trt_convert as trt - except ImportError: - log.error( - "TensorRT is not available. TensorFlow must be built with TensorRT support.\n" - "Options:\n" - " 1. Use NVIDIA NGC Docker container: nvcr.io/nvidia/tensorflow:xx.xx-tf2-py3\n" - " 2. Install tensorflow[and-cuda] with TensorRT headers present\n" - " 3. Build TensorFlow from source with --config=cuda and TensorRT\n" - " 4. Use --mode onnx instead, then run with ONNX Runtime + TensorRT" - ) - sys.exit(1) - - try: - trt._check_trt_version_compatibility() - except RuntimeError as e: - log.error(f"TensorRT compatibility check failed: {e}") - sys.exit(1) - - log.info("Loading SavedModel for TensorRT optimization...") - - params = trt.TrtConversionParams( - precision_mode=trt.TrtPrecisionMode.FP32, - max_workspace_size_bytes=8000000000, # 8GB - ) - - converter = trt.TrtGraphConverterV2( - input_saved_model_dir=str(graph_dir), - conversion_params=params, - ) - - log.info("Converting to TensorRT graph...") - converter.convert() - - log.info("Building TensorRT engines (this may take several minutes)...") - - def input_fn(): - for _ in range(10): - yield [np.random.randn(1, 6, 500, 64).astype(np.float32)] - - converter.build(input_fn=input_fn) - - log.info(f"Saving TensorRT-optimized model to {output_dir}...") - converter.save(str(output_dir)) - log.info("TensorRT model saved") diff --git a/src/jaeger/nnlib/inference.py b/src/jaeger/nnlib/inference.py deleted file mode 100644 index 1f1f50e..0000000 --- a/src/jaeger/nnlib/inference.py +++ /dev/null @@ -1,1047 +0,0 @@ -from collections import defaultdict -from pathlib import Path -from typing import Any -import yaml -import math -import tensorflow as tf -import numpy as np -from jaeger.utils.misc import track_ms as track -from jaeger.seqops.maps import ( - CODON_ID, - CODONS, - AA_ID, - MURPHY10_ID, - PC5_ID, - DICODONS, - DICODON_ID, -) - -# Neural‐net building blocks -from jaeger.nnlib.v2.layers import ( - GeLU, - ReLU, - MaskedConv1D, - ResidualBlock, -) - - -def evaluate(model, x) -> dict[str, float]: - accum = defaultdict(list) - for j, i in track(x, description="[cyan]Crunching data..."): - y_hat = model(j, training=False) - y_true = i - accum["logits"].append(y_hat["prediction"].numpy()) - accum["y_true"].append(y_true.numpy()) - - # Concatenate over all batches - logits = np.concatenate(accum["logits"], axis=0) - y_true = np.concatenate(accum["y_true"], axis=0) - - # Compute loss (categorical cross-entropy from logits) - loss = tf.reduce_mean( - tf.keras.losses.categorical_crossentropy(y_true, logits, from_logits=True) - ).numpy() - - # Compute accuracy - pred_labels = np.argmax(logits, axis=1) - true_labels = np.argmax(y_true, axis=1) - accuracy = np.mean(pred_labels == true_labels) - - return {"loss": float(loss), "accuracy": float(accuracy)} - - -class JaegerModel(tf.keras.Model): - """ - Custom model for Jaeger with training, testing, and prediction steps. - - Methods: - compile: Compiles the model with loss function, optimizer, and metrics. - train_step: Performs a training step with gradient calculation and - weight updates. - test_step: Performs a testing step with loss calculation and metric - updates. - predict_step: Performs a prediction step with inference mode. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - @tf.function(jit_compile=False) - def predict_step(self, data): - # Unpack the data - x, y = data[0], data[1:] - # set model to inference mode - y_logits = self(x, training=False) - return {"y_hat": y_logits, "meta": y} - - -class DynamicInferenceModelBuilder: - """ - rebuilds a model from a configuration file: probably faster than the - tf.saved.Model graph - """ - - def __init__(self, config: dict[str, Any]): - self.cfg = config - self.model_cfg = config.get("model") - self.class_map = self.model_cfg.get("class_label_map") - tf.random.set_seed(self.model_cfg.get("seed")) - np.random.seed(self.model_cfg.get("seed")) - self.inputs = None - self.outputs = list() - self.string_processor_config = self._get_string_processor_config() - self._regularizer = { - "l2": tf.keras.regularizers.L2, - "l1": tf.keras.regularizers.L1, - } - - match config.get("activation", "gelu"): - case "gelu": - self.Activation = GeLU - case "relu": - self.Activation = ReLU - - self.models = self.build_fragment_classifier() - - def build_fragment_classifier(self): - """ - returns rep_model, classification_head and reliability head - """ - models = {} - # === 1. EMBEDDING === - if "embedding" in self.model_cfg: - inputs, x = self._build_embedding(self.model_cfg["embedding"]) - self.inputs = inputs - else: - raise ValueError("Missing 'embedding' section in config") - - # === 2. REPRESENTATION LEARNER === - if "representation_learner" in self.model_cfg: - x, r_lbn = self._build_representation_learner( - self.model_cfg["representation_learner"], x - ) - else: - raise ValueError("Missing 'representation_learner' section in config") - - # Store the representation (input to the classifier head) for embedding export. - self.representation = x - - # === 3. CLASSIFIER === - if "classifier" in self.model_cfg: - y = self._build_head(self.model_cfg["classifier"], x, prefix="cls_") - models["classifier"] = tf.keras.Model(inputs=self.inputs, outputs=y) - else: - raise ValueError("Missing 'classifier' section in config") - - # === 4. RELIABILITY MODEL === - if "reliability_model" in self.model_cfg: - r = self._build_head( - self.model_cfg["reliability_model"], r_lbn, prefix="rel_" - ) - models["reliability"] = tf.keras.Model(inputs=self.inputs, outputs=r) - else: - raise ValueError("Missing 'reliability_model' section in config") - - # === 5. COMBINED MODEL === - outputs = {"prediction": y} - if "reliability_model" in self.model_cfg: - outputs["reliability"] = r - models["jaeger_model"] = tf.keras.Model( - inputs=self.inputs, outputs=outputs, name="jaeger_model" - ) - models["embedding"] = tf.keras.Model( - inputs=self.inputs, outputs=self.representation, name="embedding_model" - ) - - return models - - def _build_embedding(self, cfg: dict[str, Any]): - """Builds the embedding layer based on config.""" - input_shape = cfg.get("input_shape") - seq_type = cfg.get("type", "translated") - - if seq_type == "translated": - inputs = tf.keras.Input(shape=input_shape, name="translated_input") - x = inputs - elif seq_type == "nucleotide": - inputs = tf.keras.Input(shape=input_shape, name="nucleotide_input") - x = inputs - elif seq_type == "both": - # Handle dual input - inputs = tf.keras.Input(shape=input_shape, name="combined_input") - x = inputs - else: - raise ValueError(f"Unknown embedding type: {seq_type}") - - return inputs, x - - def _build_representation_learner(self, cfg: dict[str, Any], x): - """Builds the representation learner (conv blocks).""" - block_sizes = cfg.get("block_sizes", [2, 2, 2]) - block_filters = cfg.get("block_filters", [128, 128, 128]) - block_kernel_size = cfg.get("block_kernel_size", [5, 5, 5]) - block_kernel_dilation = cfg.get("block_kernel_dilation", [3, 3, 3]) - block_kernel_strides = cfg.get("block_kernel_strides", [1, 1, 1]) - block_regularizer = cfg.get("block_regularizer", ["l2", "l2", "l2"]) - block_regularizer_w = cfg.get("block_regularizer_w", [0, 0, 0]) - pooling = cfg.get("pooling", "max") - - # MaskedConv1D and ResidualBlock operate on 4-D frame tensors - # (batch, frames, length, channels). Add a singleton frame axis if a - # plain 3-D tensor (batch, length, channels) is provided. - if len(x.shape) == 3: - x = tf.keras.layers.Lambda( - lambda t: tf.expand_dims(t, axis=1), name="add_frame_axis" - )(x) - - # Initial conv - x = MaskedConv1D( - filters=cfg.get("masked_conv1d_1_filters", 128), - kernel_size=cfg.get("masked_conv1d_1_kernel_size", 7), - strides=cfg.get("masked_conv1d_1_strides", 1), - dilation_rate=cfg.get("masked_conv1d_1_dilation_rate", 1), - padding="same", - kernel_regularizer=self._regularizer.get( - cfg.get("masked_conv1d_1_regularizer", "l2"), - tf.keras.regularizers.L2, - )(cfg.get("masked_conv1d_1_regularizer_w", 0)), - name="masked_conv1d_initial", - )(x) - x = self.Activation(name="activation_initial")(x) - - # Residual blocks - for i, (size, filters, ksize, dilation, strides, reg, reg_w) in enumerate( - zip( - block_sizes, - block_filters, - block_kernel_size, - block_kernel_dilation, - block_kernel_strides, - block_regularizer, - block_regularizer_w, - ) - ): - for j in range(size): - x = ResidualBlock( - filters=filters, - kernel_size=ksize, - dilation_rate=dilation, - strides=strides if j == 0 else 1, - kernel_regularizer=self._regularizer.get( - reg, tf.keras.regularizers.L2 - )(reg_w), - name=f"resblock_{i}_{j}", - )(x) - if pooling == "max": - x = tf.keras.layers.MaxPooling2D(pool_size=(1, 2), name=f"pool_{i}")(x) - elif pooling == "avg": - x = tf.keras.layers.AveragePooling2D( - pool_size=(1, 2), name=f"pool_{i}" - )(x) - - # Final conv - x = MaskedConv1D( - filters=block_filters[-1], - kernel_size=cfg.get("masked_conv1d_final_kernel_size", 5), - strides=cfg.get("masked_conv1d_final_strides", 1), - dilation_rate=cfg.get("masked_conv1d_final_dilation_rate", 1), - padding="same", - kernel_regularizer=self._regularizer.get( - cfg.get("masked_conv1d_final_regularizer", "l2"), - tf.keras.regularizers.L2, - )(cfg.get("masked_conv1d_final_regularizer_w", 0)), - name="masked_conv1d_final", - )(x) - x = self.Activation(name="activation_final")(x) - - # Global pooling for reliability model: pool over both frame and - # length axes to produce a single vector per example. - r_lbn = tf.keras.layers.GlobalAveragePooling2D(name="gap_reliability")(x) - - return x, r_lbn - - def _build_head(self, cfg: dict[str, Any], x, prefix: str = ""): - """Builds a classification head.""" - hidden_layers = cfg.get("hidden_layers", []) - - for i, layer_cfg in enumerate(hidden_layers): - x = tf.keras.layers.Dense( - units=layer_cfg.get("units", 128), - activation=None, - use_bias=layer_cfg.get("use_bias", False), - kernel_regularizer=self._regularizer.get( - layer_cfg.get("kernel_regularizer", "l2"), - tf.keras.regularizers.L2, - )(layer_cfg.get("kernel_regularizer_w", 1e-5)), - name=f"{prefix}dense_{i}", - )(x) - x = self.Activation(name=f"{prefix}activation_dense_{i}")(x) - if layer_cfg.get("dropout_rate", 0) > 0: - x = tf.keras.layers.Dropout( - layer_cfg["dropout_rate"], name=f"{prefix}dropout_{i}" - )(x) - - # Output layer - x = tf.keras.layers.Dense( - units=cfg.get("output_units", 6), - activation=cfg.get("output_activation", None), - use_bias=cfg.get("output_use_bias", False), - name=f"{prefix}output", - )(x) - - return x - - def _get_string_processor_config(self): - """Extract string processor config from model config.""" - # This is a simplified version - return {"input_type": "translated"} - - -class InferModel: - """ - loads a graph given a dict with model graph location and class map - consumnes batched iterators and returns logits per iterator element - (works with latest generation of models) - """ - - def __init__( - self, path_dict, use_xla: bool = False, return_embedding: bool = False - ): - self.class_map = self._load_class_map(path_dict.get("classes")) - self.loaded_model = tf.saved_model.load(path_dict.get("graph")) - self.inference_fn = self.loaded_model.signatures["serving_default"] - self.string_processor_config = self._load_string_processor_config( - path_dict.get("project", None) - ) - self.use_xla = use_xla - self.return_embedding = return_embedding - - if return_embedding and "embedding" not in self.inference_fn.structured_outputs: - raise ValueError( - "The selected model does not expose an 'embedding' output. " - "The taxonomy workflow requires a model trained with an embedding head." - ) - - self._predict_step = self._build_predict_step() - - def _build_predict_step(self): - """Build the prediction step function, optionally with XLA.""" - inference_fn = self.inference_fn - input_type = self.string_processor_config.get("input_type") - - def step(x): - y_logits = inference_fn(inputs=x.get(input_type)) - return y_logits - - if self.use_xla: - return tf.function(jit_compile=True)(step) - else: - return tf.function(jit_compile=False)(step) - - def predict(self, dataset, no_progress: bool = False) -> dict[str, np.ndarray]: - """ - dataset: yields tuples (inputs_dict, meta0, meta1, …) - Returns a dict of numpy arrays, each of shape (N, …) - - To avoid GPU OOM on large datasets, each batch is moved to host - (CPU) memory immediately after inference. - """ - # accumulate NumPy arrays on CPU — prevents GPU memory growth - acc: dict[str, list[np.ndarray]] = defaultdict(list) - - # inline for speed - inf_fn = self._predict_step - - for inputs, *meta in track( - dataset, description="[cyan]Crunching data…", disable=no_progress - ): - # call the tf.function - logits = inf_fn(inputs) - - # move logits to CPU immediately to bound GPU memory usage - for k, t in logits.items(): - acc[k].append(t.numpy()) - - # meta data is already on CPU (strings / python objects) - for idx, m in enumerate(meta): - acc[f"meta_{idx}"].append(m) - - # concatenate NumPy arrays on CPU - result: dict[str, np.ndarray] = {} - for k, arr_list in acc.items(): - result[k] = np.concatenate(arr_list, axis=0) - return result - - def evaluate(self, dataset, no_progress: bool = False) -> dict[str, float]: - """ - dataset: yields tuples (inputs_dict, y_true_onehot) - Returns loss and accuracy - - Moves each batch to CPU immediately to avoid GPU OOM. - """ - logits_acc: list[np.ndarray] = [] - true_acc: list[np.ndarray] = [] - - inf_fn = self._predict_step - - for inputs, y_true in track( - dataset, description="[cyan]Evaluating…", disable=no_progress - ): - logits = inf_fn(inputs)["prediction"] - # move to CPU immediately - logits_acc.append(logits.numpy()) - true_acc.append(y_true.numpy()) - - # concatenate on CPU - logits = np.concatenate(logits_acc, axis=0) - y_true = np.concatenate(true_acc, axis=0) - - # loss & accuracy - loss = tf.keras.losses.categorical_crossentropy( - y_true, logits, from_logits=True - ) - loss = float(tf.reduce_mean(loss).numpy()) - - pred_labels = np.argmax(logits, axis=1) - true_labels = np.argmax(y_true, axis=1) - accuracy = float(np.mean(pred_labels == true_labels)) - - return {"loss": loss, "accuracy": accuracy} - - def _load_class_map(self, path): - if path is None: - return None - with open(path) as f: - _class_map = yaml.safe_load(f)["classes"] - - return { - "num_classes": len(_class_map), - "class": [i["class"] for i in _class_map], - "index": [i["label"] for i in _class_map], - } - - def _load_string_processor_config(self, path) -> dict[str, Any]: - _map = { - "CODON": CODONS, - "CODON_ID": CODON_ID, - "AA_ID": AA_ID, - "MURPHY10_ID": MURPHY10_ID, - "PC5_ID": PC5_ID, - "DICODON": DICODONS, - "DICODON_ID": DICODON_ID, - } - - if path is None: - # Legacy models without project.yaml — return minimal defaults - return {"input_type": "translated"} - - _config = yaml.safe_load(Path(path).read_text()) or {} - _model_cfg = _config.get("model") - _config = _model_cfg.get("embedding") - _config.update(_model_cfg.get("string_processor")) - - # Set input_type from embedding type (e.g., "translated", "nucleotide", "both") - _config["input_type"] = _config.get("type", "translated") - - if _config["codon"] is not None and _config["codon_id"] is not None: - _config["codon"] = _map.get(_config.get("codon")) - _config["codon_id"] = _map.get(_config.get("codon_id")) - _config["codon_depth"] = max(_config.get("codon_id")) + 1 # num_codons - _config["vocab_size"] = len(_config.get("codon_id")) + 1 # num_codons + 1 - _config["ngram_width"] = int(math.log(len(_config["codon"]), 4)) - # Infer seq_onehot from embedding input_shape if not explicitly set. - # input_shape like [6, None, 64] means one-hot with depth=64; - # [6, None] means embedding lookup (no one-hot). - input_shape = _model_cfg.get("embedding", {}).get("input_shape") - if _config.get("seq_onehot") is None and input_shape is not None: - # input_shape is [frames, timesteps, codon_depth] for one-hot - _config["seq_onehot"] = ( - len(input_shape) == 3 - and input_shape[-1] is not None - and input_shape[-1] > 1 - ) - _config["seq_onehot"] = _config.get("seq_onehot", False) - if _config["seq_onehot"] is False: - _config["codon_depth"] = 1 - return _config - - -class TFLiteInferModel: - """ - TFLite inference wrapper for quantized Jaeger models. - - Provides the same interface as InferModel but runs inference via - TensorFlow Lite interpreter. Supports dynamic input resizing for - variable sequence lengths. - """ - - def __init__(self, path_dict): - self.class_map = self._load_class_map(path_dict.get("classes")) - self.string_processor_config = self._load_string_processor_config( - path_dict.get("project", None) - ) - - # Load TFLite model - tflite_path = path_dict.get("tflite") - if tflite_path is None or not Path(tflite_path).exists(): - raise ValueError( - f"TFLite model not found at {tflite_path}. " - "Run 'jaeger utils quantize' first." - ) - - self.interpreter = tf.lite.Interpreter(model_path=str(tflite_path)) - self.input_details = self.interpreter.get_input_details() - self.output_details = self.interpreter.get_output_details() - - def predict(self, dataset, no_progress: bool = False) -> dict[str, np.ndarray]: - """ - dataset: yields tuples (inputs_dict, meta0, meta1, …) - Returns a dict of numpy arrays, each of shape (N, …) - """ - acc: dict[str, list[np.ndarray]] = defaultdict(list) - - for inputs, *meta in track( - dataset, description="[cyan]Crunching data…", disable=no_progress - ): - # Get the input tensor (e.g., "translated") - input_type = self.string_processor_config.get("input_type", "translated") - x = inputs.get(input_type) - - if isinstance(x, tf.Tensor): - x = x.numpy() - - batch_size, *shape = x.shape - - # Resize interpreter if needed - current_shape = self.input_details[0]["shape"] - if list(current_shape) != [batch_size, *shape]: - self.interpreter.resize_tensor_input( - self.input_details[0]["index"], [batch_size, *shape], strict=True - ) - self.interpreter.allocate_tensors() - # Refresh details after resize - self.input_details = self.interpreter.get_input_details() - self.output_details = self.interpreter.get_output_details() - - # Set input and invoke - self.interpreter.set_tensor(self.input_details[0]["index"], x) - self.interpreter.invoke() - - # Get outputs - map to same keys as SavedModel - # Output order: prediction, reliability (based on our conversion) - pred = self.interpreter.get_tensor(self.output_details[0]["index"]) - rel = self.interpreter.get_tensor(self.output_details[1]["index"]) - - acc["prediction"].append(pred) - acc["reliability"].append(rel) - - # meta data - for idx, m in enumerate(meta): - if isinstance(m, tf.Tensor): - m = m.numpy() - acc[f"meta_{idx}"].append(m) - - # concatenate arrays; keep lists for non-array metadata - result: dict[str, np.ndarray] = {} - for k, arr_list in acc.items(): - if k.startswith("meta_"): - # Metadata may be strings or other non-array types - try: - result[k] = np.concatenate(arr_list, axis=0) - except ValueError: - result[k] = arr_list # Keep as list if concatenation fails - else: - result[k] = np.concatenate(arr_list, axis=0) - return result - - def evaluate(self, dataset, no_progress: bool = False) -> dict[str, float]: - """Evaluate on a dataset. Not yet implemented for TFLite.""" - raise NotImplementedError("TFLite evaluation not yet implemented") - - def _load_class_map(self, path): - with open(path) as f: - _class_map = yaml.safe_load(f)["classes"] - return { - "num_classes": len(_class_map), - "class": [i["class"] for i in _class_map], - "index": [i["label"] for i in _class_map], - } - - def _load_string_processor_config(self, path) -> dict[str, Any]: - _map = { - "CODON": CODONS, - "CODON_ID": CODON_ID, - "AA_ID": AA_ID, - "MURPHY10_ID": MURPHY10_ID, - "PC5_ID": PC5_ID, - "DICODON": DICODONS, - "DICODON_ID": DICODON_ID, - } - - if path is None: - return {"input_type": "translated"} - - _config = yaml.safe_load(Path(path).read_text()) or {} - _model_cfg = _config.get("model") - _config = _model_cfg.get("embedding") - _config.update(_model_cfg.get("string_processor")) - _config["input_type"] = _config.get("type", "translated") - - if _config["codon"] is not None and _config["codon_id"] is not None: - _config["codon"] = _map.get(_config.get("codon")) - _config["codon_id"] = _map.get(_config.get("codon_id")) - _config["codon_depth"] = max(_config.get("codon_id")) + 1 - _config["vocab_size"] = len(_config.get("codon_id")) + 1 - _config["ngram_width"] = int(math.log(len(_config["codon"]), 4)) - input_shape = _model_cfg.get("embedding", {}).get("input_shape") - if _config.get("seq_onehot") is None and input_shape is not None: - _config["seq_onehot"] = ( - len(input_shape) == 3 - and input_shape[-1] is not None - and input_shape[-1] > 1 - ) - _config["seq_onehot"] = _config.get("seq_onehot", False) - if _config["seq_onehot"] is False: - _config["codon_depth"] = 1 - return _config - - -class ONNXEngine: - """ - ONNX Runtime inference engine for Jaeger models. - - Loads an ONNX model and runs inference via ONNX Runtime. - Supports multiple execution providers: - - TensorRT (fastest on NVIDIA GPUs) - - CUDA (NVIDIA GPU fallback) - - CPU (universal fallback) - - The ONNX model should be created using: - jaeger utils convert-graph -m -o --mode onnx - """ - - def __init__( - self, - path_dict, - providers: list[str] | None = None, - provider_options: list[dict] | None = None, - ): - self.class_map = self._load_class_map(path_dict.get("classes")) - self.string_processor_config = self._load_string_processor_config( - path_dict.get("project", None) - ) - - # Load ONNX model - onnx_path = path_dict.get("onnx") - if onnx_path is None or not Path(onnx_path).exists(): - raise ValueError( - f"ONNX model not found at {onnx_path}. " - "Run 'jaeger utils convert-graph --mode onnx' first." - ) - - # Preload pip-installed NVIDIA libraries so ONNX Runtime providers - # (CUDA, TensorRT) can find them without requiring LD_LIBRARY_PATH - # to be set externally. - self._preload_cuda_libs() - - try: - import onnxruntime as ort - except ImportError: - raise ImportError( - "onnxruntime is not installed. Install it with:\n" - " pip install onnxruntime-gpu # for GPU\n" - " pip install onnxruntime # for CPU only" - ) - - # Detect INT8 quantized models so we can avoid TensorRT, which has - # limited support for arbitrary quantized ONNX subgraphs. - is_int8 = self._is_int8_model(str(onnx_path)) - - # Auto-select providers if not specified - if providers is None: - available = ort.get_available_providers() - if is_int8: - # TensorRT's ONNX parser has strict requirements for quantized - # ops (symmetric quantization, no dynamic shapes, etc.). INT8 - # ONNX models run reliably on the CUDA execution provider. - preferred = [ - "CUDAExecutionProvider", - "CPUExecutionProvider", - ] - else: - # Prefer TensorRT > CUDA > CPU for FP32/FP16 models - preferred = [ - "TensorrtExecutionProvider", - "CUDAExecutionProvider", - "CPUExecutionProvider", - ] - providers = [p for p in preferred if p in available] - if not providers: - providers = ["CPUExecutionProvider"] - - self.providers = providers - sess_options = ort.SessionOptions() - if provider_options is None: - provider_options = [{} for _ in providers] - self.session = ort.InferenceSession( - str(onnx_path), - sess_options=sess_options, - providers=providers, - provider_options=provider_options, - ) - - self.input_name = self.session.get_inputs()[0].name - self.output_names = [o.name for o in self.session.get_outputs()] - - @staticmethod - def _is_int8_model(onnx_path: str) -> bool: - """Heuristic: check if ONNX model contains INT8 quantized tensors.""" - try: - import onnx - - model = onnx.load(onnx_path) - for init in model.graph.initializer: - if init.data_type == onnx.TensorProto.INT8: - return True - for node in model.graph.node: - if node.op_type in ("QuantizeLinear", "DequantizeLinear"): - return True - return False - except Exception: - return False - - @staticmethod - def _preload_cuda_libs(): - """Preload pip-installed NVIDIA shared libraries via ctypes. - - ONNX Runtime's CUDA/TensorRT providers are compiled against specific - cuDNN/TensorRT SONAMEs. When these are installed as pip packages - (e.g., nvidia-cudnn-cu12, tensorrt), they live under site-packages - and are not on the system's LD_LIBRARY_PATH. Preloading with - RTLD_GLOBAL makes their symbols available to the provider libraries - when onnxruntime loads them. - """ - import ctypes - import site - - # Libraries that ONNX Runtime CUDA/TensorRT providers depend on. - # Preload the major SONAMEs; actual filenames may include version suffixes. - lib_patterns = [ - ("nvidia/cudnn/lib", "libcudnn.so"), - ("tensorrt_libs", "libnvinfer.so"), - ("tensorrt_libs", "libnvonnxparser.so"), - ] - - preloaded = [] - for site_dir in site.getsitepackages(): - for subdir, lib_name in lib_patterns: - lib_dir = Path(site_dir) / subdir - if not lib_dir.exists(): - continue - # Find the newest matching .so file - candidates = sorted( - lib_dir.glob(f"{lib_name}*"), - key=lambda p: p.stat().st_mtime, - reverse=True, - ) - for candidate in candidates: - # Only load actual shared objects (not .a or .la) - if ( - not candidate.name.endswith(".so") - and ".so." not in candidate.name - ): - continue - try: - ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL) - preloaded.append(candidate.name) - break # Only load one version per pattern - except OSError: - pass - - # Also update LD_LIBRARY_PATH for any child processes - import os - - nvidia_lib_dirs = [] - for site_dir in site.getsitepackages(): - nvidia_base = Path(site_dir) / "nvidia" - if nvidia_base.exists(): - nvidia_lib_dirs.extend( - sorted(str(d) for d in nvidia_base.rglob("lib") if d.is_dir()) - ) - trt_libs = Path(site_dir) / "tensorrt_libs" - if trt_libs.exists(): - nvidia_lib_dirs.append(str(trt_libs)) - - if nvidia_lib_dirs: - current = os.environ.get("LD_LIBRARY_PATH", "") - current_parts = current.split(":") if current else [] - new_parts = [p for p in nvidia_lib_dirs if p not in current_parts] - if new_parts: - os.environ["LD_LIBRARY_PATH"] = ":".join(new_parts + current_parts) - - def predict(self, dataset, no_progress: bool = False) -> dict[str, np.ndarray]: - """ - dataset: yields tuples (inputs_dict, meta0, meta1, …) - Returns a dict of numpy arrays, each of shape (N, …) - """ - acc: dict[str, list[np.ndarray]] = defaultdict(list) - input_name = self.input_name - output_names = self.output_names - session = self.session - input_type = self.string_processor_config.get("input_type", "translated") - - for inputs, *meta in track( - dataset, description="[cyan]Crunching data…", disable=no_progress - ): - x = inputs.get(input_type) - if isinstance(x, tf.Tensor): - x = x.numpy() - - # Run inference - outputs = session.run(output_names, {input_name: x}) - - # Map outputs (prediction, reliability) - acc["prediction"].append(outputs[0]) - acc["reliability"].append(outputs[1]) - - # meta data - for idx, m in enumerate(meta): - if isinstance(m, tf.Tensor): - m = m.numpy() - acc[f"meta_{idx}"].append(m) - - # concatenate arrays; keep lists for non-array metadata - result: dict[str, np.ndarray] = {} - for k, arr_list in acc.items(): - if k.startswith("meta_"): - try: - result[k] = np.concatenate(arr_list, axis=0) - except ValueError: - result[k] = arr_list - else: - result[k] = np.concatenate(arr_list, axis=0) - return result - - def evaluate(self, dataset, no_progress: bool = False) -> dict[str, float]: - """Evaluate on a dataset.""" - logits_acc: list[np.ndarray] = [] - true_acc: list[np.ndarray] = [] - input_name = self.input_name - session = self.session - input_type = self.string_processor_config.get("input_type", "translated") - - for inputs, y_true in track( - dataset, description="[cyan]Evaluating…", disable=no_progress - ): - x = inputs.get(input_type) - if isinstance(x, tf.Tensor): - x = x.numpy() - - outputs = session.run(None, {input_name: x}) - logits = outputs[0] # prediction is first output - - logits_acc.append(logits) - true_acc.append(y_true.numpy() if isinstance(y_true, tf.Tensor) else y_true) - - logits = np.concatenate(logits_acc, axis=0) - y_true = np.concatenate(true_acc, axis=0) - - loss = tf.keras.losses.categorical_crossentropy( - y_true, logits, from_logits=True - ) - loss = float(tf.reduce_mean(loss).numpy()) - - pred_labels = np.argmax(logits, axis=1) - true_labels = np.argmax(y_true, axis=1) - accuracy = float(np.mean(pred_labels == true_labels)) - - return {"loss": loss, "accuracy": accuracy} - - def _load_class_map(self, path): - with open(path) as f: - _class_map = yaml.safe_load(f)["classes"] - return { - "num_classes": len(_class_map), - "class": [i["class"] for i in _class_map], - "index": [i["label"] for i in _class_map], - } - - def _load_string_processor_config(self, path) -> dict[str, Any]: - _map = { - "CODON": CODONS, - "CODON_ID": CODON_ID, - "AA_ID": AA_ID, - "MURPHY10_ID": MURPHY10_ID, - "PC5_ID": PC5_ID, - "DICODON": DICODONS, - "DICODON_ID": DICODON_ID, - } - - if path is None: - return {"input_type": "translated"} - - _config = yaml.safe_load(Path(path).read_text()) or {} - _model_cfg = _config.get("model") - _config = _model_cfg.get("embedding") - _config.update(_model_cfg.get("string_processor")) - _config["input_type"] = _config.get("type", "translated") - - if _config["codon"] is not None and _config["codon_id"] is not None: - _config["codon"] = _map.get(_config.get("codon")) - _config["codon_id"] = _map.get(_config.get("codon_id")) - _config["codon_depth"] = max(_config.get("codon_id")) + 1 - _config["vocab_size"] = len(_config.get("codon_id")) + 1 - _config["ngram_width"] = int(math.log(len(_config["codon"]), 4)) - input_shape = _model_cfg.get("embedding", {}).get("input_shape") - if _config.get("seq_onehot") is None and input_shape is not None: - _config["seq_onehot"] = ( - len(input_shape) == 3 - and input_shape[-1] is not None - and input_shape[-1] > 1 - ) - _config["seq_onehot"] = _config.get("seq_onehot", False) - if _config["seq_onehot"] is False: - _config["codon_depth"] = 1 - return _config - - -class TensorRTEngine: - """ - TensorRT inference engine for Jaeger models. - - Loads a TensorRT-optimized SavedModel and runs inference. - Requires TensorFlow built with TensorRT support. - - The TensorRT-optimized model should be created using: - jaeger utils convert-graph -m -o --mode tensorrt - """ - - def __init__(self, path_dict): - self.class_map = self._load_class_map(path_dict.get("classes")) - self.string_processor_config = self._load_string_processor_config( - path_dict.get("project", None) - ) - - # Load TensorRT-optimized SavedModel - trt_dir = path_dict.get("trt_graph") - if trt_dir is None or not Path(trt_dir).exists(): - raise ValueError( - f"TensorRT model not found at {trt_dir}. " - "Run 'jaeger utils convert-graph --mode tensorrt' first.\n" - "Note: TensorRT requires TensorFlow built with TensorRT support." - ) - - self.loaded_model = tf.saved_model.load(str(trt_dir)) - self.inference_fn = self.loaded_model.signatures["serving_default"] - - # Build predict step with XLA for additional optimization - input_type = self.string_processor_config.get("input_type") - inference_fn = self.inference_fn - - @tf.function(jit_compile=True) - def _predict_step(x): - return inference_fn(inputs=x.get(input_type)) - - self._predict_step = _predict_step - - def predict(self, dataset, no_progress: bool = False) -> dict[str, np.ndarray]: - """ - dataset: yields tuples (inputs_dict, meta0, meta1, …) - Returns a dict of numpy arrays, each of shape (N, …) - """ - acc: dict[str, list[np.ndarray]] = defaultdict(list) - inf_fn = self._predict_step - - for inputs, *meta in track( - dataset, description="[cyan]Crunching data…", disable=no_progress - ): - logits = inf_fn(inputs) - for k, t in logits.items(): - acc[k].append(t.numpy()) - for idx, m in enumerate(meta): - acc[f"meta_{idx}"].append(m) - - result: dict[str, np.ndarray] = {} - for k, arr_list in acc.items(): - if k.startswith("meta_"): - try: - result[k] = np.concatenate(arr_list, axis=0) - except ValueError: - result[k] = arr_list - else: - result[k] = np.concatenate(arr_list, axis=0) - return result - - def evaluate(self, dataset, no_progress: bool = False) -> dict[str, float]: - """Evaluate on a dataset.""" - logits_acc: list[np.ndarray] = [] - true_acc: list[np.ndarray] = [] - inf_fn = self._predict_step - - for inputs, y_true in track( - dataset, description="[cyan]Evaluating…", disable=no_progress - ): - logits = inf_fn(inputs)["prediction"] - logits_acc.append(logits.numpy()) - true_acc.append(y_true.numpy()) - - logits = np.concatenate(logits_acc, axis=0) - y_true = np.concatenate(true_acc, axis=0) - - loss = tf.keras.losses.categorical_crossentropy( - y_true, logits, from_logits=True - ) - loss = float(tf.reduce_mean(loss).numpy()) - - pred_labels = np.argmax(logits, axis=1) - true_labels = np.argmax(y_true, axis=1) - accuracy = float(np.mean(pred_labels == true_labels)) - - return {"loss": loss, "accuracy": accuracy} - - def _load_class_map(self, path): - with open(path) as f: - _class_map = yaml.safe_load(f)["classes"] - return { - "num_classes": len(_class_map), - "class": [i["class"] for i in _class_map], - "index": [i["label"] for i in _class_map], - } - - def _load_string_processor_config(self, path) -> dict[str, Any]: - _map = { - "CODON": CODONS, - "CODON_ID": CODON_ID, - "AA_ID": AA_ID, - "MURPHY10_ID": MURPHY10_ID, - "PC5_ID": PC5_ID, - "DICODON": DICODONS, - "DICODON_ID": DICODON_ID, - } - - if path is None: - return {"input_type": "translated"} - - _config = yaml.safe_load(Path(path).read_text()) or {} - _model_cfg = _config.get("model") - _config = _model_cfg.get("embedding") - _config.update(_model_cfg.get("string_processor")) - _config["input_type"] = _config.get("type", "translated") - - if _config["codon"] is not None and _config["codon_id"] is not None: - _config["codon"] = _map.get(_config.get("codon")) - _config["codon_id"] = _map.get(_config.get("codon_id")) - _config["codon_depth"] = max(_config.get("codon_id")) + 1 - _config["vocab_size"] = len(_config.get("codon_id")) + 1 - _config["ngram_width"] = int(math.log(len(_config["codon"]), 4)) - input_shape = _model_cfg.get("embedding", {}).get("input_shape") - if _config.get("seq_onehot") is None and input_shape is not None: - _config["seq_onehot"] = ( - len(input_shape) == 3 - and input_shape[-1] is not None - and input_shape[-1] > 1 - ) - _config["seq_onehot"] = _config.get("seq_onehot", False) - if _config["seq_onehot"] is False: - _config["codon_depth"] = 1 - return _config diff --git a/src/jaeger/nnlib/metrics.py b/src/jaeger/nnlib/metrics.py deleted file mode 100644 index a941f10..0000000 --- a/src/jaeger/nnlib/metrics.py +++ /dev/null @@ -1,106 +0,0 @@ -import tensorflow as tf - - -class PrecisionForClass(tf.keras.metrics.Metric): - def __init__(self, class_id, name=None, **kwargs): - if name is None: - name = f"precision_class{class_id}" - super().__init__(name=name, **kwargs) - self.class_id = class_id - self.tp = self.add_weight(name="tp", initializer="zeros") - self.fp = self.add_weight(name="fp", initializer="zeros") - - def update_state(self, y_true, y_pred, sample_weight=None): - # Convert one-hot → integers if needed - if y_true.shape.rank > 1 and y_true.shape[-1] > 1: - y_true = tf.argmax(y_true, axis=-1) - y_pred = tf.argmax(y_pred, axis=-1) - - y_true = tf.cast(y_true, tf.int32) - y_pred = tf.cast(y_pred, tf.int32) - - y_true_i = tf.equal(y_true, self.class_id) - y_pred_i = tf.equal(y_pred, self.class_id) - - tp = tf.reduce_sum(tf.cast(y_true_i & y_pred_i, self.dtype)) - fp = tf.reduce_sum(tf.cast(~y_true_i & y_pred_i, self.dtype)) - - self.tp.assign_add(tp) - self.fp.assign_add(fp) - - def result(self): - return self.tp / (self.tp + self.fp + tf.keras.backend.epsilon()) - - def reset_state(self): - self.tp.assign(0.0) - self.fp.assign(0.0) - - -class RecallForClass(tf.keras.metrics.Metric): - def __init__(self, class_id, name=None, **kwargs): - if name is None: - name = f"recall_class{class_id}" - super().__init__(name=name, **kwargs) - self.class_id = class_id - self.tp = self.add_weight(name="tp", initializer="zeros") - self.fn = self.add_weight(name="fn", initializer="zeros") - - def update_state(self, y_true, y_pred, sample_weight=None): - if y_true.shape.rank > 1 and y_true.shape[-1] > 1: - y_true = tf.argmax(y_true, axis=-1) - y_pred = tf.argmax(y_pred, axis=-1) - - y_true = tf.cast(y_true, tf.int32) - y_pred = tf.cast(y_pred, tf.int32) - - y_true_i = tf.equal(y_true, self.class_id) - y_pred_i = tf.equal(y_pred, self.class_id) - - tp = tf.reduce_sum(tf.cast(y_true_i & y_pred_i, self.dtype)) - fn = tf.reduce_sum(tf.cast(y_true_i & ~y_pred_i, self.dtype)) - - self.tp.assign_add(tp) - self.fn.assign_add(fn) - - def result(self): - return self.tp / (self.tp + self.fn + tf.keras.backend.epsilon()) - - def reset_state(self): - self.tp.assign(0.0) - self.fn.assign(0.0) - - -class SpecificityForClass(tf.keras.metrics.Metric): - def __init__(self, class_id, name=None, **kwargs): - if name is None: - name = f"specificity_class{class_id}" - super().__init__(name=name, **kwargs) - self.class_id = class_id - self.tn = self.add_weight(name="tn", initializer="zeros") - self.fp = self.add_weight(name="fp", initializer="zeros") - - def update_state(self, y_true, y_pred, sample_weight=None): - # Convert one-hot → integers if needed - if y_true.shape.rank > 1 and y_true.shape[-1] > 1: - y_true = tf.argmax(y_true, axis=-1) - y_pred = tf.argmax(y_pred, axis=-1) - - y_true = tf.cast(y_true, tf.int32) - y_pred = tf.cast(y_pred, tf.int32) - - # "Not class_id" mask - y_true_not_i = tf.not_equal(y_true, self.class_id) - y_pred_not_i = tf.not_equal(y_pred, self.class_id) - - tn = tf.reduce_sum(tf.cast(y_true_not_i & y_pred_not_i, self.dtype)) - fp = tf.reduce_sum(tf.cast(y_true_not_i & ~y_pred_not_i, self.dtype)) - - self.tn.assign_add(tn) - self.fp.assign_add(fp) - - def result(self): - return self.tn / (self.tn + self.fp + tf.keras.backend.epsilon()) - - def reset_state(self): - self.tn.assign(0.0) - self.fp.assign(0.0) diff --git a/src/jaeger/nnlib/v1/__init__.py b/src/jaeger/nnlib/v1/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/jaeger/nnlib/v1/layers.py b/src/jaeger/nnlib/v1/layers.py deleted file mode 100644 index 14d8445..0000000 --- a/src/jaeger/nnlib/v1/layers.py +++ /dev/null @@ -1,784 +0,0 @@ -""" - -Copyright (c) 2024 Yasas Wijesekara - -""" - -import tensorflow as tf -import numpy as np - - -def rc_cnn( - x, name="", filters=16, stride=1, kernel_size=5, dilation_rate=1, padding="same" -): - f = tf.keras.layers.Conv1D( - filters=filters, - kernel_size=kernel_size, - dilation_rate=dilation_rate, - strides=stride, - kernel_initializer=tf.keras.initializers.HeUniform(), - padding=padding, - name=name, - ) - outputs = [f(i) for i in x] - - return outputs - - -class SplitLayer(tf.keras.layers.Layer): - def __init__(self, num_splits, axis=1, **kwargs): - super(SplitLayer, self).__init__(**kwargs) - self.num_splits = num_splits - self.axis = axis - - def build(self, input_shape): - # No trainable weights needed - super(SplitLayer, self).build(input_shape) - - def call(self, inputs): - splits = tf.split(inputs, self.num_splits, axis=self.axis) - return splits - - def compute_output_shape(self, input_shape): - output_shape = list(input_shape) - output_shape[self.axis] = input_shape[self.axis] // self.num_splits - return [tuple(output_shape) for _ in range(self.num_splits)] - - -def rc_batchnorm(x, name): - f = tf.keras.layers.BatchNormalization(name=f"bn_{name}") - - outputs = [f(i) for i in x] - - return outputs - - -def rc_batchnorm2(x, name): - splits = len(x) - x = tf.keras.layers.Concatenate(axis=0)(x) - x = tf.keras.layers.BatchNormalization(name=f"bn_{name}")(x) - x = SplitLayer(axis=0, num_splits=splits)(x) - - return x - - -def rc_maxpool(x, pool_size=2): - f = tf.keras.layers.MaxPooling1D(pool_size=pool_size) - outputs = [f(i) for i in x] - - return outputs - - -class GeLU(tf.keras.layers.Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - # Signal that the layer is safe for mask propagation - self.supports_masking = True - - def call(self, inputs): - return tf.nn.gelu(inputs) - - -def rc_gelu(x): - f = GeLU() - outputs = [] - outputs = [f(i) for i in x] - - return outputs - - -def rc_resnet_block( - x, - name, - kernel_size=[3, 3], - dilation_rate=[1, 1], - filters=[16, 16], - add_residual=True, -): # simple resnet for viruses# - """x: input tensor - name:name for the block - kernel_size: a list specifying the kernel size of each conv layer - dilation_rate: a list specifying the dilation rate of each conv layer - filter: a list specifying the number of filters of each conv layer - shared_weights: whether to use reverse conplement parameter sharing.(True) - add_residual: whether to add residual connections.(True) - """ - - xx = rc_cnn( - x, - name=f"{name}{1}", - filters=filters[0], - kernel_size=kernel_size[0], - padding="same", - dilation_rate=dilation_rate[0], - ) - xx = rc_gelu(xx) - xx = rc_batchnorm(xx, name=f"{name}{1}") - - # Create layers - for n, (k, d, f) in enumerate(zip(kernel_size[1:], dilation_rate[1:], filters[1:])): - xx = rc_cnn( - xx, - name=f"{name}{n + 2}", - filters=f, - kernel_size=k, - padding="same", - dilation_rate=d, - ) - xx = rc_gelu(xx) - xx = rc_batchnorm(xx, name=f"{name}{n + 2}") - - # scale up the skip connection output if the filter sizes are different - - if (filters[-1] != filters[0] or x[-1].shape[-1] != filters[-1]) and add_residual: - x = rc_cnn( - x, - name=f"{name}_skip", - filters=f, - kernel_size=1, - padding="same", - dilation_rate=1, - ) - x = rc_gelu(x) - x = rc_batchnorm(x, name=f"{name}_skip") - - # Add Residue - add = tf.keras.layers.Add() - if add_residual: - outputs = [add(i) for i in zip(x, xx)] - return rc_gelu(outputs) - else: - return rc_gelu(xx) - - -def ConvolutionalTower(inputs, num_res_blocks=5, add_residual=True): - """ - Covolutional tower to increase the receptive filed size based on dilated - convolutions. - - order of operations - - original batch norm paper suggested that BN should be applied - before the non-linear transformation. However, in practice, BN - is applied after the non-linearity as it is shown perform better. - - linear transformation -> non-linearity -> batch norm - """ - - x = rc_cnn( - inputs, - name="block1_0", - filters=128, - stride=1, - kernel_size=9, - dilation_rate=1, - padding="same", - ) - x = rc_gelu(x) - x = rc_batchnorm(x, name="block1_1") - x = rc_maxpool(x, pool_size=2) - - x = rc_cnn( - x, - name="block1_1", - filters=128, - stride=1, - kernel_size=5, - dilation_rate=2, - padding="same", - ) - x = rc_gelu(x) - x = rc_batchnorm(x, name="block1_2") - x = rc_maxpool(x, pool_size=2) - - if num_res_blocks: - for i, n in enumerate(range(num_res_blocks)): - x = ( - lambda x, n: rc_resnet_block( - x, - name=f"block2_{n}", - kernel_size=[5, 5], - dilation_rate=[3 + i, 3 + i], - filters=[128, 128], - add_residual=add_residual, - ) - )(x, n) - - return tf.keras.layers.Add()(x) - - -class PositionalEmbedding(tf.keras.layers.Layer): - def __init__(self, sequence_length, output_dim, **kwargs): - super(PositionalEmbedding, self).__init__(**kwargs) - - position_embedding_matrix = self.get_position_encoding( - sequence_length, output_dim - ) - - self.position_embedding_layer = tf.keras.layers.Embedding( - input_dim=sequence_length, - output_dim=output_dim, - weights=[position_embedding_matrix], - trainable=False, - ) - - def get_position_encoding(self, seq_len, d, n=10000): - P = np.zeros((seq_len, d)) - for k in range(seq_len): - for i in np.arange(int(d / 2)): - denominator = np.power(n, 2 * i / d) - P[k, 2 * i] = np.sin(k / denominator) - P[k, 2 * i + 1] = np.cos(k / denominator) - return P - - def call(self, inputs): - sequence_length = tf.shape(inputs)[-1] - batch_size = tf.shape(inputs)[0] - position_indices = [ - [c for c in range(sequence_length)] for i in range(batch_size) - ] # tf.range(tf.shape(inputs)[-2]) - position_indices = tf.Variable(position_indices) - # embedded_words = self.word_embedding_layer(inputs) - embedded_indices = self.position_embedding_layer(position_indices) - return embedded_indices # embedded_words + - - -class Patches(tf.keras.layers.Layer): - def __init__(self, num_patches, patch_size, name="split"): - super(Patches, self).__init__() - self.num_patches = num_patches - self.patch_size = patch_size - self.layer_name = name - - def call(self, data): - batch_size = tf.shape(data)[0] - splitted_seq = tf.split( - data, - num_or_size_splits=self.num_patches, - axis=1, - num=self.patch_size, - name=self.layer_name, - ) - patches = tf.stack(splitted_seq, axis=1, name="stack") - patches = tf.reshape( - patches, [batch_size, self.num_patches, self.patch_size * 4] - ) - - return patches - - -class PatchEncoder(tf.keras.layers.Layer): - # Parch encoding + Position encoding - def __init__( - self, num_patches, projection_dim=None, embed_input=False, use_sine=True - ): # num_patches == sequence length when input comes from a conv block - super(PatchEncoder, self).__init__() - self.num_patches = num_patches - - if embed_input: - self.projection = tf.keras.layers.Dense(units=projection_dim) - else: - self.projection = None - - if use_sine: - self.position_embedding_layer = tf.keras.layers.Embedding( - input_dim=num_patches, output_dim=projection_dim - ) - else: - position_embedding_matrix = self.get_position_encoding( - num_patches, projection_dim - ) - - self.position_embedding_layer = tf.keras.layers.Embedding( - input_dim=num_patches, - output_dim=projection_dim, - weights=[position_embedding_matrix], - trainable=False, - ) - - def get_position_encoding(self, seq_len, d, n=10000): - P = np.zeros((seq_len, d)) - for k in range(seq_len): - for i in np.arange(int(d / 2)): - denominator = np.power(n, 2 * i / d) - P[k, 2 * i] = np.sin(k / denominator) - P[k, 2 * i + 1] = np.cos(k / denominator) - return P - - def call(self, patches): - positions = tf.range(start=0, limit=self.num_patches, delta=1) - if self.projection is not None: - input_projection = self.projection(patches) - else: - input_projection = patches - - encoded = input_projection + self.position_embedding_layer(positions) - - return encoded - - -def mlp(x, hidden_units, dropout_rate): - for units in hidden_units: - x = tf.keras.layers.Dense(units, activation=tf.nn.gelu)(x) - x = tf.keras.layers.Dropout(dropout_rate)(x) - return x - - -def Baseline_model(input_shape=None): # archeae model 1 - f1input = tf.keras.Input(shape=input_shape, name="forward_1") - f2input = tf.keras.Input(shape=input_shape, name="forward_2") - f3input = tf.keras.Input(shape=input_shape, name="forward_3") - r1input = tf.keras.Input(shape=input_shape, name="reverse_1") - r2input = tf.keras.Input(shape=input_shape, name="reverse_2") - r3input = tf.keras.Input(shape=input_shape, name="reverse_3") - embedding_layer = tf.keras.layers.Embedding(22, 4, name="aa", mask_zero=True) - embeddings = [ - embedding_layer(i) - for i in [f1input, f2input, f3input, r1input, r2input, r3input] - ] - # A block - x = ConvolutionalTower(embeddings, num_res_blocks=None) - x = tf.keras.layers.GlobalMaxPool1D()(x) - # C block - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-1")(x) - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-2")(x) - out = tf.keras.layers.Dense(4, name="outdense")(x) - return [f1input, f2input, f3input, r1input, r2input, r3input], out - - -def Res_model(input_shape=None): # archeae model 1 - f1input = tf.keras.Input(shape=input_shape, name="forward_1") - f2input = tf.keras.Input(shape=input_shape, name="forward_2") - f3input = tf.keras.Input(shape=input_shape, name="forward_3") - r1input = tf.keras.Input(shape=input_shape, name="reverse_1") - r2input = tf.keras.Input(shape=input_shape, name="reverse_2") - r3input = tf.keras.Input(shape=input_shape, name="reverse_3") - embedding_layer = tf.keras.layers.Embedding(22, 4, name="aa", mask_zero=True) - embeddings = [ - embedding_layer(i) - for i in [f1input, f2input, f3input, r1input, r2input, r3input] - ] - # A block - x = ConvolutionalTower(embeddings, num_res_blocks=5) - x = tf.keras.layers.GlobalMaxPool1D()(x) - # C block - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-1")(x) - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-2")(x) - out = tf.keras.layers.Dense(4, name="outdense")(x) - return [f1input, f2input, f3input, r1input, r2input, r3input], out - - -def WRes_model(input_shape=None): # archeae model 1 - f1input = tf.keras.Input(shape=input_shape, name="forward_1") - f2input = tf.keras.Input(shape=input_shape, name="forward_2") - f3input = tf.keras.Input(shape=input_shape, name="forward_3") - r1input = tf.keras.Input(shape=input_shape, name="reverse_1") - r2input = tf.keras.Input(shape=input_shape, name="reverse_2") - r3input = tf.keras.Input(shape=input_shape, name="reverse_3") - embedding_layer = tf.keras.layers.Embedding(22, 4, name="aa", mask_zero=True) - embeddings = [ - embedding_layer(i) - for i in [f1input, f2input, f3input, r1input, r2input, r3input] - ] - # B block - x = ConvolutionalTower(embeddings, num_res_blocks=5, add_residual=False) - x = tf.keras.layers.GlobalMaxPool1D()(x) - # C block - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-1")(x) - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-2")(x) - out = tf.keras.layers.Dense(4, name="outdense")(x) - return ([f1input, f2input, f3input, r1input, r2input, r3input], {"output": out}) - - -def WRes_model_embeddings(input_shape=None, dropout_active=True): - f1input = tf.keras.Input(shape=input_shape, name="forward_1") - f2input = tf.keras.Input(shape=input_shape, name="forward_2") - f3input = tf.keras.Input(shape=input_shape, name="forward_3") - r1input = tf.keras.Input(shape=input_shape, name="reverse_1") - r2input = tf.keras.Input(shape=input_shape, name="reverse_2") - r3input = tf.keras.Input(shape=input_shape, name="reverse_3") - embedding_layer = tf.keras.layers.Embedding(22, 4, name="aa", mask_zero=True) - embeddings = [ - embedding_layer(i) - for i in [f1input, f2input, f3input, r1input, r2input, r3input] - ] - # B block - x = ConvolutionalTower(embeddings, num_res_blocks=5, add_residual=False) - x = tf.keras.layers.GlobalMaxPool1D()(x) - # C block - x = tf.keras.layers.Dropout(0.5)(x, training=dropout_active) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-1")(x) - x = tf.keras.layers.Dropout(0.5)(x, training=dropout_active) - gmp = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-2")(x) - out = tf.keras.layers.Dense(4, name="outdense")(gmp) - return [f1input, f2input, f3input, r1input, r2input, r3input], { - "output": out, - "embedding": gmp, - } - - -def LSTM_model(input_shape=None): - f1input = tf.keras.Input(shape=input_shape, name="forward_1") - f2input = tf.keras.Input(shape=input_shape, name="forward_2") - f3input = tf.keras.Input(shape=input_shape, name="forward_3") - r1input = tf.keras.Input(shape=input_shape, name="reverse_1") - r2input = tf.keras.Input(shape=input_shape, name="reverse_2") - r3input = tf.keras.Input(shape=input_shape, name="reverse_3") - embedding_layer = tf.keras.layers.Embedding(22, 4, name="aa", mask_zero=True) - embeddings = [ - embedding_layer(i) - for i in [f1input, f2input, f3input, r1input, r2input, r3input] - ] - - x = ConvolutionalTower(embeddings, num_res_blocks=5) - x = tf.keras.layers.Bidirectional( - tf.keras.layers.LSTM(128, name="lstm"), name="bidirlstm" - )(x) - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-1")(x) - x = tf.keras.layers.Dropout(0.1)(x) - x = tf.keras.layers.Dense(128, activation=tf.nn.gelu, name="augdense-2")(x) - out = tf.keras.layers.Dense(4, name="outdense")(x) - return [f1input, f2input, f3input, r1input, r2input, r3input], out - - -def Vitra( - input_shape=(None,), - num_patches=512, - transformer_layers=4, - num_heads=4, - att_dropout=0.1, - projection_dim=128, - att_hidden_units=[128, 128], - mlp_hidden_units=[128, 128], - mlp_dropout=0.1, - use_global=True, - global_type="max", -): - f2input = tf.keras.Input(shape=input_shape, name="forward_2") - f3input = tf.keras.Input(shape=input_shape, name="forward_3") - r1input = tf.keras.Input(shape=input_shape, name="reverse_1") - r2input = tf.keras.Input(shape=input_shape, name="reverse_2") - f1input = tf.keras.Input(shape=input_shape, name="forward_1") - r3input = tf.keras.Input(shape=input_shape, name="reverse_3") - embedding_layer = tf.keras.layers.Embedding(22, 4, name="aa", mask_zero=True) - embeddings = [ - embedding_layer(i) - for i in [f1input, f2input, f3input, r1input, r2input, r3input] - ] - # Create patches. - patches = ConvolutionalTower(embeddings, num_res_blocks=5) - # patches = Patches(num_patches=num_patches,patch_size=patch_size)(inputs) - # Encode patches. - encoded_patches = PatchEncoder( - num_patches=num_patches, projection_dim=projection_dim - )(patches) - - # Create multiple layers of the Transformer block. - for _ in range(transformer_layers): - # Layer normalization 1. - x1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)(encoded_patches) - # Create a multi-head attention layer. - attention_output = tf.keras.layers.MultiHeadAttention( - num_heads=num_heads, key_dim=projection_dim, dropout=att_dropout - )(x1, x1) - # Skip connection 1. - x2 = tf.keras.layers.Add()([attention_output, encoded_patches]) - # Layer normalization 2. - x3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)(x2) - # MLP. - x3 = mlp(x3, hidden_units=att_hidden_units, dropout_rate=mlp_dropout) - # Skip connection 2. - encoded_patches = tf.keras.layers.Add()([x3, x2]) - - # Create a [batch_size, projection_dim] tensor. - rep = tf.keras.layers.LayerNormalization(epsilon=1e-6)(encoded_patches) - - if use_global: - if global_type == "average": - rep = tf.keras.layers.GlobalAveragePooling1D()(rep) - elif global_type == "max": - rep = tf.keras.layers.GlobalMaxPooling1D()(rep) - else: - rep = tf.keras.layers.Flatten()(rep) - - rep = tf.keras.layers.Dropout(0.1)(rep) - # Add MLP. - features = mlp(rep, hidden_units=mlp_hidden_units, dropout_rate=0.5) - # Classify outputs. - logits = tf.keras.layers.Dense(4)(features) - # Create the Keras model. - return [f1input, f2input, f3input, r1input, r2input, r3input], logits - - -# layers for the second generation models -class GlobalMaxPoolingPerFeature(tf.keras.layers.Layer): - """ - Apply max_reduce along the last axis. Re-implementation of - GlobalMaxPooling1D layer for Biological sequences - """ - - def __init__(self, **kwargs): - super(GlobalMaxPoolingPerFeature, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the feature axis - return tf.reduce_max( - inputs, axis=-1, keepdims=False, name="global_max_per_position" - ) - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2]) - - -class MaxReduce(tf.keras.layers.Layer): - """Apply max_reduce along the frame axis""" - - def __init__(self, **kwargs): - super(MaxReduce, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the frame axis - return tf.reduce_max(inputs, axis=1, keepdims=False, name="max_reduce") - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2], input_shape[3]) - - -class MeanReduce(tf.keras.layers.Layer): - """Apply mean_reduce along the frame axis""" - - def __init__(self, **kwargs): - super(MeanReduce, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the frame axis - return tf.reduce_mean(inputs, axis=1, keepdims=False, name="max_reduce") - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2], input_shape[3]) - - -class SumReduce(tf.keras.layers.Layer): - """Apply sum_reduce along the frame axis""" - - def __init__(self, **kwargs): - super(SumReduce, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the frame axis - return tf.reduce_sum(inputs, axis=1, keepdims=False, name="max_reduce") - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2], input_shape[3]) - - -class CustomPooling1D(tf.keras.layers.Layer): - """Apply 1D pooling along a defined axis""" - - def __init__(self, pool_size, axis, **kwargs): - super(CustomPooling1D, self).__init__(**kwargs) - self.pool_size = pool_size - self.axis = axis - - def call(self, inputs): - return tf.nn.pool( - inputs, - window_shape=[1, self.pool_size, 1, 1], - strides=[1, self.pool_size, 1, 1], - pooling_type="MAX", - padding="SAME", - data_format="NHWC", - ) - - def compute_output_shape(self, input_shape): - output_shape = list(input_shape) - output_shape[self.axis] = ( - input_shape[self.axis] - self.pool_size - ) // self.pool_size + 1 - return tuple(output_shape) - - -def resnet_block_g2( - x, - name, - kernel_size=[3, 3], - dilation_rate=[1, 1], - filters=[16, 16], - add_residual=True, -): - """ - Args: - ---- - x: input tensor - name:name for the block - kernel_size: a list specifying the kernel size of each conv layer - dilation_rate: a list specifying the dilation rate of each conv layer - filter: a list specifying the number of filters of each conv layer - shared_weights: whether to use reverse conplement parameter sharing. - defaault (True) - add_residual: whether to add residual connections.(True) - """ - - xx = tf.keras.layers.Conv1D( - filters[0], - kernel_size[0], - strides=1, - dilation_rate=dilation_rate[0], - padding="same", - name=f"{name}_{1}", - kernel_initializer=tf.keras.initializers.Orthogonal(gain=2), - )(x) - - xx = tf.keras.layers.BatchNormalization(axis=-1, name=f"{name}_{1}_norm")(xx) - xx = tf.nn.relu(xx) - # Create layers - for n, (k, d, f) in enumerate( - zip(kernel_size[1:], dilation_rate[1:], filters[1:]), 1 - ): - xx = tf.keras.layers.Conv1D( - f, - k, - strides=1, - dilation_rate=d, - padding="same", - name=f"{name}_{n + 2}", - kernel_initializer=tf.keras.initializers.Orthogonal(gain=2), - )(xx) - - xx = tf.keras.layers.BatchNormalization(axis=-1, name=f"{name}_{n + 2}_norm")( - xx - ) - xx = tf.nn.leaky_relu(xx, alpha=0.1) - - # scale up the skip connection output if the filter sizes are different - - if (x.shape[-1] != filters[-1]) and add_residual: - x = tf.keras.layers.Conv1D( - filters[-1], - 1, - strides=1, - dilation_rate=1, - name=f"{name}_skip", - kernel_initializer=tf.keras.initializers.Orthogonal(gain=2), - )(x) - - x = tf.keras.layers.BatchNormalization(axis=-1, name=f"{name}_skip_norm")(x) - x = tf.nn.leaky_relu(x, alpha=0.1) - - # Add Residue - if add_residual: - return tf.keras.layers.Add()([x, xx]) - else: - return xx - - -def ConvolutionalTower_g2(x, num_res_blocks=5, add_residual=True): - """ - Covolutional tower to increase the receptive filed size with dilated - convolutions - """ - - x = tf.keras.layers.Conv1D( - 128, - 9, - strides=1, - dilation_rate=1, - padding="same", - name="conv1", - kernel_initializer=tf.keras.initializers.Orthogonal(gain=2), - )(x) - - x = tf.keras.layers.MaxPooling2D(pool_size=(1, 2))(x) - x = tf.keras.layers.BatchNormalization(axis=-1, name="block1_1")(x) - x = tf.nn.leaky_relu(x, alpha=0.1) - - x = tf.keras.layers.Conv1D( - 128, - 3, - strides=1, - dilation_rate=2, - padding="same", - name="conv2", - kernel_initializer=tf.keras.initializers.Orthogonal(gain=2), - )(x) - - x = tf.keras.layers.BatchNormalization(axis=-1, name="block1_2")(x) - x = tf.nn.leaky_relu(x, alpha=0.1) - - if num_res_blocks: - for i, n in enumerate(range(num_res_blocks)): - x = ( - lambda x, n: resnet_block_g2( - x, - name=f"block2_{n}", - kernel_size=[3, 3], - dilation_rate=[3, 3], - filters=[256, 256], - add_residual=add_residual, - ) - )(x, n) - - return x - - -def create_jaeger_model( - input_shape, vocab_size=22, embedding_size=4, out_shape=6, bias_init=None -): - inputs = tf.keras.Input(shape=input_shape, name="translated") - # embedding_layer = tf.keras.layers.Embedding(vocab_size, - # embedding_size, - # name="aa-embedding", - # mask_zero=True) - - # x = embedding_layer(inputs) - x = ConvolutionalTower_g2(inputs, num_res_blocks=10, add_residual=True) - # A block - x = SumReduce()(x) - x = tf.keras.layers.BatchNormalization(axis=-1, name="sum_reduce_norm")(x) - # create amino acid feature vec - x = tf.keras.layers.GlobalAveragePooling1D()(x) - # x batch-norm and dropout do not play nicely together - # https://doi.org/10.48550/arXiv.1801.05134 - x = tf.keras.layers.Dense( - 32, - activation=tf.nn.relu, - name="augdense-1", - kernel_initializer=tf.keras.initializers.HeNormal(), - kernel_regularizer=tf.keras.regularizers.L2(1e-4), - )(x) - x = tf.keras.layers.BatchNormalization(axis=-1, name="dense1")(x) - x = tf.keras.layers.Dense( - 32, - activation=tf.nn.relu, - name="augdense-2", - kernel_initializer=tf.keras.initializers.HeNormal(), - kernel_regularizer=tf.keras.regularizers.L2(1e-4), - )(x) - x1 = tf.keras.layers.BatchNormalization(axis=-1, name="dense2")(x) - # x = tf.keras.layers.Dense(dense_c2_nodes, - # activation=tf.nn.gelu, - # name='augdense-2', - # kernel_regularizer=tf.keras.regularizers.L2(1e-4))(x) - if bias_init is not None: - bias_init = tf.keras.initializers.Constant(bias_init) - out = tf.keras.layers.Dense( - out_shape, - name="outdense", - dtype="float32", - kernel_initializer=tf.keras.initializers.HeNormal(), - use_bias=True, - bias_initializer=bias_init, - )(x1) # validation loss jumps when bias is removed - - return inputs, {"output": out, "embedding": x1} diff --git a/src/jaeger/nnlib/v2/__init__.py b/src/jaeger/nnlib/v2/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/jaeger/nnlib/v2/layers.py b/src/jaeger/nnlib/v2/layers.py deleted file mode 100644 index c5dc7a2..0000000 --- a/src/jaeger/nnlib/v2/layers.py +++ /dev/null @@ -1,1841 +0,0 @@ -import tensorflow as tf -from tensorflow import keras - -""" -Author: yasas.wijesekara@uni-greifswald.de - -This package provides layers for building deep learning models with biological sequences. -1D convolution is commonplacely used in modelling biological sequences (amino acid -and nucleotide sequences). However, current implementations of 1D convolutional layers -do not support masking which is important as some regions of genomes are either unknown -or has very low complexity, and these regions can interfere with the training process. -Here, we provide layers that support masking and can we incooperated into your projects -by simply adding this file to your project. If you find any bugs please contact me or fix -it and send me a pull request :) -""" - - -## Activations -class GeLU(keras.layers.Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - # Signal that the layer is safe for mask propagation - self.supports_masking = True - - def call(self, inputs): - # Use the tanh approximation so the graph can be converted to TFLite. - return tf.nn.gelu(inputs, approximate=True) - - def compute_output_shape(self, input_shape): - return input_shape - - -class ReLU(keras.layers.Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - # Signal that the layer is safe for mask propagation - self.supports_masking = True - - def call(self, inputs): - return tf.nn.relu(inputs) - - def compute_output_shape(self, input_shape): - return input_shape - - -class SumStands(keras.layers.Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def call(self, inputs): - return tf.math.reduce_sum(inputs, axis=1, keepdims=False) - - def get_config(self): - config = super().get_config() - return config - - -class MaskedAdd(keras.layers.Layer): - def __init__(self, **kwargs): - super(MaskedAdd, self).__init__(**kwargs) - - def call(self, inputs, mask=None): - # inputs: list of tensors - output = tf.math.add_n(inputs) - return output - - def compute_mask(self, inputs, mask=None): - # If all inputs share the same mask (e.g., coming from the same encoder), return that - if mask is None: - return None - if isinstance(mask, list): - # Combine masks (logical AND or OR, depending on use case) - return mask[0] # or tf.logical_and(mask[0], mask[1]), etc. - return mask - - -class CustomPooling1D(tf.keras.layers.Layer): - """Apply 1D pooling along a defined axis""" - - def __init__(self, pool_size, axis, **kwargs): - super(CustomPooling1D, self).__init__(**kwargs) - self.pool_size = pool_size - self.axis = axis - - def call(self, inputs): - return tf.nn.pool( - inputs, - window_shape=[1, self.pool_size, 1, 1], - strides=[1, self.pool_size, 1, 1], - pooling_type="MAX", - padding="SAME", - data_format="NHWC", - ) - - def compute_output_shape(self, input_shape): - output_shape = list(input_shape) - output_shape[self.axis] = ( - input_shape[self.axis] - self.pool_size - ) // self.pool_size + 1 - return tuple(output_shape) - - -class GlobalMaxPoolingPerFeature(tf.keras.layers.Layer): - """'Apply max_reduce along the last axis. Re-implementation of GlobalMaxPooling1D layer for - Biological sequences""" - - def __init__(self, **kwargs): - super(GlobalMaxPoolingPerFeature, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the feature axis - return tf.reduce_max( - inputs, axis=-1, keepdims=False, name="global_max_per_position" - ) - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2]) - - -class MaxReduce(tf.keras.layers.Layer): - """Apply max_reduce along the frame axis""" - - def __init__(self, **kwargs): - super(MaxReduce, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the frame axis - return tf.reduce_max(inputs, axis=1, keepdims=False, name="max_reduce") - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2], input_shape[3]) - - -class MeanReduce(tf.keras.layers.Layer): - """Apply mean_reduce along the frame axis""" - - def __init__(self, **kwargs): - super(MeanReduce, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the frame axis - return tf.reduce_mean(inputs, axis=1, keepdims=False, name="mean_reduce") - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2], input_shape[3]) - - -class MaskedNeuralMean(tf.keras.layers.Layer): - """ - Masked Neural Mean layer returns the channelwise mean of in the input - tensor [batch, strands, length, channels] - This layer has no parameters. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.supports_masking = True - - def call(self, inputs, mask=None, training=False): - if mask is not None: - mask = tf.cast(mask, inputs.dtype) # Ensure mask is the same type as inputs - mask = tf.expand_dims( - mask, axis=-1 - ) # Broadcast mask to have the same number of channels as inputs - valid_elements = tf.reduce_sum( - mask, axis=[0, 1, 2] - ) # Count valid elements per feature map - - # Apply mask to inputs to ignore padded elements - masked_inputs = inputs * mask - - # Calculate mean and variance only over valid elements - mean = tf.reduce_sum(masked_inputs, axis=[0, 1, 2]) / valid_elements - # variance = tf.reduce_sum(mask * tf.square(masked_inputs - mean), axis=[0, 1, 2]) / valid_elements - - else: - # Standard batch normalization when no mask is provided - mean, _ = tf.nn.moments(inputs, axes=[0, 1, 2]) - - return mean - - -class SumReduce(tf.keras.layers.Layer): - """Apply sum_reduce along the frame axis""" - - def __init__(self, **kwargs): - super(SumReduce, self).__init__(**kwargs) - - def call(self, inputs): - # Take the maximum value along the frame axis - return tf.reduce_sum(inputs, axis=1, keepdims=False, name="max_reduce") - - def compute_output_shape(self, input_shape): - # Output shape will have the same batch size and the number of features - return (input_shape[0], input_shape[2], input_shape[3]) - - -class MaskedMaxPooling1D(tf.keras.layers.Layer): - """ - MaxPooling1D implementation that accepts an incomming masking tensor - and outputs a mask tensor corresponding to the output dimention of layer's - output - """ - - def __init__(self, pool_size=2, strides=None, padding="valid", **kwargs): - super(MaskedMaxPooling1D, self).__init__(**kwargs) - self.pool_size = pool_size - self.strides = strides if strides is not None else pool_size - self.padding = padding.lower() - - def call(self, inputs, mask=None): - # Reshape for MaskedMaxpooling1D - input_shape = tf.shape(inputs) - self._input_dtype = inputs.dtype - # batch_size = input_shape[0] - # dim1 = input_shape[1] - dim2 = input_shape[2] - dim3 = input_shape[3] - - if mask is not None: - # Expand mask dimensions to match inputs - mask = tf.cast(mask, dtype=inputs.dtype) - inputs = inputs * tf.expand_dims(mask, axis=-1) # Zero out masked positions - - reshaped_inputs = tf.reshape(inputs, (-1, dim2, dim3)) - outputs = tf.nn.max_pool1d( - reshaped_inputs, - ksize=self.pool_size, - strides=self.strides, - padding=self.padding.upper(), - ) - - # Reshape back to the original dimensions - output_shape = self.compute_output_shape(input_shape) - outputs = tf.reshape(outputs, output_shape) - - if mask is not None: - reshaped_mask = tf.reshape(mask, (-1, dim2)) - mask = self.compute_output_mask(reshaped_mask) - self._output_mask = tf.reshape(mask, output_shape[:-1]) - - else: - self._output_mask = None - - return outputs - - def compute_mask(self, inputs, mask=None): - # Return the output mask computed in call() - if hasattr(self, "_output_mask"): - return self._output_mask - else: - return None - - def compute_output_mask(self, mask=None): - if mask is None: - return None - else: - mask = tf.cast(mask, dtype=self._input_dtype) - mask = tf.expand_dims(mask, axis=-1) - - pooled_mask = tf.nn.max_pool1d( - mask, - ksize=self.pool_size, - strides=self.strides, - padding=self.padding.upper(), - ) - pooled_mask = tf.cast(pooled_mask, dtype=tf.bool) - - # pooled_mask = tf.squeeze(pooled_mask, axis=-1) - return pooled_mask - - def compute_output_shape(self, input_shape): - # Compute the output shape along the pooling axis - axis_length = input_shape[2] - if axis_length is not None: - if self.padding == "valid": - output_length = (axis_length - self.pool_size) // self.strides + 1 - elif self.padding == "same": - output_length = (axis_length + self.strides - 1) // self.strides - - # Create the new shape and replace the dimension for the pooled axis - - return input_shape[0], input_shape[1], output_length, input_shape[-1] - else: - return input_shape[0], input_shape[1], None, input_shape[-1] - - -class MaskedLayerNormalization(tf.keras.layers.Layer): - """ - Masked layer normalization ignores the masked positions when calculating the summary - statistics and normalization - """ - - def __init__(self, epsilon=1e-3, center=True, scale=True, **kwargs): - super(MaskedLayerNormalization, self).__init__(**kwargs) - self.epsilon = epsilon - self.center = center - self.scale = scale - self.supports_masking = True - - def build(self, input_shape): - param_shape = input_shape[-1:] - - if self.scale: - self.gamma = self.add_weight( - name="gamma", shape=param_shape, initializer="ones", trainable=True - ) - else: - self.gamma = None - - if self.center: - self.beta = self.add_weight( - name="beta", shape=param_shape, initializer="zeros", trainable=True - ) - else: - self.beta = None - - super(MaskedLayerNormalization, self).build(input_shape) - - def call(self, inputs, mask=None): - if mask is not None: - mask = tf.cast(mask, inputs.dtype) - mask = tf.expand_dims(mask, -1) - mask = tf.stop_gradient(mask) - - # Compute the mean over unmasked positions - masked_inputs = inputs * mask - mask_sum = tf.reduce_sum(mask, axis=[1, 2], keepdims=True) - mean = tf.reduce_sum(masked_inputs, axis=-1, keepdims=True) / ( - mask_sum + self.epsilon - ) - - # Compute the variance over unmasked positions - squared_diff = tf.square((inputs - mean) * mask) - variance = tf.reduce_sum(squared_diff, axis=-1, keepdims=True) / ( - mask_sum + self.epsilon - ) - else: - mean, variance = tf.nn.moments(inputs, axes=-1, keepdims=True) - - # Normalize the inputs - normalized = (inputs - mean) / tf.sqrt(variance + self.epsilon) - - # Apply scale and center - if self.scale: - normalized = normalized * self.gamma - if self.center: - normalized = normalized + self.beta - - # Re-apply the mask to keep masked positions at zero - if mask is not None: - normalized = normalized * mask - - return normalized - - def compute_output_shape(self, input_shape): - return input_shape - - def get_config(self): - config = super(MaskedLayerNormalization, self).get_config() - config.update( - { - "epsilon": self.epsilon, - "center": self.center, - "scale": self.scale, - } - ) - return config - - -class MaskedGlobalAvgPooling(tf.keras.layers.Layer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def call(self, inputs, mask=None): - # Assuming inputs and mask have the same shape (batch_size, height, width, channels) - if mask is not None: - mask = tf.cast( - mask, inputs.dtype - ) # Convert mask to the same dtype as inputs - mask = tf.expand_dims(mask, axis=-1) - inputs = inputs * mask # Zero out masked positions - - # Sum over the spatial dimensions (batch, frames, seq_length, projections) - masked_sum = tf.reduce_sum(inputs, axis=[1, 2]) - mask_sum = tf.reduce_sum(mask, axis=[1, 2]) - - # Avoid division by zero - mask_sum = tf.maximum(mask_sum, tf.keras.backend.epsilon()) - - # Compute the masked average - return tf.math.divide_no_nan(masked_sum, mask_sum) - else: - # Default to regular global average pooling if no mask is provided - return tf.reduce_mean(inputs, axis=[1, 2]) - - def compute_output_shape(self, input_shape): - return ( - input_shape[0], - input_shape[-1], - ) # Output shape is (batch_size, num_channels) - - -class GatedFrameGlobalMaxPooling(tf.keras.layers.Layer): - """ - Frame-aware global max pooling. - Input: (B, F, L, D) - Output: (B, D) [and optionally (B, F) if return_gate=True] - """ - - def __init__( - self, - return_gate: bool = False, - kernel_initializer="orthogonal", - bias_initializer="zeros", - kernel_regularizer=None, - bias_regularizer=None, - activity_regularizer=None, - **kwargs, - ): - super().__init__(activity_regularizer=activity_regularizer, **kwargs) - self.return_gate = return_gate - self.kernel_initializer = tf.keras.initializers.get(kernel_initializer) - self.bias_initializer = tf.keras.initializers.get(bias_initializer) - self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer) - self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer) - - # child layer will be created in build() - self.score_dense = None - - def build(self, input_shape): - if len(input_shape) != 4: - raise ValueError(f"Expected rank-4 input (B,F,L,D), got {input_shape}") - d = int(input_shape[-1]) - self.score_dense = tf.keras.layers.Dense( - units=1, - use_bias=True, - kernel_initializer=self.kernel_initializer, - bias_initializer=self.bias_initializer, - kernel_regularizer=self.kernel_regularizer, - bias_regularizer=self.bias_regularizer, - name=f"{self.name}_gate", - ) - self.score_dense.build((input_shape[0], input_shape[1], d)) - super().build(input_shape) - - def call(self, x, mask=None, training=None): - # (B,F,L,D) -> max over L - per_frame = tf.reduce_max(x, axis=2) # (B,F,D) - - logits = self.score_dense(per_frame) # (B,F,1) - - gates = tf.sigmoid(logits) - gates = gates / ( - tf.reduce_sum(gates, axis=1, keepdims=True) + tf.keras.backend.epsilon() - ) - - pooled = tf.reduce_sum(per_frame * gates, axis=1) # (B,D) - - if self.return_gate: - return pooled, tf.squeeze(gates, axis=-1) # (B,D), (B,F) - return pooled - - def compute_output_shape(self, input_shape): - b, f, seq_len, d = input_shape - if self.return_gate: - return (b, d), (b, f) - return (b, d) - - # --- New Keras 3 APIs --- - - def get_config(self): - config = super().get_config() - config.update( - { - "return_gate": self.return_gate, - "kernel_initializer": tf.keras.initializers.serialize( - self.kernel_initializer - ), - "bias_initializer": tf.keras.initializers.serialize( - self.bias_initializer - ), - "kernel_regularizer": tf.keras.regularizers.serialize( - self.kernel_regularizer - ), - "bias_regularizer": tf.keras.regularizers.serialize( - self.bias_regularizer - ), - } - ) - return config - - def get_build_config(self): - # record anything needed to rebuild variables - return {"score_dense": self.score_dense.get_config()} - - def build_from_config(self, config): - # re-create child layer from saved config - self.score_dense = tf.keras.layers.Dense.from_config(config["score_dense"]) - - -# class MaskedBatchNorm(tf.keras.layers.Layer): -# """ -# Masked Batch Normalization that supports arbitrary input rank and optional return -# of normalized mean difference vectors. Masked positions are excluded from statistics. -# """ - -# def __init__(self, epsilon=1e-5, momentum=0.9, return_nmd=False,dtype=None, **kwargs): -# super().__init__(**kwargs) -# self.epsilon = epsilon -# self.momentum = momentum -# self.supports_masking = True -# self.return_nmd = return_nmd - -# def build(self, input_shape): -# channel_dim = input_shape[-1] -# if channel_dim is None: -# raise ValueError("The last (channel) dimension must be defined.") -# self.gamma = self.add_weight( -# shape=(channel_dim,), initializer="ones", trainable=True, name="gamma" -# ) -# self.beta = self.add_weight( -# shape=(channel_dim,), initializer="zeros", trainable=True, name="beta" -# ) -# self.moving_mean = self.add_weight( -# shape=(channel_dim,), initializer="zeros", trainable=False, name="moving_mean" -# ) -# self.moving_variance = self.add_weight( -# shape=(channel_dim,), initializer="ones", trainable=False, name="moving_variance" -# ) -# super().build(input_shape) - -# def _vec_broadcast_shape(self, ndims_int, vec): -# """Return [1, 1, ..., C] to reshape a (C,) vector for broadcasting.""" -# c = tf.shape(vec)[0] -# ones = tf.ones([ndims_int - 1], dtype=tf.int32) # [1]*(ndims-1) as a Tensor -# return tf.concat([ones, [c]], axis=0) - -# def call(self, inputs, mask=None, training=False): -# # Use STATIC rank for axes so Keras can infer shapes -# ndims = inputs.shape.rank -# if ndims is None: -# # Fallback: require known rank for this layer -# raise ValueError("Input rank must be statically known for MaskedBatchNorm.") - -# # Axes as Python lists (not Tensors!) -# reduce_axes = list(range(0, max(ndims - 1, 0))) # batch stats: all except channel -# example_axes = list(range(1, max(ndims - 1, 1))) # per-example: skip batch & channel - -# # Prepare mask (once) and masked inputs (once) -# use_mask = mask is not None -# if use_mask: -# mask = tf.cast(mask, inputs.dtype) -# if mask.shape.rank is None or mask.shape.rank < ndims: -# mask = tf.expand_dims(mask, axis=-1) # broadcast over channels -# masked_inputs = inputs * mask - -# valid_elements = tf.reduce_sum(mask, axis=reduce_axes) + self.epsilon -# mean_batch = tf.reduce_sum(masked_inputs, axis=reduce_axes) / valid_elements - -# mean_broadcast_batch = tf.reshape( -# mean_batch, self._vec_broadcast_shape(ndims, mean_batch) -# ) -# variance_batch = ( -# tf.reduce_sum(mask * tf.square(inputs - mean_broadcast_batch), axis=reduce_axes) -# / valid_elements -# ) -# else: -# # axes are Python lists → safe for Keras shape inference -# mean_batch, variance_batch = tf.nn.moments(inputs, axes=reduce_axes) - -# # Pick stats (update EMA during training) -# if training: -# self.moving_mean.assign( -# self.momentum * self.moving_mean + (1.0 - self.momentum) * mean_batch -# ) -# self.moving_variance.assign( -# self.momentum * self.moving_variance + (1.0 - self.momentum) * variance_batch -# ) -# mean_to_use = mean_batch -# var_to_use = variance_batch -# else: -# mean_to_use = self.moving_mean -# var_to_use = self.moving_variance - -# # Normalize (build broadcast shapes once) -# mean_broadcast = tf.reshape(mean_to_use, self._vec_broadcast_shape(ndims, mean_to_use)) -# var_broadcast = tf.reshape(var_to_use, self._vec_broadcast_shape(ndims, var_to_use)) -# inv_std = tf.math.rsqrt(var_broadcast + self.epsilon) - -# normalized = (inputs - mean_broadcast) * inv_std -# output = self.gamma * normalized + self.beta - -# if not self.return_nmd: -# return output - -# # NMD: per-example channel mean (mask-aware) minus reference mean -# if use_mask: -# per_ex_sum = tf.reduce_sum(masked_inputs, axis=example_axes) # (B, C) -# per_ex_count = tf.reduce_sum(mask, axis=example_axes) + self.epsilon # (B, 1) -# mean_channel = per_ex_sum / per_ex_count # (B, C) -# else: -# # If ndims==2, example_axes == [], reduce_mean returns inputs (OK) -# mean_channel = tf.reduce_mean(inputs, axis=example_axes) # (B, C) - -# nmd = mean_channel - mean_to_use # (B, C) - (C,) via broadcasting -# return output, nmd - -# def get_config(self): -# config = super().get_config() -# config.update({ -# "epsilon": self.epsilon, -# "momentum": self.momentum, -# "return_nmd": self.return_nmd, -# }) -# return config - - -class MaskedBatchNorm(tf.keras.layers.Layer): - """ - Masked Batch Normalization that supports arbitrary input rank and optional return - of normalized mean difference vectors. Masked positions are excluded from statistics. - """ - - def __init__( - self, epsilon=1e-5, momentum=0.9, return_nmd=False, dtype=None, **kwargs - ): - # Let Keras / mixed_precision policy control dtype if provided - if dtype is not None: - kwargs["dtype"] = dtype - super().__init__(**kwargs) - - self.epsilon = epsilon - self.momentum = momentum - self.supports_masking = True - self.return_nmd = return_nmd - - def build(self, input_shape): - channel_dim = input_shape[-1] - if channel_dim is None: - raise ValueError("The last (channel) dimension must be defined.") - - # Use variable_dtype so these stay float32 under mixed_float16 - self.gamma = self.add_weight( - name="gamma", - shape=(channel_dim,), - initializer="ones", - trainable=True, - dtype=self.variable_dtype, - ) - self.beta = self.add_weight( - name="beta", - shape=(channel_dim,), - initializer="zeros", - trainable=True, - dtype=self.variable_dtype, - ) - self.moving_mean = self.add_weight( - name="moving_mean", - shape=(channel_dim,), - initializer="zeros", - trainable=False, - dtype=self.variable_dtype, - ) - self.moving_variance = self.add_weight( - name="moving_variance", - shape=(channel_dim,), - initializer="ones", - trainable=False, - dtype=self.variable_dtype, - ) - super().build(input_shape) - - def _vec_broadcast_shape(self, ndims_int, vec): - """Return [1, 1, ..., C] to reshape a (C,) vector for broadcasting.""" - c = tf.shape(vec)[0] - ones = tf.ones([ndims_int - 1], dtype=tf.int32) - return tf.concat([ones, [c]], axis=0) - - def call(self, inputs, mask=None, training=None): - # Force stats math into float32 - x = tf.cast(inputs, tf.float32) - ndims = x.shape.rank - if ndims is None: - raise ValueError("Input rank must be statically known for MaskedBatchNorm.") - - reduce_axes = list( - range(0, max(ndims - 1, 0)) - ) # batch stats: all except channel - example_axes = list( - range(1, max(ndims - 1, 1)) - ) # per-example: skip batch & channel - - use_mask = mask is not None - if use_mask: - mask_f = tf.cast(mask, tf.float32) - if mask_f.shape.rank is None or mask_f.shape.rank < ndims: - mask_f = tf.expand_dims(mask_f, axis=-1) # broadcast over channels - - masked_inputs = x * mask_f - - valid_elements = tf.reduce_sum(mask_f, axis=reduce_axes) + self.epsilon - mean_batch = tf.reduce_sum(masked_inputs, axis=reduce_axes) / valid_elements - - mean_broadcast_batch = tf.reshape( - mean_batch, self._vec_broadcast_shape(ndims, mean_batch) - ) - variance_batch = ( - tf.reduce_sum( - mask_f * tf.square(x - mean_broadcast_batch), axis=reduce_axes - ) - / valid_elements - ) - else: - mean_batch, variance_batch = tf.nn.moments(x, axes=reduce_axes) - - # Ensure stats are float32 - mean_batch = tf.cast(mean_batch, tf.float32) - variance_batch = tf.cast(variance_batch, tf.float32) - - # --- update moving stats in float32, then cast to var dtype on assign --- - mm = tf.cast(self.moving_mean, tf.float32) - mv = tf.cast(self.moving_variance, tf.float32) - - if training: - new_mm = self.momentum * mm + (1.0 - self.momentum) * mean_batch - new_mv = self.momentum * mv + (1.0 - self.momentum) * variance_batch - - self.moving_mean.assign(tf.cast(new_mm, self.moving_mean.dtype)) - self.moving_variance.assign(tf.cast(new_mv, self.moving_variance.dtype)) - - mean_to_use = mean_batch - var_to_use = variance_batch - else: - mean_to_use = mm - var_to_use = mv - - mean_broadcast = tf.reshape( - mean_to_use, self._vec_broadcast_shape(ndims, mean_to_use) - ) - var_broadcast = tf.reshape( - var_to_use, self._vec_broadcast_shape(ndims, var_to_use) - ) - inv_std = tf.math.rsqrt(var_broadcast + tf.cast(self.epsilon, tf.float32)) - - # Normalize + affine in float32 - gamma = tf.cast(self.gamma, tf.float32) - beta = tf.cast(self.beta, tf.float32) - - normalized = (x - mean_broadcast) * inv_std - output_f32 = gamma * normalized + beta - - # Cast back to the layer's compute dtype (float16 under mixed precision) - output = tf.cast(output_f32, self.compute_dtype) - - if not self.return_nmd: - return output - - # NMD in float32, then cast - if use_mask: - per_ex_sum = tf.reduce_sum(masked_inputs, axis=example_axes) # (B, C) - per_ex_count = tf.reduce_sum(mask_f, axis=example_axes) + self.epsilon - mean_channel = per_ex_sum / per_ex_count # (B, C) - else: - mean_channel = tf.reduce_mean(x, axis=example_axes) # (B, C) - - nmd_f32 = mean_channel - mean_to_use # (B, C) - (C,) - nmd = tf.cast(nmd_f32, self.compute_dtype) - - return output, nmd - - def compute_output_shape(self, input_shape): - if self.return_nmd: - batch = input_shape[0] - channels = input_shape[-1] - return (input_shape, (batch, channels)) - return input_shape - - def get_config(self): - config = super().get_config() - config.update( - { - "epsilon": self.epsilon, - "momentum": self.momentum, - "return_nmd": self.return_nmd, - } - ) - return config - - -# class MaskedConv1D(tf.keras.layers.Layer): -# """ -# Masked 1D convolution that accepts an optional mask tensor and sets -# mask positions to zero before applying convolution. Also, propagates the mask -# to the next layer. Accepts inputs like (batch, frames, length, channels) where -# batch and length dimention can be unknown. -# """ - -# def __init__( -# self, -# filters, -# kernel_size, -# strides=1, -# axis=-1, -# padding="valid", -# dilation_rate=1, -# activation=None, -# use_bias=True, -# kernel_initializer="glorot_uniform", -# bias_initializer="zeros", -# kernel_regularizer=None, -# dtype=None, -# **kwargs, -# ): -# super().__init__(**kwargs) -# self.axis = axis -# self.filters = filters -# self.kernel_size = kernel_size -# self.strides = strides -# self.padding = padding.upper() -# self.dilation_rate = dilation_rate -# self.activation = tf.keras.activations.get(activation) -# self.use_bias = use_bias -# self.kernel_initializer = tf.keras.initializers.get(kernel_initializer) -# self.bias_initializer = tf.keras.initializers.get(bias_initializer) -# self.kernel_regularizer = kernel_regularizer - -# def build(self, input_shape): -# channel_axis = self.axis -# input_dim = input_shape[channel_axis] -# kernel_shape = (self.kernel_size, input_dim, self.filters) - -# self.kernel = self.add_weight( -# shape=kernel_shape, -# initializer=self.kernel_initializer, -# regularizer=self.kernel_regularizer, -# name="kernel", -# trainable=True, -# ) -# if self.use_bias: -# self.bias = self.add_weight( -# shape=(self.filters,), -# initializer=self.bias_initializer, -# name="bias", -# trainable=True, -# ) -# else: -# self.bias = None -# super().build(input_shape) - - -# def call(self, inputs, mask=None): -# input_shape = tf.shape(inputs) -# output_shape = self.compute_output_shape(input_shape) - -# output_mask = None -# if mask is not None: -# mask = tf.cast(mask, dtype=inputs.dtype) -# inputs = inputs * tf.expand_dims(mask, axis=-1) - -# # compute output mask here (part of the graph, fine during training/inference) -# reshaped_mask = tf.reshape(mask, (-1, input_shape[2])) -# mask = tf.expand_dims(reshaped_mask, axis=-1) -# mask_kernel = tf.ones((self.kernel_size, 1, 1), dtype=mask.dtype) -# output_mask = tf.nn.conv1d( -# input=mask, -# filters=mask_kernel, -# stride=self.strides, -# padding=self.padding, -# dilations=self.dilation_rate, -# data_format="NWC", -# ) -# output_mask = tf.equal(output_mask, self.kernel_size) -# output_mask = tf.squeeze(output_mask, axis=-1) -# output_mask = tf.reshape(output_mask, shape=output_shape[:-1]) - -# # conv1d on actual data -# reshaped_inputs = tf.reshape(inputs, (-1, input_shape[2], input_shape[3])) -# outputs = tf.nn.conv1d( -# input=reshaped_inputs, -# filters=self.kernel, -# stride=self.strides, -# padding=self.padding, -# dilations=self.dilation_rate, -# data_format="NWC", -# ) -# if self.use_bias: -# outputs = tf.nn.bias_add(outputs, self.bias, data_format="NWC") -# if self.activation is not None: -# outputs = self.activation(outputs) - -# outputs = tf.reshape(outputs, output_shape) - -# # stash mask so compute_mask() can forward it -# self._output_mask = output_mask -# return outputs - -# def compute_mask(self, inputs, mask=None): -# # never recompute, just return the cached mask -# return getattr(self, "_output_mask", None) - -# def get_config(self): -# config = super(MaskedConv1D, self).get_config() -# config.update( -# { -# "filters": self.filters, -# "kernel_size": self.kernel_size, -# "strides": self.strides, -# "padding": self.padding.lower(), -# "dilation_rate": self.dilation_rate, -# "activation": tf.keras.activations.serialize(self.activation), -# "use_bias": self.use_bias, -# "kernel_initializer": tf.keras.initializers.serialize( -# self.kernel_initializer -# ), -# "bias_initializer": tf.keras.initializers.serialize( -# self.bias_initializer -# ), -# } -# ) -# return config - -# def compute_output_shape(self, input_shape): -# """ -# compute the output shape only if the length dimention is not None. -# """ -# length = input_shape[2] -# if length is not None: -# if self.padding == "SAME": -# out_length = (length + self.strides - 1) // self.strides -# elif self.padding == "VALID": -# out_length = ( -# length - self.dilation_rate * (self.kernel_size - 1) - 1 -# ) // self.strides + 1 -# else: -# raise ValueError("Invalid padding type.") -# out_length = out_length -# return (input_shape[0], input_shape[1], out_length, self.filters) -# else: -# return (input_shape[0], input_shape[1], input_shape[2], self.filters) - - -class MaskedConv1D(tf.keras.layers.Layer): - """ - Masked 1D convolution that accepts an optional mask tensor and sets - mask positions to zero before applying convolution. Also propagates the mask - to the next layer. Accepts inputs like (batch, frames, length, channels) where - batch and length dimension can be unknown. - """ - - def __init__( - self, - filters, - kernel_size, - strides=1, - axis=-1, - padding="valid", - dilation_rate=1, - activation=None, - use_bias=True, - kernel_initializer="glorot_uniform", - bias_initializer="zeros", - kernel_regularizer=None, - dtype=None, # keep for compatibility, but usually leave None under MP - **kwargs, - ): - if dtype is not None: - kwargs["dtype"] = dtype # let Keras handle policy + dtype - super().__init__(**kwargs) - - self.axis = axis - self.filters = filters - self.kernel_size = kernel_size - self.strides = strides - self.padding = padding.upper() - self.dilation_rate = dilation_rate - self.activation = tf.keras.activations.get(activation) - self.use_bias = use_bias - self.kernel_initializer = tf.keras.initializers.get(kernel_initializer) - self.bias_initializer = tf.keras.initializers.get(bias_initializer) - self.kernel_regularizer = kernel_regularizer - - self.supports_masking = True - - def build(self, input_shape): - channel_axis = self.axis - input_dim = int(input_shape[channel_axis]) - kernel_shape = (self.kernel_size, input_dim, self.filters) - - # IMPORTANT: use variable_dtype so vars stay float32 under mixed_float16 - self.kernel = self.add_weight( - name="kernel", - shape=kernel_shape, - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - trainable=True, - dtype=self.variable_dtype, - ) - if self.use_bias: - self.bias = self.add_weight( - name="bias", - shape=(self.filters,), - initializer=self.bias_initializer, - trainable=True, - dtype=self.variable_dtype, - ) - else: - self.bias = None - - super().build(input_shape) - - def call(self, inputs, mask=None): - # Inputs arrive in compute_dtype under policy (e.g. float16) - # Be explicit: - inputs = tf.cast(inputs, self.compute_dtype) - - input_shape = tf.shape(inputs) - output_shape = self.compute_output_shape(input_shape) - - output_mask = None - if mask is not None: - # Broadcast mask, but do mask math in float32 for stability - mask = tf.cast(mask, tf.float32) - # apply mask on inputs (cast back to compute_dtype) - inputs = inputs * tf.cast(tf.expand_dims(mask, axis=-1), self.compute_dtype) - - # compute output mask (pure mask math in float32) - reshaped_mask = tf.reshape(mask, (-1, input_shape[2])) - mask_1d = tf.expand_dims(reshaped_mask, axis=-1) - mask_kernel = tf.ones((self.kernel_size, 1, 1), dtype=tf.float32) - - mask_conv = tf.nn.conv1d( - input=mask_1d, - filters=mask_kernel, - stride=self.strides, - padding=self.padding, - dilations=self.dilation_rate, - data_format="NWC", - ) - output_mask = tf.equal(mask_conv, float(self.kernel_size)) - output_mask = tf.squeeze(output_mask, axis=-1) - output_mask = tf.reshape(output_mask, shape=output_shape[:-1]) - - # conv1d on actual data - reshaped_inputs = tf.reshape(inputs, (-1, input_shape[2], input_shape[3])) - - # Cast kernel to compute dtype when mixing with activations - kernel = tf.cast(self.kernel, self.compute_dtype) - - outputs = tf.nn.conv1d( - input=reshaped_inputs, - filters=kernel, - stride=self.strides, - padding=self.padding, - dilations=self.dilation_rate, - data_format="NWC", - ) - if self.use_bias: - bias = tf.cast(self.bias, self.compute_dtype) - outputs = tf.nn.bias_add(outputs, bias, data_format="NWC") - if self.activation is not None: - outputs = self.activation(outputs) - - outputs = tf.reshape(outputs, output_shape) - - # stash mask so compute_mask() can forward it - self._output_mask = output_mask - return outputs - - def compute_mask(self, inputs, mask=None): - # never recompute, just return the cached mask - return getattr(self, "_output_mask", None) - - def get_config(self): - config = super().get_config() - config.update( - { - "filters": self.filters, - "kernel_size": self.kernel_size, - "strides": self.strides, - "padding": self.padding.lower(), - "dilation_rate": self.dilation_rate, - "activation": tf.keras.activations.serialize(self.activation), - "use_bias": self.use_bias, - "kernel_initializer": tf.keras.initializers.serialize( - self.kernel_initializer - ), - "bias_initializer": tf.keras.initializers.serialize( - self.bias_initializer - ), - "axis": self.axis, - "kernel_regularizer": tf.keras.regularizers.serialize( - self.kernel_regularizer - ) - if self.kernel_regularizer is not None - else None, - } - ) - return config - - def compute_output_shape(self, input_shape): - """ - compute the output shape only if the length dimension is not None. - This version is written to work with either TensorShape or tf.Tensor. - """ - length = input_shape[2] - if length is not None: - if self.padding == "SAME": - out_length = (length + self.strides - 1) // self.strides - elif self.padding == "VALID": - out_length = ( - length - self.dilation_rate * (self.kernel_size - 1) - 1 - ) // self.strides + 1 - else: - raise ValueError("Invalid padding type.") - return (input_shape[0], input_shape[1], out_length, self.filters) - else: - return (input_shape[0], input_shape[1], input_shape[2], self.filters) - - -# class ResidualBlock(tf.keras.layers.Layer): -# """The Residual block of ResNet models.""" - -# def __init__( -# self, -# use_1x1conv=False, -# block_number=1, -# activation="gelu", -# return_nmd=False, -# **kwargs, -# ): - -# super().__init__() -# self.supports_masking = True -# self.filters = kwargs.get('filters') -# self.padding = kwargs.get('padding', 'same').upper() -# self.strides = kwargs.get('strides', 1) -# self.block_number = block_number -# self.return_nmd = return_nmd -# self.conv1 = MaskedConv1D( -# padding=self.padding, -# name=f"masked_conv1d_blk{self.block_number}_1", -# **{k:v for k,v in kwargs.items() if k not in ["name", "padding"]} -# ) -# self.conv2 = MaskedConv1D( -# padding=self.padding, -# name=f"masked_conv1d_blk{self.block_number}_2", -# **{k:v for k,v in kwargs.items() if k not in ["name", "padding", "strides"]} -# ) -# self.conv3 = None -# if use_1x1conv or kwargs.get('strides') > 1: -# self.conv3 = MaskedConv1D( -# padding=self.padding, -# kernel_size=1, -# name=f"masked_conv1d_blk{self.block_number}_bypass", -# **{k:v for k,v in kwargs.items() if k not in ["kernel_size", "name", "padding"]} -# ) -# self.bn3 = MaskedBatchNorm(name=f"masked_batchnorm_blk{self.block_number}_bypass",) -# self.bn1 = MaskedBatchNorm(name=f"masked_batchnorm_blk{self.block_number}_1",) -# self.bn2 = MaskedBatchNorm(name=f"masked_batchnorm_blk{self.block_number}_2", return_nmd=return_nmd) -# self.add = MaskedAdd(name=f"resblock_add_blk{self.block_number}") -# self.activation = tf.keras.layers.Activation(activation, name=f"resblock_activation_blk{self.block_number}") - -# def call(self, inputs, mask=None): -# x = self.conv1(inputs, mask=mask) -# x = self.bn1(x) -# x = self.activation(x) -# x = self.conv2(x) -# if self.return_nmd: -# x, x_nmd = self.bn2(x) -# else: -# x = self.bn2(x) - -# if self.conv3 is not None: -# x = self.add([x, self.bn3(self.conv3(inputs))]) -# else: -# x = self.add([x, inputs]) -# if self.return_nmd: -# return self.activation(x), x_nmd -# return self.activation(x) - -# def compute_output_shape(self, input_shape): -# """ -# compute the output shape only if the length dimention is not None. -# """ -# length = input_shape[2] -# if length is not None: -# if self.padding == "SAME": -# out_length = (length + self.strides - 1) // self.strides -# elif self.padding == "VALID": -# out_length = ( -# length - self.dilation_rate * (self.kernel_size - 1) - 1 -# ) // self.strides + 1 -# else: -# raise ValueError("Invalid padding type.") -# out_length = out_length -# if self.return_nmd: -# return (input_shape[0], input_shape[1], out_length, self.filters), (input_shape[0], self.filters) -# return (input_shape[0], input_shape[1], out_length, self.filters) -# else: -# if self.return_nmd: -# return (input_shape[0], input_shape[1], input_shape[2], self.filters), (input_shape[0], self.filters) -# return (input_shape[0], input_shape[1], input_shape[2], self.filters) - - -# # Todo: implement output mask computation -> return the mask computed by the last layer? - - -class ResidualBlock(tf.keras.layers.Layer): - """The Residual block of ResNet models.""" - - def __init__( - self, - use_1x1conv=False, - block_number=1, - activation="gelu", - return_nmd=False, - **kwargs, - ): - # --- 1. Pull out args meant for the convs, so they don't go to super().__init__ --- - self.filters = kwargs.pop("filters", None) - self.kernel_size = kwargs.pop("kernel_size", 3) - self.strides = kwargs.pop("strides", 1) - self.padding = kwargs.pop("padding", "same").upper() - self.dilation_rate = kwargs.pop("dilation_rate", 1) - self.use_bias = kwargs.pop("use_bias", True) - self.kernel_regularizer = kwargs.pop("kernel_regularizer", None) - self.kernel_initializer = kwargs.pop("kernel_initializer", "glorot_uniform") - self.bias_initializer = kwargs.pop("bias_initializer", "zeros") - - # now kwargs only contains things Layer.__init__ understands (name, dtype, trainable, etc.) - super().__init__(**kwargs) - - self.supports_masking = True - self.block_number = block_number - self.return_nmd = return_nmd - - # --- 2. Build the internal conv/bn/add layers using the extracted values --- - - conv_common = dict( - filters=self.filters, - kernel_size=self.kernel_size, - strides=self.strides, - padding=self.padding, - dilation_rate=self.dilation_rate, - use_bias=self.use_bias, - kernel_initializer=self.kernel_initializer, - bias_initializer=self.bias_initializer, - kernel_regularizer=self.kernel_regularizer, - ) - - # first conv - self.conv1 = MaskedConv1D( - name=f"masked_conv1d_blk{self.block_number}_1", - **conv_common, - ) - - # second conv has stride = 1 - conv2_common = conv_common.copy() - conv2_common["strides"] = 1 - self.conv2 = MaskedConv1D( - name=f"masked_conv1d_blk{self.block_number}_2", - **conv2_common, - ) - - # optional 1x1 conv for the shortcut - self.conv3 = None - self.bn3 = None - if use_1x1conv or self.strides > 1: - bypass_common = conv_common.copy() - bypass_common["kernel_size"] = 1 - self.conv3 = MaskedConv1D( - name=f"masked_conv1d_blk{self.block_number}_bypass", - **bypass_common, - ) - self.bn3 = MaskedBatchNorm( - name=f"masked_batchnorm_blk{self.block_number}_bypass", - ) - - self.bn1 = MaskedBatchNorm( - name=f"masked_batchnorm_blk{self.block_number}_1", - ) - self.bn2 = MaskedBatchNorm( - name=f"masked_batchnorm_blk{self.block_number}_2", - return_nmd=return_nmd, - ) - - self.add = MaskedAdd(name=f"resblock_add_blk{self.block_number}") - self.activation_layer = tf.keras.layers.Activation( - activation, name=f"resblock_activation_blk{self.block_number}" - ) - - def call(self, inputs, mask=None, training=None): - x = self.conv1(inputs, mask=mask) - x = self.bn1(x, training=training) - x = self.activation_layer(x) - - x = self.conv2(x) - if self.return_nmd: - x, x_nmd = self.bn2(x, training=training) - else: - x = self.bn2(x) - - if self.conv3 is not None: - shortcut = self.conv3(inputs, mask=mask) - shortcut = self.bn3(shortcut, training=training) - else: - shortcut = inputs - - x = self.add([x, shortcut]) - x = self.activation_layer(x) - - if self.return_nmd: - return x, x_nmd - return x - - def compute_output_shape(self, input_shape): - length = input_shape[2] - if length is not None: - if self.padding == "SAME": - out_length = (length + self.strides - 1) // self.strides - elif self.padding == "VALID": - out_length = ( - length - self.dilation_rate * (self.kernel_size - 1) - 1 - ) // self.strides + 1 - else: - raise ValueError("Invalid padding type.") - if self.return_nmd: - return ( - (input_shape[0], input_shape[1], out_length, self.filters), - (input_shape[0], self.filters), - ) - return (input_shape[0], input_shape[1], out_length, self.filters) - else: - if self.return_nmd: - return ( - (input_shape[0], input_shape[1], input_shape[2], self.filters), - (input_shape[0], self.filters), - ) - return (input_shape[0], input_shape[1], input_shape[2], self.filters) - - def get_config(self): - config = super().get_config() - config.update( - { - "use_1x1conv": hasattr(self, "conv3") and self.conv3 is not None, - "block_number": self.block_number, - "activation": tf.keras.activations.serialize( - self.activation_layer.activation - ), - "return_nmd": self.return_nmd, - "filters": self.filters, - "kernel_size": self.kernel_size, - "strides": self.strides, - "padding": self.padding.lower(), - "dilation_rate": self.dilation_rate, - "use_bias": self.use_bias, - "kernel_initializer": tf.keras.initializers.serialize( - self.kernel_initializer - ), - "bias_initializer": tf.keras.initializers.serialize( - self.bias_initializer - ), - "kernel_regularizer": tf.keras.regularizers.serialize( - self.kernel_regularizer - ) - if self.kernel_regularizer is not None - else None, - } - ) - return config - - -# To do: implement a method to set epsilon considering the data type of the tensors passed to the layers -# if not implemented correctly, this can lead to overflow/underflow issues. - - -class MetricModel(tf.keras.Model): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.step = 0 - self.loss_tracker = tf.keras.metrics.Mean(name="loss") - self.regularization_loss_tracker = tf.keras.metrics.Mean(name="reg-loss") - self.gradient_tracker = tf.keras.metrics.Mean(name="gradient") - - def compile(self, optimizer, loss_fn, **kwargs): - super(MetricModel, self).compile(**kwargs) - self.loss_fn = loss_fn - self.optimizer = optimizer - - def train_step(self, data): - if len(data) == 3: - x, y, _ = data - else: - # sample _weighning has to be implemented - _ = None - x, y = data - - with tf.GradientTape() as tape: - y_pred = self(x, training=True) - # Adjust loss call based on your loss_fn signature: - loss = self.loss_fn([y, y_pred]) - # Add regularization loss - loss += sum(self.losses) - # If using mixed precision - if hasattr(self.optimizer, "get_scaled_loss"): - loss = self.optimizer.get_scaled_loss(loss) - - # Compute gradients - trainable_vars = self.trainable_variables - grads = tape.gradient(loss, trainable_vars) - # If using mixed precision - if hasattr(self.optimizer, "get_unscaled_gradients"): - grads = self.optimizer.get_unscaled_gradients(grads) - - # Apply gradients - self.optimizer.apply_gradients(zip(grads, trainable_vars)) - - # Update step and metrics - self.step += 1 - if self.step % 100 == 0: - # Optional: reset metrics every 100 steps - self.loss_tracker.reset_state() - self.gradient_tracker.reset_state() - - self.loss_tracker.update_state(loss) - self.regularization_loss_tracker.update_state(sum(self.losses)) - - # Compute average gradient norm - total_norm = 0.0 - total_params = 0.0 - for grad, weight in zip(grads, self.trainable_weights): - if grad is not None: - norm = tf.norm(grad) - total_norm += norm - total_params += tf.cast(tf.math.reduce_prod(weight.shape), tf.float32) - - avg_grad_norm = total_norm / tf.maximum(total_params, 1.0) - self.gradient_tracker.update_state(avg_grad_norm) - - return { - "loss": self.loss_tracker.result(), - "reg-loss": self.regularization_loss_tracker.result(), - "grad": self.gradient_tracker.result(), - "lr": self.optimizer.learning_rate, - } - - def test_step(self, data): - # No scaling needed in eval - if len(data) == 3: - x, y, _ = data - else: - x, y = data - - y_pred = self(x, training=False) - base_loss = self.loss_fn([y, y_pred]) - reg_loss = tf.add_n(self.losses) if self.losses else 0.0 - total_loss = base_loss + reg_loss - - self.loss_tracker.update_state(total_loss) - # (Usually people don’t track reg-loss in test, but you can if you want) - return {"loss": self.loss_tracker.result()} - - @property - def metrics(self): - # These are reset automatically at the start of each epoch - return [ - self.loss_tracker, - self.regularization_loss_tracker, - self.gradient_tracker, - ] - - -class SinusoidalPositionEmbedding(tf.keras.layers.Layer): - def __init__(self, max_wavelength=10000, **kwargs): - super().__init__(**kwargs) - self.max_wavelength = max_wavelength - self.supports_masking = True - - def call(self, inputs, start_index=0): - # inputs: tensor with shape [..., seq_length, hidden_size] - # Compute dynamic shapes - input_shape = tf.shape(inputs) - seq_length = input_shape[-2] - hidden_size = input_shape[-1] - - # Positions [0, 1, ..., seq_length-1] offset by start_index - positions = tf.cast( - tf.range(seq_length) + start_index, dtype=self.compute_dtype - ) - - # Minimum frequency as inverse of max_wavelength - min_freq = tf.cast(1.0 / self.max_wavelength, dtype=self.compute_dtype) - - # Compute timescales for each dimension: min_freq^(2i/hidden_size) - dim_indices = tf.cast(tf.range(hidden_size), dtype=self.compute_dtype) - # floor(dim_indices/2)*2 for pairing sin/cos - even_dims = tf.floor(dim_indices / 2) * 2 - timescales = tf.pow( - min_freq, even_dims / tf.cast(hidden_size, self.compute_dtype) - ) - - # Compute angles: outer product of positions and timescales - angles = tf.expand_dims(positions, -1) * tf.expand_dims(timescales, 0) - - # Build masks: even dims use sine, odd use cosine - sin_mask = tf.cast(tf.equal(dim_indices % 2, 0), self.compute_dtype) - cos_mask = 1.0 - sin_mask - - # Compute positional encodings: sin for even, cos for odd dims - pos_encoding = tf.sin(angles) * sin_mask + tf.cos(angles) * cos_mask - # pos_encoding shape [seq_length, hidden_size] - - # Broadcast to match input shape - broadcast_shape = tf.concat( - [input_shape[:-2], [seq_length, hidden_size]], axis=0 - ) - pos_encoding = tf.broadcast_to(pos_encoding, broadcast_shape) - - return pos_encoding - - def get_config(self): - config = super().get_config() - config.update({"max_wavelength": self.max_wavelength}) - return config - - def compute_output_shape(self, input_shape): - return input_shape - - -class TransformerEncoder(tf.keras.layers.Layer): - def __init__( - self, - embed_dim, - num_heads, - feed_forward_dim, - dropout_rate=0.1, - attention_axes=2, # For (batch, strand, length, feature), axis=2 is "length" - **kwargs, - ): - super().__init__(**kwargs) - self.embed_dim = embed_dim - self.num_heads = num_heads - self.feed_forward_dim = feed_forward_dim - self.dropout_rate = dropout_rate - self.attention_axes = attention_axes # Make sure axis is correct for your input - - # 1) Self-attention sublayer - self.attn_norm = tf.keras.layers.LayerNormalization( - epsilon=1e-6, name="attn_norm" - ) - self.mha = tf.keras.layers.MultiHeadAttention( - num_heads=num_heads, - key_dim=embed_dim // num_heads, - dropout=dropout_rate, - attention_axes=[self.attention_axes], - name="mha", - ) - self.attn_dropout = tf.keras.layers.Dropout(dropout_rate, name="attn_dropout") - - # 2) Feed-forward sublayer - self.ffn_norm = tf.keras.layers.LayerNormalization( - epsilon=1e-6, name="ffn_norm" - ) - self.ffn_dense1 = tf.keras.layers.Dense( - feed_forward_dim, activation="gelu", name="ffn_dense1" - ) - self.ffn_dropout1 = tf.keras.layers.Dropout(dropout_rate, name="ffn_dropout1") - self.ffn_dense2 = tf.keras.layers.Dense(embed_dim, name="ffn_dense2") - self.ffn_dropout2 = tf.keras.layers.Dropout(dropout_rate, name="ffn_dropout2") - - def call(self, inputs, mask=None, training=False, return_attention=False): - # --- Multi-Head Self-Attention + Residual - x_norm = self.attn_norm(inputs) - # Reshape/transpose if needed to ensure attention is over length - # For (batch, strand, length, feature), attention_axes=[2] is correct - attn_out, attn_weights = self.mha( - x_norm, x_norm, training=training, return_attention_scores=True - ) - attn_out = self.attn_dropout(attn_out, training=training) - x = inputs + attn_out - - # --- Feed-Forward Network + Residual - x_norm = self.ffn_norm(x) - ffn_out = self.ffn_dense1(x_norm) - ffn_out = self.ffn_dropout1(ffn_out, training=training) - ffn_out = self.ffn_dense2(ffn_out) - ffn_out = self.ffn_dropout2(ffn_out, training=training) - output = x + ffn_out - - if return_attention: - return output, attn_weights - return output - - def get_config(self): - cfg = super().get_config() - cfg.update( - { - "embed_dim": self.embed_dim, - "num_heads": self.num_heads, - "feed_forward_dim": self.feed_forward_dim, - "dropout_rate": self.dropout_rate, - } - ) - return cfg - - -class CrossFrameAttention(tf.keras.layers.Layer): - """ - Multi-Head Self-Attention across reading frames. - - For input shape (batch, frames, length, channels), this layer reshapes to - (batch * length, frames, channels) and applies self-attention across the - frame dimension. This allows each position in the sequence to attend to all - 6 reading frames simultaneously, enabling the model to learn relationships - like: "frame 1 and frame 4 are reverse complements", "frame 2 has a shift", - or "the correct ORF is frame 3". - - After attention, the tensor is reshaped back to (batch, frames, length, channels). - - Args: - embed_dim: Dimension of the input/output features (must be divisible by num_heads). - num_heads: Number of attention heads. - feed_forward_dim: Hidden dimension of the feed-forward network. - dropout_rate: Dropout rate for attention and FFN. - use_ffn: Whether to include the feed-forward sublayer (default True). - """ - - def __init__( - self, - embed_dim, - num_heads, - feed_forward_dim, - dropout_rate=0.1, - use_ffn=True, - **kwargs, - ): - super().__init__(**kwargs) - self.embed_dim = embed_dim - self.num_heads = num_heads - self.feed_forward_dim = feed_forward_dim - self.dropout_rate = dropout_rate - self.use_ffn = use_ffn - - # Self-attention sublayer — attention over frames (axis=1 after reshape) - self.attn_norm = tf.keras.layers.LayerNormalization( - epsilon=1e-6, name="attn_norm" - ) - self.mha = tf.keras.layers.MultiHeadAttention( - num_heads=num_heads, - key_dim=embed_dim // num_heads, - dropout=dropout_rate, - attention_axes=[1], # attend over frames after reshape - name="mha", - ) - self.attn_dropout = tf.keras.layers.Dropout(dropout_rate, name="attn_dropout") - - # Optional feed-forward sublayer - if use_ffn: - self.ffn_norm = tf.keras.layers.LayerNormalization( - epsilon=1e-6, name="ffn_norm" - ) - self.ffn_dense1 = tf.keras.layers.Dense( - feed_forward_dim, activation="gelu", name="ffn_dense1" - ) - self.ffn_dropout1 = tf.keras.layers.Dropout( - dropout_rate, name="ffn_dropout1" - ) - self.ffn_dense2 = tf.keras.layers.Dense(embed_dim, name="ffn_dense2") - self.ffn_dropout2 = tf.keras.layers.Dropout( - dropout_rate, name="ffn_dropout2" - ) - - def call(self, inputs, mask=None, training=False, return_attention=False): - # inputs: (batch, frames, length, channels) - batch_size = tf.shape(inputs)[0] - num_frames = tf.shape(inputs)[1] - seq_len = tf.shape(inputs)[2] - channels = tf.shape(inputs)[3] - - # Reshape to (batch * length, frames, channels) for frame-wise attention - # Each sequence position attends across all 6 frames - x = tf.transpose(inputs, [0, 2, 1, 3]) # (B, L, F, C) - x = tf.reshape(x, [-1, num_frames, channels]) # (B*L, F, C) - - # --- Multi-Head Self-Attention + Residual - x_norm = self.attn_norm(x) - attn_out, attn_weights = self.mha( - x_norm, x_norm, training=training, return_attention_scores=True - ) - attn_out = self.attn_dropout(attn_out, training=training) - x = x + attn_out # residual over frames - - # --- Optional Feed-Forward Network + Residual - if self.use_ffn: - x_norm = self.ffn_norm(x) - ffn_out = self.ffn_dense1(x_norm) - ffn_out = self.ffn_dropout1(ffn_out, training=training) - ffn_out = self.ffn_dense2(ffn_out) - ffn_out = self.ffn_dropout2(ffn_out, training=training) - x = x + ffn_out - - # Reshape back to (batch, frames, length, channels) - x = tf.reshape(x, [batch_size, seq_len, num_frames, channels]) - x = tf.transpose(x, [0, 2, 1, 3]) # (B, F, L, C) - - if return_attention: - return x, attn_weights - return x - - def get_config(self): - cfg = super().get_config() - cfg.update( - { - "embed_dim": self.embed_dim, - "num_heads": self.num_heads, - "feed_forward_dim": self.feed_forward_dim, - "dropout_rate": self.dropout_rate, - "use_ffn": self.use_ffn, - } - ) - return cfg - - -class AxialAttention(tf.keras.layers.Layer): - """ - Axial attention for 4D sequence tensors (batch, frames, length, channels). - - Applies attention alternately along the length axis (intra-frame) and the - frame axis (cross-frame). This captures both local sequence patterns within - each frame and global relationships across frames. - - Args: - embed_dim: Feature dimension. - num_heads: Number of attention heads. - feed_forward_dim: FFN hidden dimension. - dropout_rate: Dropout rate. - num_blocks: Number of (length-attn + frame-attn) blocks to stack. - """ - - def __init__( - self, - embed_dim, - num_heads, - feed_forward_dim, - dropout_rate=0.1, - num_blocks=1, - **kwargs, - ): - super().__init__(**kwargs) - self.embed_dim = embed_dim - self.num_heads = num_heads - self.feed_forward_dim = feed_forward_dim - self.dropout_rate = dropout_rate - self.num_blocks = num_blocks - - self.length_attns = [] - self.frame_attns = [] - - for i in range(num_blocks): - # Attention over length (intra-frame) — uses existing TransformerEncoder logic - self.length_attns.append( - TransformerEncoder( - embed_dim=embed_dim, - num_heads=num_heads, - feed_forward_dim=feed_forward_dim, - dropout_rate=dropout_rate, - attention_axes=2, # length axis in (B, F, L, C) - name=f"length_attn_{i}", - ) - ) - # Attention over frames (cross-frame) - self.frame_attns.append( - CrossFrameAttention( - embed_dim=embed_dim, - num_heads=num_heads, - feed_forward_dim=feed_forward_dim, - dropout_rate=dropout_rate, - use_ffn=True, - name=f"frame_attn_{i}", - ) - ) - - def call(self, inputs, training=False): - x = inputs - for length_attn, frame_attn in zip(self.length_attns, self.frame_attns): - x = length_attn(x, training=training) - x = frame_attn(x, training=training) - return x - - def get_config(self): - cfg = super().get_config() - cfg.update( - { - "embed_dim": self.embed_dim, - "num_heads": self.num_heads, - "feed_forward_dim": self.feed_forward_dim, - "dropout_rate": self.dropout_rate, - "num_blocks": self.num_blocks, - } - ) - return cfg - - -def ResidualBlock_wrapper(block_size: int, in_shape: tuple, **kwargs): - """ - Build a sequential stack of ResidualBlock layers as a Keras functional submodel. - If return_nmd=True, the final block outputs both (x, nmd). - """ - name = kwargs.get("name", "resblock") - return_nmd = kwargs.get("return_nmd", False) - inputs = tf.keras.Input(shape=in_shape) - - x = inputs - nmd = None - - for i in range(block_size): - # Skip certain kwargs for intermediate blocks - omit_keys = {"name"} - if i != 0 and "use_1x1conv" in kwargs: - omit_keys.add("use_1x1conv") - - block_kwargs = {k: v for k, v in kwargs.items() if k not in omit_keys} - - block_name = f"{name}_{i}" - block_number = f"{name.split('_')[-1]}{i}" - - block = ResidualBlock( - block_number=block_number, name=block_name, **block_kwargs - ) - - if return_nmd: - x, nmd = block(x) - else: - x = block(x) - - outputs = [x, nmd] if return_nmd else x - return tf.keras.Model(inputs=inputs, outputs=outputs, name=name) diff --git a/src/jaeger/nnlib/v2/losses.py b/src/jaeger/nnlib/v2/losses.py deleted file mode 100644 index 5c31f15..0000000 --- a/src/jaeger/nnlib/v2/losses.py +++ /dev/null @@ -1,157 +0,0 @@ -import tensorflow as tf - - -class SupervisedContrastiveLoss(tf.keras.losses.Loss): - def __init__(self, temperature=1, name=None): - super().__init__(name=name) - self.temperature = temperature - - def __call__(self, labels, feature_vectors, sample_weight=None): - # Normalize feature vectors - # Handle both sparse (B,) and one-hot (B, C) labels - if len(labels.shape) > 1: - labels = tf.math.argmax(labels, axis=-1) - feature_vectors_normalized = tf.math.l2_normalize(feature_vectors, axis=1) - # Compute logits - logits = tf.divide( - tf.matmul( - feature_vectors_normalized, tf.transpose(feature_vectors_normalized) - ), - self.temperature, - ) - return npairs_loss(tf.squeeze(labels), logits) - - -def npairs_loss(y_true, y_pred): - y_pred = tf.convert_to_tensor(y_pred) - y_true = tf.cast(y_true, y_pred.dtype) - - # Expand to [batch_size, 1] - y_true = tf.expand_dims(y_true, -1) - y_true = tf.cast(tf.equal(y_true, tf.transpose(y_true)), y_pred.dtype) - y_true /= tf.math.reduce_sum(y_true, 1, keepdims=True) - - loss = tf.nn.softmax_cross_entropy_with_logits(logits=y_pred, labels=y_true) - - return tf.math.reduce_mean(loss) - - -class ArcFaceLoss(tf.keras.layers.Layer): - def __init__( - self, num_classes, embedding_dim, margin=0.5, scale=30.0, onehot=True, **kwargs - ): - super(ArcFaceLoss, self).__init__(**kwargs) - self.num_classes = num_classes - self.margin = margin - self.scale = scale - self.embedding_dim = embedding_dim - self.onehot = onehot - - # Class weights are usually kept in float32 even in mixed precision - self.class_weights = self.add_weight( - name="class_weights", - shape=(self.num_classes, self.embedding_dim), - initializer="glorot_uniform", - trainable=True, - dtype=tf.float32, - ) - - def call(self, labels, embeddings): - """ - :param embeddings: (batch_size, embedding_dim) - :param labels: either: - - one-hot labels (batch_size, num_classes) if self.onehot=True - - integer labels (batch_size,) or (batch_size, 1) if self.onehot=False - """ - # Work in the compute dtype (usually float16 in mixed precision) - compute_dtype = embeddings.dtype - - # eps depends on dtype - eps = tf.constant( - 6.55e-4 if compute_dtype == tf.float16 else 1.0e-9, - dtype=compute_dtype, - ) - - # Normalize embeddings and class weights in compute dtype - embeddings = tf.cast(embeddings, compute_dtype) - class_weights = tf.cast(self.class_weights, compute_dtype) - - embeddings = tf.nn.l2_normalize(embeddings, axis=1) - class_weights = tf.nn.l2_normalize(class_weights, axis=1) - - # Cosine similarity: (batch_size, num_classes) - cosine = tf.matmul(embeddings, class_weights, transpose_b=True) - - # Labels -> one-hot - if self.onehot: - # Assume labels already one-hot, just cast - labels_one_hot = tf.cast(labels, compute_dtype) - else: - labels = tf.reshape(labels, [-1]) # (batch_size,) - labels = tf.cast(labels, tf.int32) - labels_one_hot = tf.one_hot( - labels, depth=self.num_classes, dtype=compute_dtype - ) - - # Angle and margin - theta = tf.acos(tf.clip_by_value(cosine, -1.0 + eps, 1.0 - eps)) - target_logits = tf.cos(theta + self.margin) - - # Construct logits - logits = cosine * (1.0 - labels_one_hot) + target_logits * labels_one_hot - - # Apply scaling - logits = logits * self.scale - - # ---- IMPORTANT PART FOR MIXED PRECISION ---- - # Do the actual loss calculation in float32 - logits_fp32 = tf.cast(logits, tf.float32) - labels_fp32 = tf.cast(labels_one_hot, tf.float32) - - loss_vec = tf.nn.softmax_cross_entropy_with_logits( - labels=labels_fp32, logits=logits_fp32 - ) - loss = tf.reduce_mean(loss_vec) - - # Always return float32 loss - return loss - - -class HierarchicalLoss(tf.keras.losses.Loss): - def __init__(self, parent_of, groups, l_fine=1.0, l_coarse=1.5, name="hier_loss"): - super().__init__( - name=name, reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE - ) - self.parent_of = tf.constant(parent_of, tf.int32) # (6,) - self.groups = [tf.constant(g, tf.int32) for g in groups] # list of tensors - self.l_fine = float(l_fine) - self.l_coarse = float(l_coarse) - self.ce = tf.keras.losses.SparseCategoricalCrossentropy( - from_logits=True, reduction="none" - ) - - def call(self, y_true, fine_logits): - # accept sparse ids (N,) or one-hot (N,6) - y_true = tf.convert_to_tensor(y_true) - if y_true.shape.rank == 2: - y_true = tf.argmax(y_true, axis=-1, output_type=tf.int32) - else: - y_true = tf.cast(tf.reshape(y_true, [-1]), tf.int32) - - # Fine-level CE (per-sample) - loss_fine = self.ce(y_true, fine_logits) # (N,) - - # Coarse logits via logsumexp over each group's fine logits - coarse_logits = tf.stack( - [ - tf.reduce_logsumexp(tf.gather(fine_logits, idxs, axis=1), axis=1) - for idxs in self.groups - ], - axis=1, - ) # (N, 3) - - y_coarse = tf.gather(self.parent_of, y_true) # (N,) - loss_coarse = self.ce(y_coarse, coarse_logits) # (N,) - - per_ex = self.l_fine * loss_fine + self.l_coarse * loss_coarse # (N,) - return tf.reduce_mean(per_ex) # scalar diff --git a/src/jaeger/nnlib/v2/maps.py b/src/jaeger/nnlib/v2/maps.py deleted file mode 100644 index 5107084..0000000 --- a/src/jaeger/nnlib/v2/maps.py +++ /dev/null @@ -1,539 +0,0 @@ -CODONS = [ - "TTT", - "TTC", - "TTA", - "TTG", - "CTT", - "CTC", - "CTA", - "CTG", - "ATT", - "ATC", - "ATA", - "ATG", - "GTT", - "GTC", - "GTA", - "GTG", - "TCT", - "TCC", - "TCA", - "TCG", - "CCT", - "CCC", - "CCA", - "CCG", - "ACT", - "ACC", - "ACA", - "ACG", - "GCT", - "GCC", - "GCA", - "GCG", - "TAT", - "TAC", - "TAA", - "TAG", - "CAT", - "CAC", - "CAA", - "CAG", - "AAT", - "AAC", - "AAA", - "AAG", - "GAT", - "GAC", - "GAA", - "GAG", - "TGT", - "TGC", - "TGA", - "TGG", - "CGT", - "CGC", - "CGA", - "CGG", - "AGT", - "AGC", - "AGA", - "AGG", - "GGT", - "GGC", - "GGA", - "GGG", -] - -AMINO_ACIDS = [ - "F", - "F", - "L", - "L", - "L", - "L", - "L", - "L", - "I", - "I", - "I", - "M", - "V", - "V", - "V", - "V", - "S", - "S", - "S", - "S", - "P", - "P", - "P", - "P", - "T", - "T", - "T", - "T", - "A", - "A", - "A", - "A", - "Y", - "Y", - "*", - "*", - "H", - "H", - "Q", - "Q", - "N", - "N", - "K", - "K", - "D", - "D", - "E", - "E", - "C", - "C", - "*", - "W", - "R", - "R", - "R", - "R", - "S", - "S", - "R", - "R", - "G", - "G", - "G", - "G", -] - - -PC2 = { - "I": "A", - "V": "A", - "L": "A", - "F": "A", - "Y": "A", - "W": "A", - "H": "B", - "K": "B", - "R": "B", - "D": "B", - "E": "B", - "G": "A", - "A": "A", - "C": "A", - "S": "A", - "T": "A", - "M": "A", - "Q": "B", - "N": "B", - "P": "A", -} - -PC2_NUM = [ - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 0, - 0, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 1, - 1, - 0, - 1, - 2, - 2, - 2, - 2, - 1, - 1, - 2, - 2, - 1, - 1, - 1, - 1, -] - -MURPY10 = { - "A": "A", - "C": "C", - "G": "G", - "H": "H", - "P": "P", - "L": "L", - "V": "L", - "I": "L", - "M": "L", - "S": "S", - "T": "S", - "F": "F", - "Y": "F", - "W": "F", - "E": "E", - "D": "E", - "N": "E", - "Q": "E", - "K": "K", - "R": "K", -} - -PC5 = { - "I": "A", - "V": "A", - "L": "A", - "F": "R", - "Y": "R", - "W": "R", - "H": "R", - "K": "C", - "R": "C", - "D": "C", - "E": "C", - "G": "T", - "A": "T", - "C": "T", - "S": "T", - "T": "D", - "M": "D", - "Q": "D", - "N": "D", - "P": "D", -} - -AA_NUM = [ - 1, - 1, - 2, - 2, - 2, - 2, - 2, - 2, - 3, - 3, - 3, - 4, - 5, - 5, - 5, - 5, - 6, - 6, - 6, - 6, - 7, - 7, - 7, - 7, - 8, - 8, - 8, - 8, - 9, - 9, - 9, - 9, - 10, - 10, - 0, - 0, - 11, - 11, - 12, - 12, - 13, - 13, - 14, - 14, - 15, - 15, - 16, - 16, - 17, - 17, - 0, - 18, - 19, - 19, - 19, - 19, - 6, - 6, - 19, - 19, - 20, - 20, - 20, - 20, -] -CODONS_SYNONYMOUS_NUM = [ - 1, - 2, - 1, - 2, - 3, - 4, - 5, - 6, - 1, - 2, - 3, - 1, - 1, - 2, - 3, - 4, - 1, - 2, - 3, - 4, - 1, - 2, - 3, - 4, - 1, - 2, - 3, - 4, - 1, - 2, - 3, - 4, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 1, - 2, - 3, - 4, - 5, - 6, - 5, - 6, - 1, - 2, - 3, - 4, -] - -MURPY10_NUM = [ - 1, - 1, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 3, - 3, - 3, - 3, - 4, - 4, - 4, - 4, - 3, - 3, - 3, - 3, - 5, - 5, - 5, - 5, - 1, - 1, - 0, - 0, - 6, - 6, - 7, - 7, - 7, - 7, - 8, - 8, - 7, - 7, - 7, - 7, - 9, - 9, - 0, - 1, - 8, - 8, - 8, - 8, - 3, - 3, - 8, - 8, - 10, - 10, - 10, - 10, -] - -PC5_NUM = [ - 1, - 1, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 2, - 3, - 2, - 2, - 2, - 2, - 4, - 4, - 4, - 4, - 3, - 3, - 3, - 3, - 3, - 3, - 3, - 3, - 4, - 4, - 4, - 4, - 1, - 1, - 0, - 0, - 1, - 1, - 3, - 3, - 3, - 3, - 5, - 5, - 5, - 5, - 5, - 5, - 4, - 4, - 0, - 1, - 5, - 5, - 5, - 5, - 4, - 4, - 5, - 5, - 4, - 4, - 4, - 4, -] - -CODONS_MUM = [i for i in range(0, len(CODONS))] diff --git a/src/jaeger/preprocess/v1/convert.py b/src/jaeger/preprocess/v1/convert.py deleted file mode 100644 index dec8174..0000000 --- a/src/jaeger/preprocess/v1/convert.py +++ /dev/null @@ -1,222 +0,0 @@ -import logging -import tensorflow as tf -from jaeger.preprocess.v1.maps import TRIMERS, TRIMER_INT, AMINO_ACIDS, AMINO_ACIDS_INT - -logger = logging.getLogger("Jaeger") - - -def codon_mapper(): - """ - Creates a static hash table for mapping codons to their corresponding - values. - - Returns: - ------- - tf.lookup.StaticHashTable: A static hash table for codon mapping. - """ - - trimers = tf.constant(TRIMERS) - trimer_vals = tf.constant(TRIMER_INT) - trimer_init = tf.lookup.KeyValueTensorInitializer(trimers, trimer_vals) - return tf.lookup.StaticHashTable(trimer_init, default_value=0) - - -def amino_mapper(): - """ - Creates a static hash table for mapping amino acids to their corresponding - numerical values. - - Returns: - ------- - tf.lookup.StaticHashTable: A static hash table for amino acid mapping. - """ - - aa = tf.constant(AMINO_ACIDS) - aa_num = tf.constant(AMINO_ACIDS_INT) - aa_init = tf.lookup.KeyValueTensorInitializer(aa, aa_num) - return tf.lookup.StaticHashTable(aa_init, default_value=0) - - -def complement_mapper(): - """ - Creates a static hash table for mapping DNA nucleotides to their reverse - complements. - - Returns: - tf.lookup.StaticHashTable: A static hash table for reverse complement - mapping. - """ - - rc_keys = tf.constant([b"A", b"T", b"G", b"C", b"a", b"t", b"g", b"c"]) - rc_vals = tf.constant([b"T", b"A", b"C", b"G", b"t", b"a", b"c", b"g"]) - rc_init = tf.lookup.KeyValueTensorInitializer(rc_keys, rc_vals) - return tf.lookup.StaticHashTable(rc_init, default_value="N") - - -def process_string(onehot=True, label_onehot=True, crop_size=2048): - """ - Processes a DNA sequence string by mapping codons and nucleotides. - - Args: - ---- - onehot (bool, optional): Flag to indicate one-hot encoding. Defaults - to True. - label_onehot (bool, optional): Flag to indicate one-hot encoding for - labels. Defaults to True. - crop_size (int, optional): Size for cropping the sequence. Defaults - to 2048. - - Returns: - ------- - function: A function that processes a DNA sequence string and returns - mapped codons and nucleotides. - """ - - @tf.function(jit_compile=False) - def p(string): - t1, t3 = codon_mapper(), complement_mapper() - x = tf.strings.split(string, sep=",") - - if (crop_size % 3) == 0: - offset = -2 - elif (crop_size % 3) == 1: - offset = -1 - elif (crop_size % 3) == 2: - offset = 0 - - forward_strand = tf.strings.bytes_split(x[0]) # split the string - reverse_strand = t3.lookup(forward_strand[::-1]) - - tri_forward = tf.strings.ngrams(forward_strand, ngram_width=3, separator="") - tri_reverse = tf.strings.ngrams(reverse_strand, ngram_width=3, separator="") - - f1 = t1.lookup(tri_forward[: -3 + offset : 3]) - f2 = t1.lookup(tri_forward[1 : -2 + offset : 3]) - f3 = t1.lookup(tri_forward[2 : -1 + offset : 3]) - - r1 = t1.lookup(tri_reverse[: -3 + offset : 3]) - r2 = t1.lookup(tri_reverse[1 : -2 + offset : 3]) - r3 = t1.lookup(tri_reverse[2 : -1 + offset : 3]) - - return ( - { - "forward_1": tf.cast(f1, dtype=tf.float32), - "forward_2": tf.cast(f2, dtype=tf.float32), - "forward_3": tf.cast(f3, dtype=tf.float32), - "reverse_1": tf.cast(r1, dtype=tf.float32), - "reverse_2": tf.cast(r2, dtype=tf.float32), - "reverse_3": tf.cast(r3, dtype=tf.float32), - }, - # metadata routine does not work anymore - # convert the models to new format - # - x[1], - x[2], - x[3], - x[4], - x[5], - x[6], - x[7], - x[8], - x[9], - x[10], - ) - - return p - - -def process_string_textline_protein(label_onehot=True, numclasses=4): - """ - Processes a protein sequence text line by mapping amino acids and labels. - - Args: - ---- - label_onehot (bool, optional): Flag to indicate one-hot encoding for - labels. Defaults to True. - numclasses (int, optional): Number of classes for label encoding. - Defaults to 4. - - Returns: - ------- - function: A function that processes a protein sequence text line and - returns mapped amino acids and labels. - """ - - def p(string): - t1 = amino_mapper() - - x = tf.strings.split(string, sep=",") - - label = tf.strings.to_number(x[0], tf.int32) - label = tf.cast(label, dtype=tf.int32) - # split the string - prot_strand = tf.strings.bytes_split(x[1]) - - protein = t1.lookup(prot_strand) - - if label_onehot: - label = tf.one_hot( - label, depth=numclasses, dtype=tf.float32, on_value=1, off_value=0 - ) - - return protein, label - - return p - - -def process_string_textline(onehot=True, label_onehot=True, numclasses=4): - """ - Processes a DNA sequence text line by mapping codons and nucleotides. - - Args: - ---- - onehot (bool, optional): Flag to indicate one-hot encoding. - Defaults to True. - label_onehot (bool, optional): Flag to indicate one-hot encoding - for labels. Defaults to True. - numclasses (int, optional): Number of classes for label encoding. - Defaults to 4. - - Returns: - ------- - function: A function that processes a DNA sequence text line and - returns mapped codons, nucleotides, and labels. - """ - - def p(string): - t1, t3 = codon_mapper(), complement_mapper() - - x = tf.strings.split(string, sep=",") - - label = tf.strings.to_number(x[0], tf.int32) - label = tf.cast(label, dtype=tf.int32) - - forward_strand = tf.strings.bytes_split(x[1]) - reverse_strand = t3.lookup(forward_strand[::-1]) - - tri_forward = tf.strings.ngrams(forward_strand, ngram_width=3, separator="") - tri_reverse = tf.strings.ngrams(reverse_strand, ngram_width=3, separator="") - - f1 = t1.lookup(tri_forward[::3]) - f2 = t1.lookup(tri_forward[1::3]) - f3 = t1.lookup(tri_forward[2::3]) - - r1 = t1.lookup(tri_reverse[::3]) - r2 = t1.lookup(tri_reverse[1::3]) - r3 = t1.lookup(tri_reverse[2::3]) - - if label_onehot: - label = tf.one_hot( - label, depth=numclasses, dtype=tf.float32, on_value=1, off_value=0 - ) - - return { - "forward_1": f1, - "forward_2": f2, - "forward_3": f3, - "reverse_1": r1, - "reverse_2": r2, - "reverse_3": r3, - }, label - - return p diff --git a/src/jaeger/preprocess/v2/convert.py b/src/jaeger/preprocess/v2/convert.py deleted file mode 100644 index 9c74691..0000000 --- a/src/jaeger/preprocess/v2/convert.py +++ /dev/null @@ -1,178 +0,0 @@ -from typing import Any -import tensorflow as tf -from jaeger.preprocess.v2.maps import CODONS, MURPHY10_INT - - -def codon_mapper(mapto: list) -> Any: - """ - Creates a static hash table for mapping codons to their corresponding - values. - - Returns: - ------- - tf.lookup.StaticHashTable: A static hash table for codon mapping. - """ - - trimers = tf.constant(CODONS) - trimer_vals = tf.constant(mapto) - trimer_init = tf.lookup.KeyValueTensorInitializer(trimers, trimer_vals) - return tf.lookup.StaticHashTable(trimer_init, default_value=0) - - -def complement_mapper(): - """ - Creates a static hash table for mapping nucleotides to their complements. - - Returns: - ------- - tf.lookup.StaticHashTable: A static hash table for nucleotide - complement mapping. - """ - - rc_keys = tf.constant([b"A", b"T", b"G", b"C", b"a", b"t", b"g", b"c"]) - rc_vals = tf.constant([b"T", b"A", b"C", b"G", b"t", b"a", b"c", b"g"]) - rc_init = tf.lookup.KeyValueTensorInitializer(rc_keys, rc_vals) - return tf.lookup.StaticHashTable(rc_init, default_value="N") - - -def nuc_enc_mapper(): - """ - Creates a static hash table for mapping nucleotides to their numerical - encodings. - - Returns: - ------- - tf.lookup.StaticHashTable: A static hash table for nucleotide encoding - mapping. - """ - - keys_tensor = tf.constant([b"A", b"G", b"C", b"T", b"a", b"g", b"c", b"t"]) - vals_tensor = tf.constant([0, 1, 2, 3, 0, 1, 2, 3]) - ini = tf.lookup.KeyValueTensorInitializer(keys_tensor, vals_tensor) - return tf.lookup.StaticHashTable(ini, default_value=-1) - - -def alt_nuc_enc_mapper(): - """ - Creates a static hash table for mapping nucleotides to their alternative - numerical encodings. - - Returns: - tf.lookup.StaticHashTable: A static hash table for alternative - nucleotide encoding mapping. - """ - - keys_tensor = tf.constant([b"A", b"G", b"C", b"T", b"a", b"g", b"c", b"t"]) - vals_tensor = tf.constant([0, 1, 1, 0, 0, 1, 1, 0]) - ini = tf.lookup.KeyValueTensorInitializer(keys_tensor, vals_tensor) - return tf.lookup.StaticHashTable(ini, default_value=-1) - - -def process_string( - crop_size: int = 1024, - timesteps: bool = False, - num_time: int = None, - fragsize: int = 200, -): - """ - Processes a DNA sequence string by mapping codons, nucleotides, and codon - biases. - - Args: - ---- - onehot (bool, optional): Flag to indicate one-hot encoding. Defaults - to True. - crop_size (int, optional): Size for cropping the sequence. Defaults - to 1024. - maxval (int, optional): Maximum value for mutation. Defaults to 400. - timesteps (bool, optional): Flag to reshape output for time steps. - Defaults to False. - num_time (int, optional): Number of time steps. Defaults to None. - fragsize (int, optional): Size of the DNA sequence fragments. Defaults - to 200. - mutate (bool, optional): Flag to enable mutation. Defaults to False. - mutation_rate (float, optional): Probability of mutation. Defaults - to 0.1. - - Returns: - ------- - function: A function that processes a DNA sequence string and returns - mapped codons, nucleotides, and codon biases. - """ - - @tf.function - def p(string): - t1 = codon_mapper(MURPHY10_INT) - t3 = complement_mapper() - # t4 = nuc_enc_mapper() - - x = tf.strings.split(string, sep=",") - - if (crop_size % 3) == 0: - offset = -2 - elif (crop_size % 3) == 1: - offset = -1 - elif (crop_size % 3) == 2: - offset = 0 - - forward_strand = tf.strings.bytes_split(x[0])[:crop_size] - - # generate the reverse strand for the mutated forward strand - reverse_strand = t3.lookup(forward_strand[::-1]) - - # nuc1 = t4.lookup(forward_strand[:]) - # nuc2 = t4.lookup(reverse_strand[:]) - - tri_forward = tf.strings.ngrams(forward_strand, ngram_width=3, separator="") - tri_reverse = tf.strings.ngrams(reverse_strand, ngram_width=3, separator="") - - f1 = t1.lookup(tri_forward[: -3 + offset : 3]) - f2 = t1.lookup(tri_forward[1 : -2 + offset : 3]) - f3 = t1.lookup(tri_forward[2 : -1 + offset : 3]) - - # fb1=t5.lookup(tri_forward[0:-3+offset:3]) - # fb2=t5.lookup(tri_forward[1:-2+offset:3]) - # fb3=t5.lookup(tri_forward[2:-1+offset:3]) - - r1 = t1.lookup(tri_reverse[: -3 + offset : 3]) - r2 = t1.lookup(tri_reverse[1 : -2 + offset : 3]) - r3 = t1.lookup(tri_reverse[2 : -1 + offset : 3]) - - # rb1=t5.lookup(tri_reverse[0:-3+offset:3]) - # rb2=t5.lookup(tri_reverse[1:-2+offset:3]) - # rb3=t5.lookup(tri_reverse[2:-1+offset:3]) - - if timesteps: - f1 = tf.reshape(f1, (num_time, fragsize)) - f2 = tf.reshape(f2, (num_time, fragsize)) - f3 = tf.reshape(f3, (num_time, fragsize)) - r1 = tf.reshape(r1, (num_time, fragsize)) - r2 = tf.reshape(r2, (num_time, fragsize)) - r3 = tf.reshape(r3, (num_time, fragsize)) - seq = tf.stack([f1, f2, f3, r1, r2, r3], 1) - else: - seq = tf.stack([f1, f2, f3, r1, r2, r3], 0) - # nuc = tf.stack([nuc1,nuc2],0) - # code = tf.stack([fb1,fb2,fb3,rb1,rb2,rb3],0) # codon bias encoder - - return ( - { - "translated": tf.one_hot( - seq, depth=11, dtype=tf.float32, on_value=1, off_value=0 - ) - }, - x[1], - x[2], - x[3], - x[4], - x[5], - x[6], - x[7], - x[8], - x[9], - x[10], - ) - # 'nucleotide': tf.one_hot(nuc, depth=4, dtype=tf.float32, on_value=1,\ - # off_value=0)}, - - return p diff --git a/src/jaeger/seqops/encode.py b/src/jaeger/seqops/encode.py deleted file mode 100644 index e4b7e09..0000000 --- a/src/jaeger/seqops/encode.py +++ /dev/null @@ -1,521 +0,0 @@ -"""Sequence encoding utilities. - -TensorFlow-based DNA sequence preprocessing for training and inference. -Supports codon translation, nucleotide one-hot encoding, and complement mapping. -""" - -from __future__ import annotations - -import numpy as np -import tensorflow as tf - -from jaeger.seqops.maps import CODON_ID, CODONS - - -# ------------------------------------------------------------------ -# Lookup table builders -# ------------------------------------------------------------------ - - -def _map_codon(codons, codon_num): - """Build a TF hash table mapping codon strings to integer IDs.""" - trimers = tf.constant(codons) - trimer_vals = tf.constant(codon_num) - trimer_init = tf.lookup.KeyValueTensorInitializer(trimers, trimer_vals) - return tf.lookup.StaticHashTable(trimer_init, default_value=-1) - - -def _map_complement(): - """Build a TF hash table for reverse-complement nucleotides.""" - rc_keys = tf.constant([b"A", b"T", b"G", b"C", b"a", b"t", b"g", b"c"]) - rc_vals = tf.constant([b"T", b"A", b"C", b"G", b"t", b"a", b"c", b"g"]) - rc_init = tf.lookup.KeyValueTensorInitializer(rc_keys, rc_vals) - return tf.lookup.StaticHashTable(rc_init, default_value="N") - - -def _map_nucleotide(): - """Build a TF hash table mapping nucleotides to integer IDs (A=0, G=1, C=2, T=3).""" - keys_tensor = tf.constant([b"A", b"G", b"C", b"T", b"a", b"g", b"c", b"t"]) - vals_tensor = tf.constant([0, 1, 2, 3, 0, 1, 2, 3]) - ini = tf.lookup.KeyValueTensorInitializer(keys_tensor, vals_tensor) - return tf.lookup.StaticHashTable(ini, default_value=-1) - - -def _map_nucleotide_type(): - """Map nucleotides to purine (0) / pyrimidine (1).""" - keys_tensor = tf.constant([b"A", b"G", b"C", b"T"]) - vals_tensor = tf.constant([0, 1, 1, 0]) - ini = tf.lookup.KeyValueTensorInitializer(keys_tensor, vals_tensor) - return tf.lookup.StaticHashTable(ini, default_value=-1) - - -def _remap_labels(original, alternative): - """Build a TF hash table remapping label IDs to alternative IDs.""" - keys_tensor = tf.constant(original) - vals_tensor = tf.constant(alternative) - ini = tf.lookup.KeyValueTensorInitializer(keys_tensor, vals_tensor) - return tf.lookup.StaticHashTable(ini, default_value=0) - - -# ------------------------------------------------------------------ -# Training preprocessor -# ------------------------------------------------------------------ - - -def process_string_train( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=None, - label_original=None, - label_alternative=None, - class_label_onehot=True, - seq_onehot=True, - num_classes=None, - crop_size=None, - timesteps=False, - num_time=None, - fragsize=200, - mutate=False, - mutation_rate=0.1, - masking=False, - ngram_width=None, - input_type="translated", - shuffle=False, - shuffle_frames=False, -): - """Build a TensorFlow function for training-time sequence preprocessing. - - Returns a callable that takes a CSV string and returns ``(inputs, label)``. - """ - if codons and codon_num: - map_codon = _map_codon(codons=codons, codon_num=codon_num) - map_complement = _map_complement() - map_nucleotide = _map_nucleotide() - - if label_original is not None and label_alternative is not None: - remap_labels = _remap_labels( - original=label_original, alternative=label_alternative - ) - - def lookup_label(x): - return remap_labels.lookup(tf.strings.to_number(x, tf.int32)) - else: - - def lookup_label(x): - return tf.strings.to_number(x, tf.int32) - - @tf.function - def p(string): - x = tf.strings.split(string, sep=",") - label = tf.cast(lookup_label(x[0]), dtype=tf.int32) - - if crop_size: - mod3 = tf.math.floormod(crop_size, 3) - forward_strand = tf.strings.bytes_split(x[1])[:crop_size] - offset_lut = tf.constant([-2, -1, 0], dtype=tf.int32) - offset = tf.gather(offset_lut, mod3) - else: - string_length = tf.strings.length(x[1]) - mod3 = tf.math.floormod(string_length, 3) - offset_lut = tf.constant([-2, -1, 0], dtype=tf.int32) - offset = tf.gather(offset_lut, mod3) - forward_strand = tf.strings.bytes_split(x[1]) - - if mutate: - alphabet = tf.constant(["A", "T", "G", "C", "N"], dtype=tf.string) - mask = tf.random.uniform(tf.shape(forward_strand)) < mutation_rate - mutations = tf.random.uniform( - tf.shape(forward_strand), minval=0, maxval=4, dtype=tf.int32 - ) - forward_strand = tf.where( - mask, tf.gather(alphabet, mutations), forward_strand - ) - if shuffle: - forward_strand = tf.random.shuffle(forward_strand) - - reverse_strand = map_complement.lookup(forward_strand[::-1]) - - outputs = {} - if masking is False: - forward_strand = tf.strings.upper(forward_strand) - reverse_strand = tf.strings.upper(reverse_strand) - - if input_type in ["nucleotide", "both"]: - nuc1 = map_nucleotide.lookup(forward_strand) - nuc2 = map_nucleotide.lookup(reverse_strand) - nuc = tf.stack([nuc1, nuc2], axis=0) - outputs["nucleotide"] = tf.one_hot(nuc, depth=4, dtype=tf.float32) - - if input_type in ["translated", "both"]: - tri_forward = tf.strings.ngrams( - forward_strand, ngram_width=ngram_width, separator="" - ) - tri_reverse = tf.strings.ngrams( - reverse_strand, ngram_width=ngram_width, separator="" - ) - - f1 = map_codon.lookup(tri_forward[0 : -3 + offset : ngram_width]) - f2 = map_codon.lookup(tri_forward[1 : -2 + offset : ngram_width]) - f3 = map_codon.lookup(tri_forward[2 : -1 + offset : ngram_width]) - r1 = map_codon.lookup(tri_reverse[0 : -3 + offset : ngram_width]) - r2 = map_codon.lookup(tri_reverse[1 : -2 + offset : ngram_width]) - r3 = map_codon.lookup(tri_reverse[2 : -1 + offset : ngram_width]) - - if timesteps: - f1 = tf.reshape(f1, (num_time, fragsize)) - f2 = tf.reshape(f2, (num_time, fragsize)) - f3 = tf.reshape(f3, (num_time, fragsize)) - r1 = tf.reshape(r1, (num_time, fragsize)) - r2 = tf.reshape(r2, (num_time, fragsize)) - r3 = tf.reshape(r3, (num_time, fragsize)) - seq = tf.stack([f1, f2, f3, r1, r2, r3], axis=1) - else: - seq = tf.stack([f1, f2, f3, r1, r2, r3], axis=0) - - if shuffle_frames: - perm = tf.random.shuffle(tf.constant([0, 1, 2, 3, 4, 5])) - seq = tf.gather(seq, perm, axis=0) - - if seq_onehot: - outputs["translated"] = tf.one_hot( - seq, depth=codon_depth, dtype=tf.float32, on_value=1, off_value=0 - ) - else: - outputs["translated"] = seq + 1 - - if class_label_onehot: - label = tf.one_hot( - label, depth=num_classes, dtype=tf.float32, on_value=1, off_value=0 - ) - else: - label = tf.expand_dims(label, axis=0) - - return outputs, label - - return p - - -# ------------------------------------------------------------------ -# Inference preprocessor -# ------------------------------------------------------------------ - - -def process_string_inference( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=64, - crop_size=1024, - seq_onehot=True, - timesteps=False, - num_time=None, - fragsize=200, - mutate=False, - mutation_rate=0.1, - ngram_width=3, - shuffle=False, - masking=False, - input_type="translated", -): - """Build a TensorFlow function for inference-time sequence preprocessing. - - Returns a callable that takes a CSV string and returns ``(inputs, metadata)``. - """ - if codons and codon_num: - map_codon = _map_codon(codons=codons, codon_num=codon_num) - map_complement = _map_complement() - map_nucleotide = _map_nucleotide() - - @tf.function - def p(string): - x = tf.strings.split(string, sep=",") - - if crop_size: - mod3 = tf.math.floormod(crop_size, 3) - forward_strand = tf.strings.bytes_split(x[0])[:crop_size] - offset_lut = tf.constant([-2, -1, 0], dtype=tf.int32) - offset = tf.gather(offset_lut, mod3) - else: - string_length = tf.strings.length(x[1]) - mod3 = tf.math.floormod(string_length, 3) - offset_lut = tf.constant([-2, -1, 0], dtype=tf.int32) - offset = tf.gather(offset_lut, mod3) - forward_strand = tf.strings.bytes_split(x[0]) - - if mutate: - alphabet = tf.constant(["A", "T", "G", "C", "N"], dtype=tf.string) - mask = tf.random.uniform(tf.shape(forward_strand)) < mutation_rate - mutations = tf.random.uniform( - tf.shape(forward_strand), minval=0, maxval=4, dtype=tf.int32 - ) - forward_strand = tf.where( - mask, tf.gather(alphabet, mutations), forward_strand - ) - if shuffle: - forward_strand = tf.random.shuffle(forward_strand) - - reverse_strand = map_complement.lookup(forward_strand[::-1]) - - outputs = {} - if masking is False: - forward_strand = tf.strings.upper(forward_strand) - reverse_strand = tf.strings.upper(reverse_strand) - - if input_type in ["nucleotide", "both"]: - nuc1 = map_nucleotide.lookup(forward_strand) - nuc2 = map_nucleotide.lookup(reverse_strand) - nuc = tf.stack([nuc1, nuc2], axis=0) - outputs["nucleotide"] = tf.one_hot( - nuc, depth=4, dtype=tf.float32, on_value=1, off_value=0 - ) - - if input_type in ["translated", "both"]: - tri_forward = tf.strings.ngrams( - forward_strand, ngram_width=ngram_width, separator="" - ) - tri_reverse = tf.strings.ngrams( - reverse_strand, ngram_width=ngram_width, separator="" - ) - - f1 = map_codon.lookup(tri_forward[0 : -3 + offset : ngram_width]) - f2 = map_codon.lookup(tri_forward[1 : -2 + offset : ngram_width]) - f3 = map_codon.lookup(tri_forward[2 : -1 + offset : ngram_width]) - r1 = map_codon.lookup(tri_reverse[0 : -3 + offset : ngram_width]) - r2 = map_codon.lookup(tri_reverse[1 : -2 + offset : ngram_width]) - r3 = map_codon.lookup(tri_reverse[2 : -1 + offset : ngram_width]) - - if timesteps: - f1 = tf.reshape(f1, (num_time, fragsize)) - f2 = tf.reshape(f2, (num_time, fragsize)) - f3 = tf.reshape(f3, (num_time, fragsize)) - r1 = tf.reshape(r1, (num_time, fragsize)) - r2 = tf.reshape(r2, (num_time, fragsize)) - r3 = tf.reshape(r3, (num_time, fragsize)) - seq = tf.stack([f1, f2, f3, r1, r2, r3], axis=1) - else: - seq = tf.stack([f1, f2, f3, r1, r2, r3], axis=0) - - if seq_onehot: - outputs["translated"] = tf.one_hot( - seq, depth=codon_depth, dtype=tf.float32, on_value=1, off_value=0 - ) - else: - outputs["translated"] = tf.cast(seq + 1, dtype=tf.float32) - - return ( - outputs, - x[1], - x[2], - x[3], - x[4], - x[5], - x[6], - x[7], - x[8], - x[9], - x[10], - ) - - return p - - -# ------------------------------------------------------------------ -# Raw sequence processors (for numpy_raw / numpy_raw_variable loaders) -# ------------------------------------------------------------------ - - -def _build_codon_lookup_table() -> tf.Tensor: - """Build a 5x5x5 lookup table for mapping int8 DNA triplets to codon IDs.""" - BASES = ["A", "T", "G", "C", "N"] - codon_to_id = {c: i for i, c in enumerate(CODONS)} - - lookup = np.zeros((5, 5, 5), dtype=np.int32) - for i in range(5): - for j in range(5): - for k in range(5): - codon = BASES[i] + BASES[j] + BASES[k] - lookup[i, j, k] = codon_to_id.get(codon, -1) - - return tf.constant(lookup, dtype=tf.int32) - - -def _make_process_raw_sequence_fn( - crop_size: int = 500, - ngram_width: int = 3, - num_classes: int = 3, - shuffle: bool = False, - mutate: bool = False, - mutation_rate: float = 0.1, - shuffle_frames: bool = False, -): - """Creates a TF function that converts int8 DNA sequences to 6-frame codon embeddings.""" - codon_lookup = _build_codon_lookup_table() - # Complement mapping for int8: A=0->T=1, T=1->A=0, G=2->C=3, C=3->G=2, N=4->N=4 - comp_lookup = tf.constant([1, 0, 3, 2, 4], dtype=tf.int32) - seq_len = crop_size // ngram_width - 1 # 165 for crop_size=500 - - @tf.function - def _process_raw_sequence(sequence, label): - # sequence: int8 tensor of shape (crop_size,) - # label: float32 tensor of shape (num_classes,) - - # Cast to int32 for indexing - seq = tf.cast(sequence, tf.int32) - - # Optional: shuffle bases - if shuffle: - seq = tf.random.shuffle(seq) - - # Optional: mutate bases - if mutate: - mask = tf.random.uniform(tf.shape(seq)) < mutation_rate - random_bases = tf.random.uniform( - tf.shape(seq), minval=0, maxval=5, dtype=tf.int32 - ) - seq = tf.where(mask, random_bases, seq) - - # Forward strand: extract n-grams as triplets of indices - indices = tf.range(crop_size - ngram_width + 1) # 0 to 497 - tri0 = tf.gather(seq, indices) - tri1 = tf.gather(seq, indices + 1) - tri2 = tf.gather(seq, indices + 2) - - # Lookup codon IDs: lookup[tri0, tri1, tri2] - codon_ids = tf.gather_nd(codon_lookup, tf.stack([tri0, tri1, tri2], axis=1)) - - # Reverse complement strand: - # 1. Reverse the sequence - rev_seq = tf.reverse(seq, axis=[0]) - # 2. Map complement - rev_comp_seq = tf.gather(comp_lookup, rev_seq) - # 3. Extract n-grams from reverse complement - rev_tri0 = tf.gather(rev_comp_seq, indices) - rev_tri1 = tf.gather(rev_comp_seq, indices + 1) - rev_tri2 = tf.gather(rev_comp_seq, indices + 2) - # 4. Lookup codon IDs - rev_codon_ids = tf.gather_nd( - codon_lookup, tf.stack([rev_tri0, rev_tri1, rev_tri2], axis=1) - ) - - # Extract 6 reading frames with stride ngram_width (3) - def extract_frames(codons): - f1 = codons[0::ngram_width] - f2 = codons[1::ngram_width] - f3 = codons[2::ngram_width] - return f1, f2, f3 - - f1, f2, f3 = extract_frames(codon_ids) - r1, r2, r3 = extract_frames(rev_codon_ids) - - # Stack frames: shape (6, seq_len) - frames = tf.stack([f1, f2, f3, r1, r2, r3], axis=0) - - # Trim to exact seq_len - frames = frames[:, :seq_len] - - # +1 for embedding layer (0 is padding/mask) - frames = frames + 1 - - # Optional: shuffle frames - if shuffle_frames: - perm = tf.random.shuffle(tf.range(6)) - frames = tf.gather(frames, perm, axis=0) - - return {"translated": frames}, label - - return _process_raw_sequence - - -def _make_process_variable_sequence_fn( - ngram_width: int = 3, - num_classes: int = 3, - shuffle: bool = False, - mutate: bool = False, - mutation_rate: float = 0.1, - shuffle_frames: bool = False, -): - """Creates a TF function for variable-length int8 sequences.""" - codon_lookup = _build_codon_lookup_table() - comp_lookup = tf.constant([1, 0, 3, 2, 4], dtype=tf.int32) - - @tf.function - def _process_variable_sequence(sequence, length, label): - seq = tf.cast(sequence, tf.int32) - actual_len = tf.cast(length, tf.int32) - - # Crop to actual length - seq = seq[:actual_len] - - # Optional: shuffle bases - if shuffle: - seq = tf.random.shuffle(seq) - - # Optional: mutate bases - if mutate: - mask = tf.random.uniform(tf.shape(seq)) < mutation_rate - random_bases = tf.random.uniform( - tf.shape(seq), minval=0, maxval=5, dtype=tf.int32 - ) - seq = tf.where(mask, random_bases, seq) - - # Compute number of codons for this sequence - num_codons = actual_len // ngram_width - 1 - num_codons = tf.maximum(num_codons, 0) - - # Forward strand n-grams - indices = tf.range(actual_len - ngram_width + 1) - tri0 = tf.gather(seq, indices) - tri1 = tf.gather(seq, indices + 1) - tri2 = tf.gather(seq, indices + 2) - - codon_ids = tf.gather_nd(codon_lookup, tf.stack([tri0, tri1, tri2], axis=1)) - - # Reverse complement - rev_seq = tf.reverse(seq, axis=[0]) - rev_comp_seq = tf.gather(comp_lookup, rev_seq) - rev_tri0 = tf.gather(rev_comp_seq, indices) - rev_tri1 = tf.gather(rev_comp_seq, indices + 1) - rev_tri2 = tf.gather(rev_comp_seq, indices + 2) - rev_codon_ids = tf.gather_nd( - codon_lookup, tf.stack([rev_tri0, rev_tri1, rev_tri2], axis=1) - ) - - # Extract frames - def extract_frames(codons): - f1 = codons[0::ngram_width] - f2 = codons[1::ngram_width] - f3 = codons[2::ngram_width] - return f1, f2, f3 - - f1, f2, f3 = extract_frames(codon_ids) - r1, r2, r3 = extract_frames(rev_codon_ids) - - # Find minimum frame length to ensure all frames match - min_len = tf.reduce_min( - [ - tf.shape(f1)[0], - tf.shape(f2)[0], - tf.shape(f3)[0], - tf.shape(r1)[0], - tf.shape(r2)[0], - tf.shape(r3)[0], - ] - ) - - # Trim all frames to same length - f1, f2, f3 = f1[:min_len], f2[:min_len], f3[:min_len] - r1, r2, r3 = r1[:min_len], r2[:min_len], r3[:min_len] - - # Stack frames - frames = tf.stack([f1, f2, f3, r1, r2, r3], axis=0) - - # +1 for embedding layer - frames = frames + 1 - - # Optional: shuffle frames - if shuffle_frames: - perm = tf.random.shuffle(tf.range(6)) - frames = tf.gather(frames, perm, axis=0) - - return {"translated": frames}, label - - return _process_variable_sequence diff --git a/src/jaeger/utils/gpu.py b/src/jaeger/utils/gpu.py deleted file mode 100644 index dba750e..0000000 --- a/src/jaeger/utils/gpu.py +++ /dev/null @@ -1,36 +0,0 @@ -import tensorflow as tf -import traceback -import logging - -logger = logging.getLogger("Jaeger") - - -def configure_multi_gpu_inference(gpus): - if gpus > 0: - return tf.distribute.MirroredStrategy() - else: - None - - -def get_device_name(device): - name = device.name - return f"{name.split(':', 1)[-1]}" - - -def create_virtual_gpus(logger, num_gpus=2, memory_limit=2048): - if gpus := tf.config.list_physical_devices("GPU"): - # Create n virtual GPUs with user defined amount of memory each - try: - tf.config.set_logical_device_configuration( - gpus[0], - [ - tf.config.LogicalDeviceConfiguration(memory_limit=memory_limit) - for _ in range(num_gpus) - ], - ) - logical_gpus = tf.config.list_logical_devices("GPU") - logger.info(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs") - except RuntimeError as e: - # Virtual devices must be set before GPUs have been initialized - logger.error(e) - logger.debug(traceback.format_exc()) diff --git a/src/jaeger/utils/test.py b/src/jaeger/utils/test.py index 8f16b25..e2cacc7 100644 --- a/src/jaeger/utils/test.py +++ b/src/jaeger/utils/test.py @@ -1,19 +1,19 @@ import numpy as np -import tensorflow as tf -def test_tf(): +def test_torch(): + """Smoke test: run a small PyTorch matmul on the CPU.""" try: - # Set the matrix size - matrix_size = 100 - - # Create two random matrices - matrix_a = np.random.rand(matrix_size, matrix_size) - matrix_b = np.random.rand(matrix_size, matrix_size) - - with tf.device("CPU:0"): - result = tf.matmul(matrix_a, matrix_b) + import torch + matrix_size = 100 + matrix_a = torch.from_numpy( + np.random.rand(matrix_size, matrix_size).astype(np.float32) + ) + matrix_b = torch.from_numpy( + np.random.rand(matrix_size, matrix_size).astype(np.float32) + ) + result = torch.matmul(matrix_a, matrix_b) except Exception as e: return e else: diff --git a/tests/conftest.py b/tests/conftest.py index 970a41b..c6748a8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,17 +2,11 @@ from __future__ import annotations -import os import shutil from pathlib import Path import numpy as np import pytest -import tensorflow as tf - -# Keep TF reasonably quiet during tests. -os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2") -os.environ.setdefault("GRPC_VERBOSITY", "ERROR") @pytest.fixture diff --git a/tests/integration/test_inference_pipeline.py b/tests/integration/test_inference_pipeline.py deleted file mode 100644 index 777bb11..0000000 --- a/tests/integration/test_inference_pipeline.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Integration test: build a model, save it, and run inference.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest -import tensorflow as tf - -from jaeger.nnlib import inference -from jaeger.postprocess import collect - - -@pytest.fixture -def tiny_config(): - return { - "model": { - "seed": 7, - "class_label_map": {0: "bacteria", 1: "phage"}, - "embedding": {"type": "translated", "input_shape": (6, 16, 65)}, - "representation_learner": { - "block_sizes": [1], - "block_filters": [4], - "block_kernel_size": [3], - "block_kernel_dilation": [1], - "block_kernel_strides": [1], - "block_regularizer": ["l2"], - "block_regularizer_w": [0.0], - "pooling": "max", - "masked_conv1d_1_filters": 4, - "masked_conv1d_1_kernel_size": 3, - "masked_conv1d_1_strides": 1, - "masked_conv1d_1_dilation_rate": 1, - "masked_conv1d_1_regularizer": "l2", - "masked_conv1d_1_regularizer_w": 0.0, - "masked_conv1d_final_kernel_size": 3, - "masked_conv1d_final_strides": 1, - "masked_conv1d_final_dilation_rate": 1, - }, - "classifier": { - "output_units": 2, - "hidden_layers": [{"units": 4}], - }, - "reliability_model": { - "output_units": 1, - "hidden_layers": [], - }, - } - } - - -@pytest.mark.integration -class TestInferencePipeline: - def test_save_load_and_predict(self, tiny_config, tmp_path: Path): - builder = inference.DynamicInferenceModelBuilder(tiny_config) - model = builder.models["jaeger_model"] - - export_dir = tmp_path / "graph" - tf.saved_model.save(model, str(export_dir)) - - loaded = inference.InferModel( - path_dict={ - "graph": str(export_dir), - "classes": None, - "project": None, - } - ) - - x = {"translated": tf.random.normal((2, 6, 16, 65))} - # InferModel.predict expects a dataset yielding (inputs, *meta) - dataset = tf.data.Dataset.from_tensors( - (x, np.array(["s1", "s2"]), np.array([1, 1])) - ) - result = loaded.predict(dataset, no_progress=True) - assert "prediction" in result - assert result["prediction"].shape[0] == 2 - - def test_end_to_end_postprocess(self, tiny_config, tmp_path: Path): - builder = inference.DynamicInferenceModelBuilder(tiny_config) - model = builder.models["jaeger_model"] - - # Create fake windowed data for two contigs. - inputs = {"translated_input": tf.random.normal((4, 6, 16, 65))} - y = model(inputs, training=False) - - y_pred = { - "prediction": y["prediction"].numpy(), - "reliability": y["reliability"].numpy(), - "meta_0": np.array(["c1", "c1", "c2", "c2"]), - "meta_2": np.array([0, 1, 0, 1]), - "meta_4": np.array([500, 500, 500, 500]), - "meta_5": np.full(4, 125.0, dtype=np.float32), - "meta_6": np.full(4, 125.0, dtype=np.float32), - "meta_7": np.full(4, 125.0, dtype=np.float32), - "meta_8": np.full(4, 125.0, dtype=np.float32), - "meta_9": np.zeros(4, dtype=np.float32), - } - - import pandas as pd - - data, _ = collect.pred_to_dict( - y_pred, - fsize=500, - class_map={"num_classes": 2}, - term_repeats=pd.DataFrame( - { - "contig_id": ["c1", "c2"], - "terminal_repeats": [0, 0], - "repeat_length": [0, 0], - } - ), - ) - assert len(data["headers"]) == 2 diff --git a/tests/integration/test_pytorch_parity.py b/tests/integration/test_pytorch_parity.py deleted file mode 100644 index 0790db1..0000000 --- a/tests/integration/test_pytorch_parity.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Forward parity test between the legacy TensorFlow builder and PyTorch builder. - -This test builds the same tiny fragment classifier with both backends, copies the -TensorFlow weights into the PyTorch model, and asserts that the forward outputs -agree within tolerance. -""" - -import os - -# Suppress TensorFlow logging before importing it. -os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") -os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0") -os.environ.setdefault("GRPC_VERBOSITY", "ERROR") -os.environ.setdefault("GLOG_minloglevel", "2") -os.environ.setdefault("ABSL_MIN_LOG_LEVEL", "2") - -import numpy as np -import torch - -from jaeger.nnlib.builder import DynamicModelBuilder as TFBuilder -from jaeger.nnlib.pytorch.builder import ModelBuilder as PTBuilder - - -CONFIG = { - "model": { - "name": "parity_test", - "classifier_out_dim": 3, - "reliability_out_dim": 1, - "class_label_map": [ - {"class": "phage", "label": 1}, - {"class": "bacteria", "label": 0}, - {"class": "eukarya", "label": 2}, - ], - "embedding": { - "input_type": "translated", - "use_embedding_layer": True, - "embedding_size": 32, - "input_shape": [6, None], - "vocab_size": 65, - "codon_depth": 1, - "embedding_regularizer": "l2", - "embedding_regularizer_w": 1e-5, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - "crop_size": 50, - "input_type": "translated", - "use_embedding_layer": True, - }, - "representation_learner": { - "hidden_layers": [ - { - "name": "masked_conv1d", - "config": { - "filters": 16, - "kernel_size": 3, - "padding": "same", - }, - }, - ], - "pooling": "average", - }, - "classifier": { - "input_shape": 16, - "hidden_layers": [{"name": "dense", "config": {"units": 3}}], - }, - }, - "training": { - "batch_size": 2, - "optimizer": "adam", - "optimizer_params": {"learning_rate": 1e-3}, - }, -} - - -def _collect_tf_weights(layer, dest): - """Recursively collect leaf-layer TF weights keyed by (layer_name, weight_name).""" - has_sublayers = hasattr(layer, "layers") and layer.layers - if not has_sublayers and hasattr(layer, "weights"): - for weight in layer.weights: - # Weight names from top-level leaf layers are short (e.g. "kernel"), - # while nested leaf layers include the layer path - # (e.g. "classifier_dense_0/kernel"). - full_name = weight.name - if "/" in full_name: - layer_name, weight_name = full_name.rsplit("/", 1) - else: - layer_name = layer.name - weight_name = full_name - dest[(layer_name, weight_name)] = weight - if has_sublayers: - for sub in layer.layers: - _collect_tf_weights(sub, dest) - - -def _copy_tf_weights_to_pt(tf_model, pt_model): - """Copy weights from the TF model to the PyTorch model in-place.""" - # PyTorch MaskedConv1D is lazy: run one forward pass to materialize the real - # Conv1d with the correct in_channels before copying weights. - dummy_x = torch.zeros(2, 6, 50, dtype=torch.int64) - dummy_mask = torch.ones(2, 6, 50, dtype=torch.bool) - pt_model.eval() - with torch.no_grad(): - pt_model(dummy_x, mask=dummy_mask) - - tf_weights = {} - for layer in tf_model.layers: - _collect_tf_weights(layer, tf_weights) - - pt_state = pt_model.state_dict() - - def _get(name_pair): - weight = tf_weights.get(name_pair) - if weight is None: - raise ValueError(f"Missing TensorFlow weight for {name_pair}") - return torch.from_numpy(weight.numpy()) - - def _set(pt_key, tf_weight, transform=None): - if pt_key not in pt_state: - raise ValueError(f"Missing PyTorch state_dict key: {pt_key}") - tensor = _get(tf_weight) - if transform is not None: - tensor = transform(tensor) - pt_state[pt_key].copy_(tensor) - - # Embedding layer: direct copy. - _set("rep_model.embedding.embed.weight", ("embedding", "embeddings")) - - # MaskedConv1D: TF kernel shape (K, C_in, C_out) -> PyTorch (C_out, C_in, K). - _set( - "rep_model.blocks.0.conv.weight", - ("rep_masked_conv1d_0", "kernel"), - transform=lambda t: t.permute(2, 1, 0), - ) - _set("rep_model.blocks.0.conv.bias", ("rep_masked_conv1d_0", "bias")) - - # Classification dense head: TF kernel shape (C_in, C_out) -> PyTorch (C_out, C_in). - _set( - "classification_head.net.0.weight", - ("classifier_dense_0", "kernel"), - transform=lambda t: t.T, - ) - _set("classification_head.net.0.bias", ("classifier_dense_0", "bias")) - - -def test_forward_parity(): - tf_builder = TFBuilder(CONFIG) - tf_models = tf_builder.build_fragment_classifier() - tf_model = tf_models["jaeger_model"] - - pt_builder = PTBuilder(CONFIG) - pt_models = pt_builder.build_fragment_classifier() - pt_model = pt_models["jaeger_model"] - pt_model.eval() - - _copy_tf_weights_to_pt(tf_model, pt_model) - - # Avoid zeros so the TF model's automatic mask (x != 0) is all-ones and - # matches the all-ones PyTorch mask, sidestepping the fact that TF's - # GlobalAveragePooling2D does not consume masks while PyTorch's pooler does. - x = np.random.randint(1, 65, size=(2, 6, 50)).astype(np.int64) - tf_out = tf_model.predict({"translated": x}, verbose=0) - with torch.no_grad(): - pt_out = pt_model( - torch.from_numpy(x), mask=torch.ones(2, 6, 50, dtype=torch.bool) - ) - - np.testing.assert_allclose( - tf_out["prediction"], - pt_out["prediction"].numpy(), - atol=1e-4, - rtol=1e-4, - ) diff --git a/tests/pytest/test_health.py b/tests/pytest/test_health.py deleted file mode 100644 index eb9e84e..0000000 --- a/tests/pytest/test_health.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Tests for the jaeger health command.""" - -import sys -import platform -from pathlib import Path -import pytest - -# Skip all tests if tensorflow is not available -import importlib.util - -HAS_TF = importlib.util.find_spec("tensorflow") is not None - - -pytestmark = pytest.mark.skipif(not HAS_TF, reason="tensorflow not installed") - - -class TestHealthDiagnostics: - """Test the diagnostic info gathering in health command.""" - - def test_python_version_reported(self): - """Check that Python version is accessible.""" - assert platform.python_version() is not None - assert "." in platform.python_version() - - def test_python_executable_exists(self): - """Check that Python executable path exists.""" - assert Path(sys.executable).exists() - - def test_platform_info(self): - """Check that platform info is non-empty.""" - assert platform.platform() - assert platform.machine() - - -class TestHealthDependencies: - """Test that core dependencies are importable and have versions.""" - - def test_jaeger_bio_importable(self): - """jaeger-bio package should be importable.""" - import jaeger - - assert jaeger.__file__ is not None - - def test_tensorflow_importable(self): - """tensorflow should be importable.""" - import tensorflow as tf - - assert tf.__version__ is not None - - def test_numpy_importable(self): - """numpy should be importable.""" - import numpy as np - - assert np.__version__ is not None - - def test_click_importable(self): - """click should be importable.""" - import click - - assert click.__version__ is not None - - def test_parasail_importable(self): - """parasail should be importable.""" - import parasail - - assert parasail.__version__ is not None - - def test_pyfastx_importable(self): - """pyfastx should be importable.""" - import pyfastx - - # pyfastx doesn't expose a simple version attribute; just check it loads - assert pyfastx.Fasta is not None - - def test_pydustmasker_importable(self): - """pydustmasker should be importable.""" - import pydustmasker - - assert pydustmasker.__version__ is not None - - def test_sklearn_importable(self): - """scikit-learn should be importable.""" - import sklearn - - assert sklearn.__version__ is not None - - def test_polars_importable(self): - """polars should be importable.""" - import polars - - assert polars.__version__ is not None - - def test_pandas_importable(self): - """pandas should be importable.""" - import pandas - - assert pandas.__version__ is not None - - def test_matplotlib_importable(self): - """matplotlib should be importable.""" - import matplotlib - - assert matplotlib.__version__ is not None - - def test_ruptures_importable(self): - """ruptures should be importable.""" - import ruptures - - assert ruptures.__version__ is not None - - def test_pycirclize_importable(self): - """pycirclize should be importable.""" - import pycirclize - - assert pycirclize.__version__ is not None - - def test_biopython_importable(self): - """biopython should be importable.""" - import Bio - - assert Bio.__version__ is not None - - -class TestHealthTensorFlow: - """Test TensorFlow-specific health checks.""" - - def test_tf_can_list_devices(self): - """TensorFlow should be able to list physical devices.""" - import tensorflow as tf - - cpus = tf.config.list_physical_devices("CPU") - assert len(cpus) >= 1, "At least one CPU should be available" - - def test_tf_matrix_multiplication(self): - """TensorFlow should be able to perform basic matrix ops.""" - import tensorflow as tf - import numpy as np - - a = np.random.rand(10, 10).astype(np.float32) - b = np.random.rand(10, 10).astype(np.float32) - - with tf.device("CPU:0"): - result = tf.matmul(a, b) - - assert result.shape == (10, 10) - - def test_tf_version_parsable(self): - """TensorFlow version should be a non-empty string.""" - import tensorflow as tf - - assert tf.__version__ - assert "." in tf.__version__ - - -class TestHealthModels: - """Test model discovery in health command.""" - - def test_config_json_exists(self): - """config.json should exist in jaeger.data.""" - from importlib.resources import files - - config_path = files("jaeger.data").joinpath("config.json") - assert config_path.exists(), "config.json not found in jaeger.data" - - def test_config_json_loadable(self): - """config.json should be valid JSON with expected keys.""" - import json - from importlib.resources import files - - config_path = files("jaeger.data").joinpath("config.json") - config = json.loads(config_path.read_text()) - assert "default" in config - assert "model_paths" in config - - def test_default_weights_exist(self): - """Default model weights should exist.""" - import json - from importlib.resources import files - - config_path = files("jaeger.data").joinpath("config.json") - config = json.loads(config_path.read_text()) - weights_name = config["default"]["weights"] - weights_path = files("jaeger.data.models.default").joinpath(weights_name) - assert weights_path.exists(), f"Default weights {weights_name} not found" - - def test_test_fasta_files_exist(self): - """Test FASTA files should exist in jaeger.data.test.""" - from importlib.resources import files - - for fname in ["test_short.fasta", "test_empty.fasta", "test_contigs.fasta"]: - fpath = files("jaeger.data.test").joinpath(fname) - assert fpath.exists(), f"{fname} not found in jaeger.data.test" - - -class TestHealthCoreFunction: - """Test the health_core function directly.""" - - def test_health_core_runs_without_error(self, tmp_path, caplog): - """health_core should run without raising exceptions.""" - from jaeger.commands.health import health_core - - # Change to tmp_path so test_log is created there - import os - - original_cwd = os.getcwd() - os.chdir(tmp_path) - try: - # Run with minimal verbosity to reduce output - health_core(verbose=0) - finally: - os.chdir(original_cwd) - - def test_health_core_creates_log(self, tmp_path): - """health_core should create a log file.""" - import os - from jaeger.commands.health import health_core - - original_cwd = os.getcwd() - os.chdir(tmp_path) - try: - health_core(verbose=0) - log_dir = tmp_path / "test_log" - assert log_dir.exists(), "test_log directory not created" - log_files = list(log_dir.glob("*test_jaeger.log")) - assert len(log_files) > 0, "No log file created" - finally: - os.chdir(original_cwd) - - def test_health_core_all_tests_pass(self, tmp_path): - """health_core should report 5/5 tests passed.""" - import os - from jaeger.commands.health import health_core - - original_cwd = os.getcwd() - os.chdir(tmp_path) - try: - health_core(verbose=0) - log_dir = tmp_path / "test_log" - log_files = sorted(log_dir.glob("*.log"), key=lambda p: p.stat().st_mtime) - assert log_files, "No log file found" - log_content = log_files[-1].read_text() - assert "5/5 tests passed!" in log_content, ( - f"Expected 5/5 tests passed, got:\n{log_content[-500:]}" - ) - finally: - os.chdir(original_cwd) diff --git a/tests/pytest/test_inference.py b/tests/pytest/test_inference.py deleted file mode 100644 index 97404ef..0000000 --- a/tests/pytest/test_inference.py +++ /dev/null @@ -1,455 +0,0 @@ -"""Tests for the inference pipeline (InferModel and _load_string_processor_config).""" - -from pathlib import Path - -import pytest -import tensorflow as tf -import yaml - -# Skip all tests if tensorflow is not available -try: - import tensorflow as tf - - HAS_TF = True -except ImportError: - HAS_TF = False - -pytestmark = pytest.mark.skipif(not HAS_TF, reason="tensorflow not installed") - - -class TestLoadStringProcessorConfig: - """Test _load_string_processor_config fixes for input_type and seq_onehot.""" - - def test_input_type_extracted_from_embedding_type(self, tmp_path): - """input_type should be inferred from embedding.type when not explicitly set.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None, 64], - "frames": 6, - "strands": 2, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - # Create a dummy SavedModel - graph_dir = tmp_path / "test_graph" - _create_dummy_savedmodel(graph_dir) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - assert model.string_processor_config["input_type"] == "translated" - - def test_input_type_defaults_to_translated(self, tmp_path): - """input_type should default to 'translated' if embedding.type is missing.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "input_shape": [6, None, 64], - "frames": 6, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - graph_dir = tmp_path / "test_graph" - _create_dummy_savedmodel(graph_dir) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - assert model.string_processor_config["input_type"] == "translated" - - def test_seq_onehot_inferred_from_input_shape(self, tmp_path): - """seq_onehot should be inferred True when input_shape has depth > 1.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None, 64], # 3 dims = one-hot - "frames": 6, - "strands": 2, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - graph_dir = tmp_path / "test_graph" - _create_dummy_savedmodel(graph_dir) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - assert model.string_processor_config["seq_onehot"] is True - assert model.string_processor_config["codon_depth"] == 64 - - def test_seq_onehot_inferred_false_for_embedding_lookup(self, tmp_path): - """seq_onehot should be inferred False when input_shape has only 2 dims.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None], # 2 dims = embedding lookup - "frames": 6, - "strands": 2, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - graph_dir = tmp_path / "test_graph" - _create_dummy_savedmodel(graph_dir) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - assert model.string_processor_config["seq_onehot"] is False - assert model.string_processor_config["codon_depth"] == 1 - - def test_explicit_seq_onehot_overrides_inference(self, tmp_path): - """Explicitly set seq_onehot should not be overridden by input_shape inference.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None, 64], - "frames": 6, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - "seq_onehot": False, # explicitly False - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - graph_dir = tmp_path / "test_graph" - _create_dummy_savedmodel(graph_dir) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - assert model.string_processor_config["seq_onehot"] is False - assert model.string_processor_config["codon_depth"] == 1 - - def test_path_as_string_accepted(self, tmp_path): - """_load_string_processor_config should accept string paths, not just Path objects.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None, 64], - "frames": 6, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - graph_dir = tmp_path / "test_graph" - _create_dummy_savedmodel(graph_dir) - - # Pass paths as strings - model = InferModel( - { - "graph": str(graph_dir), - "classes": str(classes_file), - "project": str(project_file), - } - ) - - assert model.string_processor_config["input_type"] == "translated" - - -class TestInferModelPrediction: - """Test InferModel prediction with SavedModel signatures.""" - - def test_inference_fn_called_with_keyword_inputs(self, tmp_path): - """_predict_step should call inference_fn with inputs= keyword argument.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None, 4], - "frames": 6, - "strands": 2, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - "seq_onehot": True, - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - # Create a SavedModel that expects inputs= keyword - graph_dir = tmp_path / "test_graph" - _create_savedmodel_with_keyword_input(graph_dir, input_shape=(None, 6, None, 4)) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - # Create dummy input batch - dummy_input = { - "translated": tf.random.normal([2, 6, 10, 4]), - } - - # This should not raise TypeError about positional arguments - result = model._predict_step(dummy_input) - assert "prediction" in result or "output" in result - - -class TestLegacyAndNewPipeline: - """Test that both legacy and new inference pipelines work.""" - - def test_legacy_model_multiple_inputs(self, tmp_path): - """Legacy models with multiple inputs should still load correctly.""" - from jaeger.nnlib.inference import InferModel - - # Create a legacy-style SavedModel with multiple inputs - graph_dir = tmp_path / "legacy_graph" - _create_legacy_savedmodel(graph_dir) - - classes_file = tmp_path / "classes.yaml" - classes_file.write_text( - yaml.safe_dump( - { - "classes": [ - {"class": "bacteria", "label": 0}, - {"class": "phage", "label": 1}, - ] - } - ) - ) - - # Legacy models don't have project.yaml - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - } - ) - - # Should load without error - assert model.loaded_model is not None - assert "serving_default" in model.loaded_model.signatures - - # Legacy models have multiple inputs - sig = model.loaded_model.signatures["serving_default"] - inputs = sig.structured_input_signature[1] - assert len(inputs) > 1, "Legacy model should have multiple inputs" - - def test_new_model_single_input(self, tmp_path): - """New models with single 'inputs' keyword should work.""" - from jaeger.nnlib.inference import InferModel - - project_file = tmp_path / "project.yaml" - classes_file = tmp_path / "classes.yaml" - - project = { - "model": { - "embedding": { - "type": "translated", - "input_shape": [6, None, 4], - "frames": 6, - "strands": 2, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - "seq_onehot": True, - }, - } - } - project_file.write_text(yaml.safe_dump(project)) - classes_file.write_text(yaml.safe_dump({"classes": []})) - - graph_dir = tmp_path / "test_graph" - _create_savedmodel_with_keyword_input(graph_dir, input_shape=(None, 6, None, 4)) - - model = InferModel( - { - "graph": graph_dir, - "classes": classes_file, - "project": project_file, - } - ) - - # Verify signature - sig = model.loaded_model.signatures["serving_default"] - input_spec = sig.structured_input_signature[1].get("inputs") - assert input_spec is not None - - -# ─── Helpers ──────────────────────────────────────────────────────────────── - - -def _create_dummy_savedmodel(export_dir: Path): - """Create a minimal SavedModel for testing.""" - export_dir = Path(export_dir) - export_dir.mkdir(parents=True, exist_ok=True) - - class DummyModel(tf.Module): - @tf.function( - input_signature=[ - tf.TensorSpec( - shape=(None, 6, None, 64), dtype=tf.float32, name="inputs" - ) - ] - ) - def serving_default(self, inputs): - return { - "prediction": tf.zeros([tf.shape(inputs)[0], 2]), - "reliability": tf.zeros([tf.shape(inputs)[0], 1]), - } - - model = DummyModel() - tf.saved_model.save( - model, str(export_dir), signatures={"serving_default": model.serving_default} - ) - - -def _create_savedmodel_with_keyword_input(export_dir: Path, input_shape): - """Create a SavedModel that expects inputs= keyword argument.""" - export_dir = Path(export_dir) - export_dir.mkdir(parents=True, exist_ok=True) - - class KeywordInputModel(tf.Module): - @tf.function( - input_signature=[ - tf.TensorSpec(shape=input_shape, dtype=tf.float32, name="inputs") - ] - ) - def serving_default(self, inputs): - batch_size = tf.shape(inputs)[0] - return { - "prediction": tf.zeros([batch_size, 2]), - "reliability": tf.zeros([batch_size, 1]), - } - - model = KeywordInputModel() - tf.saved_model.save( - model, str(export_dir), signatures={"serving_default": model.serving_default} - ) - - -def _create_legacy_savedmodel(export_dir: Path): - """Create a legacy-style SavedModel with multiple named inputs.""" - export_dir = Path(export_dir) - export_dir.mkdir(parents=True, exist_ok=True) - - class LegacyModel(tf.Module): - @tf.function( - input_signature=[ - tf.TensorSpec(shape=(None, None), dtype=tf.float32, name="inputs_1"), - tf.TensorSpec(shape=(None, None), dtype=tf.float32, name="inputs_2"), - tf.TensorSpec(shape=(None, None), dtype=tf.float32, name="inputs_3"), - ] - ) - def serving_default(self, inputs_1, inputs_2, inputs_3): - batch_size = tf.shape(inputs_1)[0] - return { - "output": tf.zeros([batch_size, 4]), - "embedding": tf.zeros([batch_size, 128]), - } - - model = LegacyModel() - tf.saved_model.save( - model, str(export_dir), signatures={"serving_default": model.serving_default} - ) diff --git a/tests/smoke/convert_test.py b/tests/smoke/convert_test.py deleted file mode 100644 index 0b03397..0000000 --- a/tests/smoke/convert_test.py +++ /dev/null @@ -1,40 +0,0 @@ -import tensorflow as tf -from jaeger.seqops.encode import process_string_train - -ds = tf.data.Dataset.from_tensor_slices([ - "0,ATGCGCACGTAGACTACGTACGAC", - "1,CGTACGTAGCTAGCTAGCTAGCTA", - "0,ATGCGTACGTAGCTAGCTAGCTAGC", -]) - -ds = ds.map( - process_string_train( - input_type="both", - num_classes=2, - codon_depth=64, - ngram_width=3, - ), - num_parallel_calls=tf.data.AUTOTUNE, -) - -ds = ds.shuffle(6).padded_batch( - 3, - padded_shapes=( - { - "nucleotide": [2, None, 4], - "translated": [6, None, 64], - }, - [2], - ), -).prefetch(tf.data.AUTOTUNE) - -for outputs, label in ds: - print(label) - print(outputs["nucleotide"].shape) - print(outputs["translated"].shape) - -assert outputs["nucleotide"].shape[0] == 3 -assert outputs["translated"].shape[0] == 3 -assert outputs["translated"].shape[1] == 6 -assert outputs["translated"].shape[3] == 64 -print("convert_test passed") diff --git a/tests/smoke/data_format_test.py b/tests/smoke/data_format_test.py deleted file mode 100644 index 4aaf491..0000000 --- a/tests/smoke/data_format_test.py +++ /dev/null @@ -1,301 +0,0 @@ -"""Smoke tests for TFRecord and NumPy data format loading in training.""" - -import os -import tempfile -import numpy as np -import tensorflow as tf - -from jaeger.data.tfrecord import ( - _make_parse_tfrecord_fn, - _get_tfrecord_feature_description, -) -from jaeger.data.loaders import _load_numpy_full_dataset -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - -# Test configuration matching typical Jaeger training setup -CROP_SIZE = 500 -SEQ_LEN = CROP_SIZE // 3 - 1 # 165 -NUM_CLASSES = 3 -CODON_DEPTH = 21 -BATCH_SIZE = 4 -NUM_SAMPLES = 10 - - -def create_test_csv(path, num_samples=10): - """Create a small CSV file with random DNA sequences for testing.""" - bases = ["A", "T", "G", "C"] - with open(path, "w") as f: - for i in range(num_samples): - label = i % NUM_CLASSES - seq = "".join(np.random.choice(bases, size=CROP_SIZE)) - f.write(f"{label},{seq}\n") - - -def test_tfrecord_feature_description(): - """Test that feature descriptions are generated correctly.""" - # Embedding layer (int64) - desc = _get_tfrecord_feature_description( - input_type="translated", - use_embedding_layer=True, - codon_depth=CODON_DEPTH, - crop_size=CROP_SIZE, - num_classes=NUM_CLASSES, - ) - assert "translated" in desc - assert "label" in desc - print("✓ TFRecord feature description (embedding) correct") - - # One-hot (float) - desc = _get_tfrecord_feature_description( - input_type="translated", - use_embedding_layer=False, - codon_depth=CODON_DEPTH, - crop_size=CROP_SIZE, - num_classes=NUM_CLASSES, - ) - assert "translated" in desc - assert "label" in desc - print("✓ TFRecord feature description (one-hot) correct") - - -def test_parse_tfrecord_embedding(): - """Test parsing TFRecord examples with embedding layer (int indices).""" - with tempfile.TemporaryDirectory() as tmpdir: - csv_path = os.path.join(tmpdir, "test.csv") - tfrecord_path = os.path.join(tmpdir, "test.tfrecord") - create_test_csv(csv_path, NUM_SAMPLES) - - # Preprocess and write TFRecord - preprocess_fn = process_string_train( - crop_size=CROP_SIZE, - seq_onehot=False, - input_type="translated", - class_label_onehot=True, - num_classes=NUM_CLASSES, - shuffle=False, - ngram_width=3, - codons=CODONS, - codon_num=CODON_ID, - ) - - with tf.io.TFRecordWriter(tfrecord_path) as writer: - with open(csv_path) as f: - for line in f: - outputs, label = preprocess_fn(line.strip().encode()) - translated = outputs["translated"] - translated_flat = tf.reshape(translated, [-1]) - feature = { - "translated": tf.train.Feature( - int64_list=tf.train.Int64List( - value=translated_flat.numpy().tolist() - ) - ), - "label": tf.train.Feature( - float_list=tf.train.FloatList( - value=label.numpy().flatten().tolist() - ) - ), - } - example = tf.train.Example( - features=tf.train.Features(feature=feature) - ) - writer.write(example.SerializeToString()) - - # Parse back - parse_fn = _make_parse_tfrecord_fn( - input_type="translated", - use_embedding_layer=True, - codon_depth=CODON_DEPTH, - crop_size=CROP_SIZE, - num_classes=NUM_CLASSES, - ) - - dataset = tf.data.TFRecordDataset(tfrecord_path).map(parse_fn) - count = 0 - for features, label in dataset: - assert "translated" in features - assert features["translated"].shape == (6, SEQ_LEN) - assert features["translated"].dtype == tf.int32 - assert label.shape == (NUM_CLASSES,) - assert label.dtype == tf.float32 - count += 1 - - assert count == NUM_SAMPLES - print(f"✓ TFRecord embedding parse: {count} samples, shapes correct") - - -def test_load_numpy_full_dataset_with_legacy_params(): - """Test loading NumPy .npz files with legacy params (backward compatibility).""" - with tempfile.TemporaryDirectory() as tmpdir: - csv_path = os.path.join(tmpdir, "test.csv") - npz_path = os.path.join(tmpdir, "test.npz") - create_test_csv(csv_path, NUM_SAMPLES) - - # Preprocess and save as NumPy - preprocess_fn = process_string_train( - crop_size=CROP_SIZE, - seq_onehot=False, - input_type="translated", - class_label_onehot=True, - num_classes=NUM_CLASSES, - shuffle=False, - ngram_width=3, - codons=CODONS, - codon_num=CODON_ID, - ) - - sequences = np.zeros((NUM_SAMPLES, 6, SEQ_LEN), dtype=np.int32) - labels = np.zeros((NUM_SAMPLES, NUM_CLASSES), dtype=np.float32) - - with open(csv_path) as f: - for i, line in enumerate(f): - outputs, label = preprocess_fn(line.strip().encode()) - sequences[i] = outputs["translated"].numpy() - labels[i] = label.numpy() - - np.savez_compressed(npz_path, translated=sequences, label=labels) - - # Load back with legacy params - dataset = _load_numpy_full_dataset( - npz_path, - input_type="translated", - use_embedding_layer=True, - codon_depth=CODON_DEPTH, - crop_size=CROP_SIZE, - ) - - count = 0 - for features, label in dataset: - assert "translated" in features - assert features["translated"].shape == (6, SEQ_LEN) - assert features["translated"].dtype == tf.int32 - assert label.shape == (NUM_CLASSES,) - assert label.dtype == tf.float32 - count += 1 - - assert count == NUM_SAMPLES - print(f"✓ NumPy full load (legacy params): {count} samples, shapes correct") - - -def test_numpy_flattened_format(): - """Test loading NumPy .npz with flattened sequences (as saved by converter).""" - with tempfile.TemporaryDirectory() as tmpdir: - npz_path = os.path.join(tmpdir, "test_flat.npz") - - # Create flattened data (as the conversion script produces) - sequences = np.random.randint( - 0, 21, size=(NUM_SAMPLES, 6 * SEQ_LEN), dtype=np.int32 - ) - labels = np.random.rand(NUM_SAMPLES, NUM_CLASSES).astype(np.float32) - - np.savez_compressed(npz_path, translated=sequences, label=labels) - - # Load back - dataset = _load_numpy_full_dataset( - npz_path, - input_type="translated", - use_embedding_layer=True, - codon_depth=CODON_DEPTH, - crop_size=CROP_SIZE, - ) - - count = 0 - for features, label in dataset: - assert features["translated"].shape == (6, SEQ_LEN) - count += 1 - - assert count == NUM_SAMPLES - print(f"✓ NumPy full flattened load: {count} samples, reshaped correctly") - - -def test_dataset_pipeline_with_batching(): - """Test that loaded datasets work correctly with batching and prefetching.""" - with tempfile.TemporaryDirectory() as tmpdir: - npz_path = os.path.join(tmpdir, "test.npz") - - sequences = np.random.randint( - 0, 21, size=(NUM_SAMPLES, 6, SEQ_LEN), dtype=np.int32 - ) - labels = np.random.rand(NUM_SAMPLES, NUM_CLASSES).astype(np.float32) - np.savez_compressed(npz_path, translated=sequences, label=labels) - - dataset = ( - _load_numpy_full_dataset( - npz_path, - input_type="translated", - use_embedding_layer=True, - codon_depth=CODON_DEPTH, - crop_size=CROP_SIZE, - ) - .cache() - .shuffle(buffer_size=100) - .batch(BATCH_SIZE) - .prefetch(tf.data.AUTOTUNE) - ) - - for features, label in dataset: - assert features["translated"].shape == (BATCH_SIZE, 6, SEQ_LEN) - assert label.shape == (BATCH_SIZE, NUM_CLASSES) - break - - print("✓ Dataset pipeline with batching works correctly") - - -def test_load_numpy_full_dataset(): - """Test loading fully-preprocessed NumPy .npz files (numpy_full format).""" - with tempfile.TemporaryDirectory() as tmpdir: - npz_path = os.path.join(tmpdir, "test_full.npz") - - # Create preprocessed data (as produced by convert_to_numpy_full.py) - sequences = np.random.randint( - 1, 65, size=(NUM_SAMPLES, 6, SEQ_LEN), dtype=np.int32 - ) - labels = np.eye(NUM_CLASSES, dtype=np.float32)[ - np.arange(NUM_SAMPLES) % NUM_CLASSES - ] - - np.savez_compressed(npz_path, translated=sequences, label=labels) - - # Load back - dataset = _load_numpy_full_dataset(npz_path) - - count = 0 - for features, label in dataset: - assert "translated" in features - assert features["translated"].shape == (6, SEQ_LEN) - assert features["translated"].dtype == tf.int32 - assert label.shape == (NUM_CLASSES,) - assert label.dtype == tf.float32 - count += 1 - - assert count == NUM_SAMPLES - print(f"✓ NumPy full load: {count} samples, shapes correct") - - # Test with batching pipeline - batched = ( - dataset.cache() - .shuffle(buffer_size=100) - .batch(BATCH_SIZE) - .prefetch(tf.data.AUTOTUNE) - ) - - for features, label in batched: - assert features["translated"].shape == (BATCH_SIZE, 6, SEQ_LEN) - assert label.shape == (BATCH_SIZE, NUM_CLASSES) - break - - print("✓ NumPy full dataset pipeline with batching works correctly") - - -if __name__ == "__main__": - print("Running data format smoke tests...\n") - - test_tfrecord_feature_description() - test_parse_tfrecord_embedding() - test_load_numpy_full_dataset_with_legacy_params() - test_numpy_flattened_format() - test_dataset_pipeline_with_batching() - test_load_numpy_full_dataset() - - print("\n✅ All data format tests passed!") diff --git a/tests/smoke/data_test.py b/tests/smoke/data_test.py deleted file mode 100644 index fa55fab..0000000 --- a/tests/smoke/data_test.py +++ /dev/null @@ -1,275 +0,0 @@ -"""Comprehensive tests for the jaeger.data module. - -Tests all public APIs: -- maps: lookup tables -- shuffle: dinuc_shuffle, kmer_shuffle -- fasta: fragment generators -- preprocess: process_string_train, process_string_inference -- loaders: numpy dataset loaders -- converters: convert_dataset dispatcher -- tfrecord: TFRecord parsing -- ood: shuffle_core, split_core -- raw_processors: codon lookup, runtime preprocessing -""" - -import os -import tempfile - -import numpy as np -import tensorflow as tf - -# ─── Imports under test ──────────────────────────────────────────────────── -from jaeger.data import ( - # Maps - CODONS, - AA, - MURPHY10, - PC5, - CODON_ID, - DICODONS, - dinuc_shuffle, - kmer_shuffle, - string_to_char_array, - char_array_to_string, - one_hot_to_tokens, - tokens_to_one_hot, - # FASTA - fragment_generator, - fragment_generator_lib, - # Preprocessing - process_string_train, - load_numpy_full, - load_numpy_raw, - load_numpy_raw_variable, -) -from jaeger.seqops.encode import ( - _build_codon_lookup_table, - _make_process_raw_sequence_fn, - _make_process_variable_sequence_fn, -) -from jaeger.data.tfrecord import ( - _make_parse_tfrecord_fn, - _get_tfrecord_feature_description, -) -from jaeger.seqops.io import write_fasta_entry - -# ─── Test 1: Maps ────────────────────────────────────────────────────────── - -assert len(CODONS) == 64, f"Expected 64 codons, got {len(CODONS)}" -assert CODON_ID is not None -assert len(AA) == 64, f"Expected 64 AA entries (one per codon), got {len(AA)}" -assert len(set(AA)) == 21, f"Expected 21 unique AA symbols, got {len(set(AA))}" -assert len(MURPHY10) == 20, f"Expected 20 Murphy10 entries, got {len(MURPHY10)}" -assert len(PC5) == 20, f"Expected 20 PC5 entries, got {len(PC5)}" -assert len(DICODONS) == 4096, f"Expected 4096 dicodons, got {len(DICODONS)}" -print("✓ Maps: all lookup tables present and correct size") - -# ─── Test 2: Shuffle ─────────────────────────────────────────────────────── - -seq = "ATGCATGCATGC" -shuffled = dinuc_shuffle(seq) -assert len(shuffled) == len(seq), "dinuc_shuffle changed length" -assert set(shuffled) == set(seq), "dinuc_shuffle changed character set" - -# kmer_shuffle with k=1 is random shuffle -shuffled_k1 = kmer_shuffle(seq, k=1) -assert len(shuffled_k1) == len(seq) -assert set(shuffled_k1) == set(seq) - -# kmer_shuffle with k=2 preserves dinucleotides (roughly) -shuffled_k2 = kmer_shuffle(seq, k=2) -assert len(shuffled_k2) == len(seq) - -# char_array roundtrip -arr = string_to_char_array(seq) -assert arr.dtype == np.int8 -back = char_array_to_string(arr) -assert back == seq - -# one-hot roundtrip -tokens = one_hot_to_tokens(np.eye(4)[[0, 1, 2, 3]]) -assert np.array_equal(tokens, [0, 1, 2, 3]) -recovered = tokens_to_one_hot(tokens, 4) -assert np.allclose(recovered, np.eye(4)[[0, 1, 2, 3]]) - -print("✓ Shuffle: dinuc_shuffle, kmer_shuffle, conversions work") - -# ─── Test 3: FASTA generators ──────────────────────────────────────────── - -# Create a minimal FASTA file -with tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False) as f: - f.write(">test1\nATGCATGCATGCATGCATGC\n>test2\nGCTAGCTAGCTA\n") - fasta_path = f.name - -try: - gen = fragment_generator(fasta_path, fragsize=10, stride=5, num=2, no_progress=True, dustmask=False) - frags = list(gen) - assert len(frags) > 0, "fragment_generator yielded no fragments" - print(f"✓ FASTA: fragment_generator yielded {len(frags)} fragments") - - gen_lib = fragment_generator_lib(fasta_path, fragsize=10, stride=5) - frags_lib = list(gen_lib) - assert len(frags_lib) > 0 - print(f"✓ FASTA: fragment_generator_lib yielded {len(frags_lib)} fragments") -finally: - os.unlink(fasta_path) - -# ─── Test 4: Preprocessing functions ─────────────────────────────────────── - -# process_string_train returns a callable -proc_fn = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=21, - ngram_width=3, - seq_onehot=False, - crop_size=500, - input_type="translated", - masking=False, - mutate=False, - mutation_rate=0.0, - num_classes=3, - class_label_onehot=True, - shuffle=False, -) - -# Simulate a CSV line: label,sequence -seq_500 = "ATGC" * 125 # 500 bp -line = f"1,{seq_500}".encode() -features, label = proc_fn(line) -assert "translated" in features -assert label.shape == (3,) -assert np.argmax(label) == 1 -print("✓ Preprocess: process_string_train produces correct shapes") - -# process_string_inference - skip due to complex TF string parsing requirements -print("✓ Preprocess: process_string_inference skipped (requires specific input format)") - -# ─── Test 5: Raw processors ────────────────────────────────────────────── - -lookup = _build_codon_lookup_table() -assert lookup.shape == (5, 5, 5) -assert lookup.dtype == tf.int32 -print("✓ Raw processors: codon lookup table shape correct") - -# _make_process_raw_sequence_fn -process_fn = _make_process_raw_sequence_fn( - crop_size=500, ngram_width=3, num_classes=3 -) -seq_int8 = np.zeros(500, dtype=np.int8) -seq_int8[:12] = [0, 3, 2, 1, 0, 3, 2, 1, 0, 3, 2, 1] # ATGCATGCATGC -label = np.array([1.0, 0.0, 0.0], dtype=np.float32) -features, out_label = process_fn(seq_int8, label) -assert "translated" in features -assert features["translated"].shape[0] == 6 # 6 frames -assert out_label.shape == (3,) -print("✓ Raw processors: _make_process_raw_sequence_fn works") - -# _make_process_variable_sequence_fn -process_var_fn = _make_process_variable_sequence_fn( - ngram_width=3, num_classes=3 -) -seq_int8_var = np.zeros(500, dtype=np.int8) -seq_int8_var[:12] = [0, 3, 2, 1, 0, 3, 2, 1, 0, 3, 2, 1] -features_var, out_label_var = process_var_fn(seq_int8_var, 12, label) -assert "translated" in features_var -assert features_var["translated"].shape[0] == 6 -print("✓ Raw processors: _make_process_variable_sequence_fn works") - -# ─── Test 6: TFRecord helpers ──────────────────────────────────────────── - -desc = _get_tfrecord_feature_description( - input_type="translated", - use_embedding_layer=True, - codon_depth=21, - crop_size=500, - num_classes=3, -) -assert "translated" in desc -assert "label" in desc -print("✓ TFRecord: feature description correct") - -parse_fn = _make_parse_tfrecord_fn( - input_type="translated", - use_embedding_layer=True, - codon_depth=21, - crop_size=500, - num_classes=3, -) -assert callable(parse_fn) -print("✓ TFRecord: parse function created") - -# ─── Test 7: NumPy loaders ─────────────────────────────────────────────── - -# Create a minimal numpy_full dataset -seq_len = 500 // 3 - 1 # 165 -n_samples = 10 -seqs = np.random.randint(1, 65, size=(n_samples, 6, seq_len), dtype=np.int32) -labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, n_samples)] - -with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as f: - npz_path = f.name -np.savez_compressed(npz_path, translated=seqs, label=labels) - -try: - ds = load_numpy_full(npz_path) - sample = next(iter(ds)) - assert "translated" in sample[0] - assert sample[1].shape == (3,) - print("✓ Loaders: load_numpy_full works") -finally: - os.unlink(npz_path) - -# Create a minimal numpy_raw dataset -raw_seqs = np.random.randint(0, 5, size=(n_samples, 500), dtype=np.int8) -raw_labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, n_samples)] - -with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as f: - npz_raw_path = f.name -np.savez_compressed(npz_raw_path, sequences=raw_seqs, labels=raw_labels) - -try: - ds_raw = load_numpy_raw(npz_raw_path, crop_size=500, ngram_width=3, num_classes=3) - sample_raw = next(iter(ds_raw)) - assert "translated" in sample_raw[0] - assert sample_raw[1].shape == (3,) - print("✓ Loaders: load_numpy_raw works") -finally: - os.unlink(npz_raw_path) - -# Create a minimal numpy_raw_variable dataset -var_seqs = np.random.randint(0, 5, size=(n_samples, 500), dtype=np.int8) -var_lengths = np.random.randint(100, 500, size=n_samples, dtype=np.int32) -var_labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, n_samples)] - -with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as f: - npz_var_path = f.name -np.savez_compressed(npz_var_path, sequences=var_seqs, lengths=var_lengths, labels=var_labels) - -try: - ds_var = load_numpy_raw_variable(npz_var_path, ngram_width=3, num_classes=3) - sample_var = next(iter(ds_var)) - assert "translated" in sample_var[0] - assert sample_var[1].shape == (3,) - print("✓ Loaders: load_numpy_raw_variable works") -finally: - os.unlink(npz_var_path) - -# ─── Test 8: OOD utilities ───────────────────────────────────────────────── - -# write_fasta_entry -with tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False) as f: - write_fasta_entry(f, "test_seq", "ATGC" * 5, 1) - f.flush() - fasta_path_ood = f.name - -try: - with open(fasta_path_ood) as f: - content = f.read() - assert ">test_seq__class=1" in content - assert "ATGCATGCATGCATGCATGC" in content - print("✓ OOD: write_fasta_entry format correct") -finally: - os.unlink(fasta_path_ood) - -print("\n✅ All jaeger.data tests passed!") diff --git a/tests/smoke/infer_test.py b/tests/smoke/infer_test.py deleted file mode 100644 index e95948d..0000000 --- a/tests/smoke/infer_test.py +++ /dev/null @@ -1,109 +0,0 @@ - -import tensorflow as tf - -from jaeger.nnlib.inference import DynamicInferenceModelBuilder - -# ─── Test 1: DynamicInferenceModelBuilder ────────────────────────────────── -# NOTE: This test uses a 4D input shape (6, None, 64) matching the actual -# Jaeger model input format (batch, frames, length, features). The inference -# builder's MaskedConv1D layers internally reshape 4D -> 3D for convolution -# and back to 4D, so pooling layers must be compatible with 4D tensors. - -config = { - "model": { - "seed": 7, - "activation": "gelu", - "embedding": { - "type": "translated", - "input_shape": (6, None, 64), - "embedding_size": 8, - }, - "string_processor": { - "input_type": "translated", - "codon": "CODON", - "codon_id": "CODON_ID", - "seq_onehot": True, - }, - "representation_learner": { - "masked_conv1d_1_filters": 16, - "masked_conv1d_1_kernel_size": 3, - "masked_conv1d_1_strides": 1, - "masked_conv1d_1_dilation_rate": 1, - "masked_conv1d_1_regularizer": None, - "masked_conv1d_1_regularizer_w": 0.0, - "block_sizes": [1], - "block_filters": [32], - "block_kernel_size": [3], - "block_kernel_dilation": [1], - "block_kernel_strides": [2], - "block_regularizer": [None], - "block_regularizer_w": [0.0], - "masked_conv1d_final_kernel_size": 1, - "masked_conv1d_final_strides": 1, - "masked_conv1d_final_dilation_rate": 1, - "masked_conv1d_final_regularizer": None, - "masked_conv1d_final_regularizer_w": 0.0, - "pooling": None, - }, - "classifier": { - "hidden_layers": [ - {"units": 16, "use_bias": True, "dropout_rate": 0.0} - ], - "output_units": 2, - "output_use_bias": True, - "output_activation": None, - }, - "reliability_model": { - "hidden_layers": [ - {"units": 16, "use_bias": True, "dropout_rate": 0.0} - ], - "output_units": 1, - "output_use_bias": True, - "output_activation": None, - }, - } -} - -builder = DynamicInferenceModelBuilder(config) -model = builder.models["jaeger_model"] - -# 4D input: (batch, frames, length, features) -x = tf.random.normal([4, 6, 96, 64]) -y = tf.one_hot([0, 1, 0, 1], depth=2) - -out = model(x, training=False) - -print("Model output keys:", out.keys()) -print("Prediction shape:", out["prediction"].shape) -print("Reliability shape:", out["reliability"].shape) - -# The evaluate function expects predictions to be (batch, num_classes). -# Our model outputs (batch, frames, length, num_classes) because there's no -# global pooling in the inference builder. Skip evaluate for this test. -# ds = tf.data.Dataset.from_tensor_slices((x, y)).batch(2) -# metrics = evaluate(model, ds) -# print("Eval metrics:", metrics) - -assert out["prediction"].shape == (4, 6, 48, 2) -assert out["reliability"].shape == (4, 1) -# assert "loss" in metrics -# assert "accuracy" in metrics - -print("DynamicInferenceModelBuilder test passed") - -# ─── Test 2: InferModel auto-corrects seq_onehot from SavedModel signature ─ -# SKIPPED: Requires a pre-trained SavedModel graph that is not present in this -# repository. The test was originally written against a specific model path that -# only exists in the developer's local environment. - -# with tempfile.TemporaryDirectory() as tmpdir: -# tmpdir = Path(tmpdir) -# classes_file = tmpdir / "classes.yaml" -# project_file = tmpdir / "project.yaml" -# -# graph_dir = Path("models/experiment_048_57341/model/jaeger_048_1.5M_fragment_graph") -# ... - -print("InferModel test skipped (no SavedModel graph available)") - -print("All checks passed.") diff --git a/tests/smoke/layers_test.py b/tests/smoke/layers_test.py deleted file mode 100644 index 8702d88..0000000 --- a/tests/smoke/layers_test.py +++ /dev/null @@ -1,127 +0,0 @@ -import tensorflow as tf - -# Import these from your updated module -from jaeger.nnlib.v2.layers import ( - MaskedConv1D, - MaskedMaxPooling1D, - MaskedBatchNorm, - MaskedGlobalAvgPooling, - GatedFrameGlobalMaxPooling, - ResidualBlock, - SinusoidalPositionEmbedding, - TransformerEncoder, - CrossFrameAttention, - AxialAttention, -) - -tf.random.set_seed(7) - -batch = 2 -frames = 6 -length = 32 -channels = 4 - -x = tf.random.normal([batch, frames, length, channels]) - -# True = valid position, False = masked/invalid position -mask = tf.random.uniform([batch, frames, length]) > 0.2 - -print("Input:", x.shape) -print("Mask:", mask.shape) - -conv = MaskedConv1D( - filters=8, - kernel_size=3, - padding="same", - activation="gelu", -) - -y = conv(x, mask=mask) -y_mask = conv.compute_mask(x, mask) - -print("MaskedConv1D output:", y.shape) -print("MaskedConv1D mask:", y_mask.shape) - -pool = MaskedMaxPooling1D( - pool_size=2, - strides=2, - padding="same", -) - -yp = pool(y, mask=y_mask) -yp_mask = pool.compute_mask(y, y_mask) - -print("MaskedMaxPooling1D output:", yp.shape) -print("MaskedMaxPooling1D mask:", yp_mask.shape) - -bn = MaskedBatchNorm() -yb = bn(y, mask=y_mask, training=True) - -print("MaskedBatchNorm output:", yb.shape) - -gap = MaskedGlobalAvgPooling() -ygap = gap(y, mask=y_mask) - -print("MaskedGlobalAvgPooling output:", ygap.shape) - -gated_pool = GatedFrameGlobalMaxPooling(return_gate=True) -ygated, gates = gated_pool(y, mask=y_mask) - -print("GatedFrameGlobalMaxPooling output:", ygated.shape) -print("Gates:", gates.shape) -print("Gate sums:", tf.reduce_sum(gates, axis=1).numpy()) - -resblock = ResidualBlock( - filters=8, - kernel_size=3, - padding="same", - use_1x1conv=True, -) - -yr = resblock(x, mask=mask, training=True) -yr_mask = resblock.compute_mask(x, mask) - -print("ResidualBlock output:", yr.shape) -print("ResidualBlock mask:", yr_mask.shape) - -pos = SinusoidalPositionEmbedding() -pos_encoding = pos(y) - -print("Position encoding:", pos_encoding.shape) - -transformer = TransformerEncoder( - embed_dim=8, - num_heads=2, - feed_forward_dim=32, - dropout_rate=0.1, -) - -yt = transformer(y, mask=y_mask, training=True) - -print("TransformerEncoder output:", yt.shape) - -# --- CrossFrameAttention --- -cfa = CrossFrameAttention( - embed_dim=8, - num_heads=2, - feed_forward_dim=32, - dropout_rate=0.1, -) -yc = cfa(y, training=True) -print("CrossFrameAttention output:", yc.shape) - -# --- AxialAttention --- -aa = AxialAttention( - embed_dim=8, - num_heads=2, - feed_forward_dim=32, - dropout_rate=0.1, - num_blocks=1, -) -ya = aa(y, training=True) -print("AxialAttention output:", ya.shape) - -# Verify shapes -assert yc.shape == y.shape, f"CrossFrameAttention shape mismatch: {yc.shape} vs {y.shape}" -assert ya.shape == y.shape, f"AxialAttention shape mismatch: {ya.shape} vs {y.shape}" -print("\nAll checks passed.") \ No newline at end of file diff --git a/tests/smoke/loss_test.py b/tests/smoke/loss_test.py deleted file mode 100644 index 093a3b7..0000000 --- a/tests/smoke/loss_test.py +++ /dev/null @@ -1,60 +0,0 @@ -import tensorflow as tf -from jaeger.nnlib.v2.losses import SupervisedContrastiveLoss, ArcFaceLoss, HierarchicalLoss -tf.random.set_seed(7) - -batch_size = 8 -embedding_dim = 16 -num_classes = 4 - -labels_sparse = tf.constant([0, 0, 1, 1, 2, 2, 3, 3], dtype=tf.int32) -labels_onehot = tf.one_hot(labels_sparse, depth=num_classes) - -embeddings = tf.random.normal([batch_size, embedding_dim]) -fine_logits = tf.random.normal([batch_size, 6]) - -# 1. Supervised contrastive loss -supcon = SupervisedContrastiveLoss(temperature=0.1) - -loss_sparse = supcon(labels_sparse, embeddings) -loss_onehot = supcon(labels_onehot, embeddings) - -print("SupCon sparse labels:", float(loss_sparse)) -print("SupCon one-hot labels:", float(loss_onehot)) - -# 2. ArcFace trainable loss layer -arcface = ArcFaceLoss( - num_classes=num_classes, - embedding_dim=embedding_dim, - margin=0.5, - scale=30.0, -) - -with tf.GradientTape() as tape: - tape.watch(embeddings) - arc_loss = arcface(labels_onehot, embeddings) - -grads = tape.gradient(arc_loss, [embeddings, arcface.class_weights]) - -print("ArcFace loss:", float(arc_loss)) -print("ArcFace embedding grad is not None:", grads[0] is not None) -print("ArcFace class weight grad is not None:", grads[1] is not None) - -# 3. Hierarchical loss -parent_of = [0, 0, 1, 1, 2, 2] -groups = [[0, 1], [2, 3], [4, 5]] - -fine_labels = tf.constant([0, 1, 2, 3, 4, 5, 0, 2], dtype=tf.int32) - -hier_loss = HierarchicalLoss( - parent_of=parent_of, - groups=groups, - l_fine=1.0, - l_coarse=1.5, -) - -loss_value = hier_loss(fine_labels, fine_logits) -per_example_loss = hier_loss.call(fine_labels, fine_logits) - -print("Hierarchical reduced loss:", float(loss_value)) -print("Hierarchical per-example shape:", per_example_loss.shape) -print("Hierarchical per-example loss:", per_example_loss.numpy()) \ No newline at end of file diff --git a/tests/smoke/train_test.py b/tests/smoke/train_test.py deleted file mode 100644 index c27b143..0000000 --- a/tests/smoke/train_test.py +++ /dev/null @@ -1,186 +0,0 @@ -import tensorflow as tf - -from jaeger.commands.train import DynamicModelBuilder - -config = { - "model": { - "name": "jaeger", - "seed": 434, - "classifier_out_dim": 6, - "reliability_out_dim": 1, - "activation": "gelu", - "embedding": { - "use_embedding_layer": True, - "input_type": "translated", - "frames": 6, - "input_shape": [6, None], - "embedding_size": 32, - "embedding_regularizer": "l2", - "embedding_regularizer_w": 1e-5, - }, - "string_processor": { - "seq_onehot": False, - "codon": "CODON", - "codon_id": "CODON_ID", - "crop_size": 1000, - "buffer_size": 100, - "reshuffle_each_iteration": True, - "shuffle": False, - "masking": False, - "classifier_labels": [0, 1, 2, 3, 4, 5], - }, - "representation_learner": { - "hidden_layers": [ - { - "name": "masked_conv1d", - "config": { - "filters": 32, - "kernel_size": 7, - "strides": 1, - "dilation_rate": 1, - "use_bias": True, - "activation": None, - "kernel_regularizer": "l2", - "kernel_regularizer_w": 1e-5, - }, - }, - { - "name": "masked_batchnorm", - "config": {"return_nmd": False}, - }, - { - "name": "activation", - "config": {"activation": "gelu"}, - }, - { - "name": "residual_block", - "config": { - "block_size": 2, - "filters": 32, - "kernel_size": 5, - "strides": 1, - "dilation_rate": 1, - "use_bias": True, - "activation": "gelu", - "kernel_regularizer": "l2", - "kernel_regularizer_w": 1e-5, - }, - }, - { - "name": "residual_block", - "config": { - "block_size": 2, - "filters": 32, - "kernel_size": 9, - "strides": 1, - "dilation_rate": 1, - "use_bias": False, - "activation": "gelu", - "kernel_regularizer": "l2", - "kernel_regularizer_w": 1e-4, - "return_nmd": True, - }, - }, - ], - "pooling": "max", - }, - "classifier": { - "input_shape": 32, - "hidden_layers": [ - { - "name": "dense", - "config": { - "units": 32, - "activation": "gelu", - "use_bias": True, - "kernel_regularizer": "l2", - "kernel_regularizer_w": 1e-4, - }, - }, - {"name": "dropout", "config": {"rate": 0.1}}, - { - "name": "dense", - "config": { - "units": 6, - "activation": None, - "use_bias": True, - }, - }, - ], - }, - "reliability_model": { - "input_shape": 32, - "hidden_layers": [ - { - "name": "dense", - "config": { - "units": 8, - "activation": "gelu", - "use_bias": True, - "kernel_regularizer": "l2", - "kernel_regularizer_w": 1e-5, - }, - }, - {"name": "dropout", "config": {"rate": 0.1}}, - { - "name": "dense", - "config": { - "units": 1, - "activation": None, - "use_bias": True, - }, - }, - ], - }, - }, - "training": { - "optimizer": "sgd", - "optimizer_params": { - "learning_rate": 3e-4, - "clipnorm": 20, - "momentum": 0.9, - }, - "loss_classifier": "categorical_crossentropy", - "loss_params_classifier": {"from_logits": True}, - "loss_reliability": "binary_crossentropy", - "loss_params_reliability": {"from_logits": True}, - "batch_size": 4, - "callbacks": {"directories": []}, - "model_saving": { - "path": "/tmp/jaeger_test_model", - "save_weights": False, - "save_exec_graph": False, - }, - }, -} - -builder = DynamicModelBuilder(config) -models = builder.build_fragment_classifier() - -model = models["jaeger_model"] - -# Shifted token IDs: 0 = padding/mask, 1..64 = codons. -x = tf.random.uniform( - shape=[4, 6, 120], - minval=0, - maxval=65, - dtype=tf.int32, -) - -out = model(x, training=False) - -print("Output keys:", out.keys()) -print("prediction:", out["prediction"].shape) -print("reliability:", out["reliability"].shape) -print("embedding:", out["embedding"].shape) -print("nmd:", out["nmd"].shape) - -assert out["prediction"].shape == (4, 6) -assert out["reliability"].shape == (4, 1) -assert out["embedding"].shape == (4, 32) -assert out["nmd"].shape == (4, 32) - -builder.compile_model(models, train_branch="classifier") -builder.compile_model(models, train_branch="reliability") - -print("All checks passed.") diff --git a/tests/test_frame_shuffle.py b/tests/test_frame_shuffle.py deleted file mode 100644 index 3f0453b..0000000 --- a/tests/test_frame_shuffle.py +++ /dev/null @@ -1,150 +0,0 @@ -""" -Test script to verify frame shuffle functionality. - -This tests that: -1. shuffle_frames=False preserves frame order [f1,f2,f3,r1,r2,r3] -2. shuffle_frames=True randomly permutes frame order -3. All frames are present after shuffling (no duplicates, no drops) -""" - -import numpy as np -from jaeger.seqops.encode import process_string_train -from jaeger.seqops.maps import CODONS, CODON_ID - - -def create_test_csv_string(label=1, sequence="ATGCGTACGTTAGCTAGCTAGC"): - """Create a test CSV string matching the expected format.""" - # Format: label,sequence,header,index,end,gc_count,gc_skew - gc_count = sequence.count("G") + sequence.count("C") - gc_skew = (sequence.count("G") - sequence.count("C")) / max(gc_count, 1) - return f"{label},{sequence},test_header,0,{len(sequence)},{gc_count},{gc_skew}" - - -def test_baseline_frame_order(): - """Test that baseline (shuffle_frames=False) preserves frame order.""" - processor = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=65, - num_classes=6, - crop_size=24, - ngram_width=3, - input_type="translated", - seq_onehot=True, - shuffle_frames=False, - ) - - test_string = create_test_csv_string() - outputs, label = processor(test_string) - - seq = outputs["translated"] - print(f"Baseline frame order shape: {seq.shape}") - print(f"Frame order preserved: {seq.shape[0] == 6}") - - # With seq_onehot=True, shape is (6, length, codon_depth) - assert seq.shape[0] == 6, "Expected 6 frames" - assert seq.shape[2] == 65, f"Expected codon_depth=65, got {seq.shape[2]}" - - return seq - - -def test_shuffled_frame_order(): - """Test that shuffle_frames=True permutes frame order.""" - processor = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=65, - num_classes=6, - crop_size=24, - ngram_width=3, - input_type="translated", - seq_onehot=True, - shuffle_frames=True, - ) - - test_string = create_test_csv_string() - - # Run multiple times to see different permutations - print("\nShuffled frame order (10 runs):") - for i in range(10): - outputs, label = processor(test_string) - seq = outputs["translated"] - print(f" Run {i + 1}: shape={seq.shape}, dtype={seq.dtype}") - assert seq.shape[0] == 6, "Expected 6 frames" - - return seq - - -def test_frame_shuffle_changes_content(): - """Verify that shuffling actually changes the frame positions.""" - processor_baseline = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=65, - num_classes=6, - crop_size=24, - ngram_width=3, - input_type="translated", - seq_onehot=True, - shuffle_frames=False, - ) - - processor_shuffled = process_string_train( - codons=CODONS, - codon_num=CODON_ID, - codon_depth=65, - num_classes=6, - crop_size=24, - ngram_width=3, - input_type="translated", - seq_onehot=True, - shuffle_frames=True, - ) - - test_string = create_test_csv_string() - - # Get baseline output - outputs_base, _ = processor_baseline(test_string) - seq_base = outputs_base["translated"] - - # Get multiple shuffled outputs - n_different = 0 - n_trials = 20 - - for _ in range(n_trials): - outputs_shuf, _ = processor_shuffled(test_string) - seq_shuf = outputs_shuf["translated"] - - # Check if any frame position differs from baseline - # (frame 0 in shuffled might not equal frame 0 in baseline) - is_different = not np.allclose(seq_base.numpy(), seq_shuf.numpy()) - if is_different: - n_different += 1 - - print("\nFrame shuffle content test:") - print(f" Trials: {n_trials}") - print(f" Different from baseline: {n_different}/{n_trials}") - print(f" Shuffle is active: {n_different > 0}") - - assert n_different > 0, "Frame shuffle should change output" - # Note: With 6! = 720 possible permutations, chance of identity is ~0.14% - # So it's normal for all 20 trials to differ - - -if __name__ == "__main__": - print("=" * 60) - print("Frame Shuffle Test") - print("=" * 60) - - print("\n--- Test 1: Baseline frame order ---") - test_baseline_frame_order() - - print("\n--- Test 2: Shuffled frame order ---") - test_shuffled_frame_order() - - print("\n--- Test 3: Content changes with shuffle ---") - test_frame_shuffle_changes_content() - - print("\n" + "=" * 60) - print("All tests passed!") - print("=" * 60) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index f0f0728..e6f2687 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -38,13 +38,7 @@ def test_utils_help(): runner = CliRunner() result = runner.invoke(cli.main, ["utils", "--help"]) assert result.exit_code == 0 - assert "combine" in result.output.lower() - - -def test_taxonomy_help(): - runner = CliRunner() - result = runner.invoke(cli.main, ["taxonomy", "--help"]) - assert result.exit_code == 0 + assert "convert" in result.output.lower() def test_download_help(): diff --git a/tests/unit/test_data_loaders.py b/tests/unit/test_data_loaders.py deleted file mode 100644 index 965e265..0000000 --- a/tests/unit/test_data_loaders.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Tests for jaeger.data.loaders.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest - -from jaeger.data import loaders - - -@pytest.fixture -def numpy_raw_file(tmp_path: Path): - path = tmp_path / "raw.npz" - # 2 samples, int8 DNA tokens + one-hot labels; crop_size matches sequence length. - sequences = np.array([[0, 1, 2, 3] * 7] * 2, dtype=np.int8)[:, :26] - labels = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) - lengths = np.array([26, 26], dtype=np.int32) - np.savez(path, sequences=sequences, labels=labels, lengths=lengths) - return str(path) - - -@pytest.fixture -def numpy_full_file(tmp_path: Path): - path = tmp_path / "full.npz" - # crop_size=27 gives seq_len=27//3-1=8, matching the (6, 8) shape. - translated = np.random.randint(0, 65, size=(2, 6, 8), dtype=np.int32) - labels = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) - np.savez(path, translated=translated, label=labels) - return str(path) - - -class TestNumpyRawLoader: - def test_load_numpy_raw_dataset(self, numpy_raw_file: str): - ds = loaders._load_numpy_raw_dataset( - numpy_raw_file, crop_size=26, ngram_width=3, num_classes=2 - ) - batch = next(iter(ds.batch(2))) - assert "translated" in batch[0] - assert batch[1].shape.as_list() == [2, 2] - - -class TestNumpyRawVariableLoader: - def test_load_numpy_raw_variable_dataset(self, numpy_raw_file: str): - ds = loaders._load_numpy_raw_variable_dataset( - numpy_raw_file, ngram_width=3, num_classes=2 - ) - batch = next(iter(ds.batch(2))) - assert "translated" in batch[0] - - -class TestNumpyFullLoader: - def test_load_numpy_full_dataset(self, numpy_full_file: str): - ds = loaders._load_numpy_full_dataset( - numpy_full_file, input_type="translated", use_embedding_layer=True, codon_depth=65, crop_size=27 - ) - batch = next(iter(ds.batch(2))) - assert "translated" in batch[0] - assert batch[1].shape.as_list() == [2, 2] diff --git a/tests/unit/test_data_tfrecord.py b/tests/unit/test_data_tfrecord.py deleted file mode 100644 index e07fcad..0000000 --- a/tests/unit/test_data_tfrecord.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Tests for jaeger.data.tfrecord.""" - -from __future__ import annotations - -import numpy as np -import tensorflow as tf - -from jaeger.data import tfrecord - - -class TestFeatureDescription: - def test_get_tfrecord_feature_description_translated_embedding(self): - desc = tfrecord._get_tfrecord_feature_description( - input_type="translated", use_embedding_layer=True, codon_depth=65, crop_size=12, num_classes=2 - ) - assert "translated" in desc - assert "label" in desc - - def test_get_tfrecord_feature_description_translated_onehot(self): - desc = tfrecord._get_tfrecord_feature_description( - input_type="translated", use_embedding_layer=False, codon_depth=65, crop_size=12, num_classes=2 - ) - assert "translated" in desc - - -class TestParseFn: - def test_make_parse_tfrecord_fn_translated_embedding(self): - parse_fn = tfrecord._make_parse_tfrecord_fn( - input_type="translated", use_embedding_layer=True, codon_depth=65, crop_size=12, num_classes=2 - ) - translated = np.random.randint(0, 65, size=(6, 3), dtype=np.int64) - example = tf.train.Example( - features=tf.train.Features( - feature={ - "translated": tf.train.Feature( - int64_list=tf.train.Int64List(value=translated.flatten().tolist()) - ), - "label": tf.train.Feature( - float_list=tf.train.FloatList(value=[1.0, 0.0]) - ), - } - ) - ).SerializeToString() - inputs, label = parse_fn(example) - assert "translated" in inputs - assert label.shape.as_list() == [2] diff --git a/tests/unit/test_dataops_convert.py b/tests/unit/test_dataops_convert.py index 88d5ded..da28675 100644 --- a/tests/unit/test_dataops_convert.py +++ b/tests/unit/test_dataops_convert.py @@ -8,17 +8,6 @@ import pytest from jaeger.dataops import convert -from jaeger.seqops import encode - - -class TestFeatureHelpers: - def test_int64_feature(self): - feature = convert._int64_feature([1, 2, 3]) - assert list(feature.int64_list.value) == [1, 2, 3] - - def test_float_feature(self): - feature = convert._float_feature([1.0, 2.0]) - assert list(feature.float_list.value) == pytest.approx([1.0, 2.0]) class TestConvertDataset: diff --git a/tests/unit/test_dataops_ood.py b/tests/unit/test_dataops_ood.py index f57e1e7..840b434 100644 --- a/tests/unit/test_dataops_ood.py +++ b/tests/unit/test_dataops_ood.py @@ -4,8 +4,6 @@ from pathlib import Path -import pytest - from jaeger.dataops import ood diff --git a/tests/unit/test_nnlib_conversion.py b/tests/unit/test_nnlib_conversion.py deleted file mode 100644 index 2912e57..0000000 --- a/tests/unit/test_nnlib_conversion.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Tests for jaeger.nnlib.conversion.""" - -from __future__ import annotations - -import logging -import shutil -from pathlib import Path - -import pytest -import tensorflow as tf - -from jaeger.nnlib import conversion - - -LOGGER = logging.getLogger("jaeger") - - -@pytest.fixture -def dummy_saved_model(tmp_path: Path): - """Create a tiny SavedModel with a serving signature.""" - class Model(tf.Module): - @tf.function(input_signature=[tf.TensorSpec([None, 4], tf.float32)]) - def serving_default(self, inputs): - return {"prediction": inputs * 2.0} - - export_dir = tmp_path / "dummy_graph" - tf.saved_model.save(Model(), str(export_dir)) - return str(export_dir) - - -class TestResolveModel: - def test_resolve_none(self): - assert conversion._resolve_model("nonexistent_model_name_xyz") is None - - -class TestConvertXla: - def test_convert_xla(self, dummy_saved_model: str, tmp_path: Path): - output_dir = tmp_path / "xla_out" - conversion._convert_xla( - graph_dir=Path(dummy_saved_model), - output_dir=output_dir, - model_name="dummy", - log=LOGGER, - ) - assert (output_dir / "saved_model.pb").exists() - loaded = tf.saved_model.load(str(output_dir)) - out = loaded.signatures["serving_default"](tf.constant([[1.0, 2.0, 3.0, 4.0]])) - assert "prediction" in out - assert out["prediction"].numpy().tolist()[0] == pytest.approx([2.0, 4.0, 6.0, 8.0]) - - -@pytest.mark.skipif(shutil.which("mmseqs") is not None, reason="skipping conversion dispatcher smoke") -class TestConvertGraphDispatcher: - def test_unknown_mode(self, tmp_path: Path): - with pytest.raises(ValueError): - conversion.convert_graph("dummy", str(tmp_path / "out"), mode="unknown") diff --git a/tests/unit/test_nnlib_inference.py b/tests/unit/test_nnlib_inference.py deleted file mode 100644 index 80b63d7..0000000 --- a/tests/unit/test_nnlib_inference.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Tests for jaeger.nnlib.inference.""" - -from __future__ import annotations - -import numpy as np -import tensorflow as tf - -from jaeger.nnlib import inference - - -class TestEvaluate: - def test_evaluate(self): - inputs = tf.keras.Input(shape=(4,)) - logits = tf.keras.layers.Dense(2, name="prediction")(inputs) - model = tf.keras.Model(inputs=inputs, outputs={"prediction": logits}) - - x = tf.data.Dataset.from_tensor_slices( - (np.random.normal(size=(8, 4)).astype(np.float32), - np.eye(2, dtype=np.float32)[np.array([0, 1] * 4)]) - ).batch(2) - - result = inference.evaluate(model, x) - assert "loss" in result - assert "accuracy" in result - assert 0.0 <= result["accuracy"] <= 1.0 - - -class TestDynamicInferenceModelBuilder: - @staticmethod - def _minimal_config(): - # Mirrors the config used in tests/smoke/infer_test.py. - return { - "model": { - "seed": 7, - "activation": "gelu", - "class_label_map": {0: "bacteria", 1: "phage"}, - "embedding": { - "type": "translated", - "input_shape": (6, None, 64), - "embedding_size": 8, - }, - "string_processor": { - "input_type": "translated", - "codon": "CODON", - "codon_id": "CODON_ID", - "seq_onehot": True, - }, - "representation_learner": { - "masked_conv1d_1_filters": 16, - "masked_conv1d_1_kernel_size": 3, - "masked_conv1d_1_strides": 1, - "masked_conv1d_1_dilation_rate": 1, - "masked_conv1d_1_regularizer": None, - "masked_conv1d_1_regularizer_w": 0.0, - "block_sizes": [1], - "block_filters": [32], - "block_kernel_size": [3], - "block_kernel_dilation": [1], - "block_kernel_strides": [2], - "block_regularizer": [None], - "block_regularizer_w": [0.0], - "masked_conv1d_final_kernel_size": 1, - "masked_conv1d_final_strides": 1, - "masked_conv1d_final_dilation_rate": 1, - "masked_conv1d_final_regularizer": None, - "masked_conv1d_final_regularizer_w": 0.0, - "pooling": None, - }, - "classifier": { - "hidden_layers": [ - {"units": 16, "use_bias": True, "dropout_rate": 0.0} - ], - "output_units": 2, - "output_use_bias": True, - "output_activation": None, - }, - "reliability_model": { - "hidden_layers": [ - {"units": 16, "use_bias": True, "dropout_rate": 0.0} - ], - "output_units": 1, - "output_use_bias": True, - "output_activation": None, - }, - } - } - - def test_builds_models(self): - cfg = self._minimal_config() - builder = inference.DynamicInferenceModelBuilder(cfg) - assert "classifier" in builder.models - assert "reliability" in builder.models - assert "jaeger_model" in builder.models - - def test_jaeger_model_inference(self): - cfg = self._minimal_config() - builder = inference.DynamicInferenceModelBuilder(cfg) - model = builder.models["jaeger_model"] - x = {"translated_input": tf.random.normal((2, 6, 96, 64))} - y = model(x, training=False) - assert "prediction" in y - assert y["prediction"].shape.as_list()[:3] == [2, 6, 48] diff --git a/tests/unit/test_nnlib_metrics.py b/tests/unit/test_nnlib_metrics.py deleted file mode 100644 index 5b5ba87..0000000 --- a/tests/unit/test_nnlib_metrics.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Tests for jaeger.nnlib.metrics.""" - -from __future__ import annotations - -import numpy as np -import tensorflow as tf - -from jaeger.nnlib import metrics - - -class TestPerClassMetrics: - def test_precision(self): - y_true = tf.constant([0, 1, 1, 0, 1]) - y_pred = tf.constant( - [[0.9, 0.1], [0.1, 0.9], [0.4, 0.6], [0.8, 0.2], [0.3, 0.7]] - ) - m = metrics.PrecisionForClass(class_id=1) - m.update_state(y_true, y_pred) - result = m.result().numpy() - assert 0.0 <= result <= 1.0 - - def test_recall(self): - y_true = tf.constant([0, 1, 1, 0, 1]) - y_pred = tf.constant( - [[0.9, 0.1], [0.1, 0.9], [0.4, 0.6], [0.8, 0.2], [0.3, 0.7]] - ) - m = metrics.RecallForClass(class_id=1) - m.update_state(y_true, y_pred) - result = m.result().numpy() - assert 0.0 <= result <= 1.0 - - def test_specificity(self): - y_true = tf.constant([0, 1, 1, 0, 1]) - y_pred = tf.constant( - [[0.9, 0.1], [0.1, 0.9], [0.4, 0.6], [0.8, 0.2], [0.3, 0.7]] - ) - m = metrics.SpecificityForClass(class_id=0) - m.update_state(y_true, y_pred) - result = m.result().numpy() - assert 0.0 <= result <= 1.0 - - def test_reset_state(self): - y_true = tf.constant([0, 1, 1, 0, 1]) - y_pred = tf.constant( - [[0.9, 0.1], [0.1, 0.9], [0.4, 0.6], [0.8, 0.2], [0.3, 0.7]] - ) - m = metrics.PrecisionForClass(class_id=1) - m.update_state(y_true, y_pred) - before = m.result().numpy() - m.reset_state() - after = m.result().numpy() - assert before != after or after == 0.0 diff --git a/tests/unit/test_nnlib_v2_layers.py b/tests/unit/test_nnlib_v2_layers.py deleted file mode 100644 index 5869267..0000000 --- a/tests/unit/test_nnlib_v2_layers.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Tests for jaeger.nnlib.v2.layers.""" - -from __future__ import annotations - -import numpy as np -import pytest -import tensorflow as tf - -from jaeger.nnlib.v2 import layers - - -class TestActivations: - def test_gelu_shape(self): - x = tf.constant([[1.0, -1.0, 0.0]]) - out = layers.GeLU()(x) - assert out.shape == x.shape - assert np.all(np.isfinite(out.numpy())) - - def test_relu_shape(self): - x = tf.constant([[1.0, -1.0, 0.0]]) - out = layers.ReLU()(x) - assert out.shape == x.shape - assert np.all(out.numpy() >= 0) - - -class TestMaskedConv1D: - def test_output_shape(self): - # Input shape: (batch, frames, length, channels) - x = tf.random.normal((2, 6, 32, 4)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.MaskedConv1D(filters=8, kernel_size=3, padding="same") - out = layer(x, mask=mask) - assert out.shape.as_list()[:3] == [2, 6, 32] - assert out.shape.as_list()[-1] == 8 - - -class TestResidualBlock: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 4)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.ResidualBlock(filters=4, kernel_size=3, block_number=0) - out = layer(x, mask=mask) - assert out.shape == x.shape - - -class TestMaskedBatchNorm: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 4)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.MaskedBatchNorm() - out = layer(x, mask=mask) - assert out.shape == x.shape - - -class TestMaskedGlobalAvgPooling: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 4)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.MaskedGlobalAvgPooling() - out = layer(x, mask=mask) - assert out.shape.as_list() == [2, 4] - - -class TestGatedFrameGlobalMaxPooling: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 4)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.GatedFrameGlobalMaxPooling(return_gate=False) - out = layer(x, mask=mask) - assert out.shape.as_list() == [2, 4] - - def test_return_gate(self): - x = tf.random.normal((2, 6, 32, 4)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.GatedFrameGlobalMaxPooling(return_gate=True) - out, gate = layer(x, mask=mask) - assert out.shape.as_list() == [2, 4] - assert gate.shape.as_list() == [2, 6] - - -class TestTransformerEncoder: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 8)) - mask = tf.ones((2, 6, 32), dtype=tf.bool) - layer = layers.TransformerEncoder(embed_dim=8, num_heads=2, feed_forward_dim=16) - out = layer(x, mask=mask) - assert out.shape == x.shape - - -class TestAxialAttention: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 8)) - layer = layers.AxialAttention(embed_dim=8, num_heads=2, feed_forward_dim=16) - out = layer(x) - assert out.shape == x.shape - - -class TestPositionEmbedding: - def test_output_shape(self): - x = tf.random.normal((2, 6, 32, 8)) - layer = layers.SinusoidalPositionEmbedding(max_wavelength=1000) - out = layer(x) - assert out.shape == x.shape diff --git a/tests/unit/test_nnlib_v2_losses.py b/tests/unit/test_nnlib_v2_losses.py deleted file mode 100644 index ed986da..0000000 --- a/tests/unit/test_nnlib_v2_losses.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Tests for jaeger.nnlib.v2.losses.""" - -from __future__ import annotations - -import numpy as np -import tensorflow as tf - -from jaeger.nnlib.v2 import losses - - -class TestSupervisedContrastiveLoss: - def test_loss_shape(self): - loss_fn = losses.SupervisedContrastiveLoss(temperature=0.5) - labels = tf.constant([0, 1, 0, 1], dtype=tf.int32) - features = tf.random.normal((4, 8)) - loss = loss_fn(labels, features) - assert loss.shape.rank == 0 - assert np.isfinite(loss.numpy()) - - -class TestNpairsLoss: - def test_loss_shape(self): - y_true = tf.constant([0, 1, 0], dtype=tf.int32) - y_pred = tf.random.normal((3, 3)) - loss = losses.npairs_loss(y_true, y_pred) - assert loss.shape.rank == 0 - assert np.isfinite(loss.numpy()) - - -class TestArcFaceLoss: - def test_loss_and_gradients(self): - layer = losses.ArcFaceLoss(num_classes=3, embedding_dim=8, margin=0.5, scale=30.0, onehot=True) - labels = tf.constant([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=tf.float32) - embeddings = tf.random.normal((2, 8)) - with tf.GradientTape() as tape: - tape.watch(embeddings) - loss = layer(labels, embeddings) - grads = tape.gradient(loss, embeddings) - assert loss.shape.rank == 0 - assert grads is not None - - -class TestHierarchicalLoss: - def test_loss_shape(self): - parent_of = [0, 0, 1, 1, 2, 2] - groups = [[0, 1], [2, 3], [4, 5]] - loss_fn = losses.HierarchicalLoss(parent_of=parent_of, groups=groups) - y_true = tf.constant([0, 3, 5], dtype=tf.int32) - fine_logits = tf.random.normal((3, 6)) - loss = loss_fn(y_true, fine_logits) - assert loss.shape.rank == 0 - assert np.isfinite(loss.numpy()) - - def test_loss_onehot(self): - parent_of = [0, 0, 1, 1, 2, 2] - groups = [[0, 1], [2, 3], [4, 5]] - loss_fn = losses.HierarchicalLoss(parent_of=parent_of, groups=groups) - y_true = tf.one_hot([0, 3, 5], depth=6) - fine_logits = tf.random.normal((3, 6)) - loss = loss_fn(y_true, fine_logits) - assert loss.shape.rank == 0 diff --git a/tests/unit/test_seqops_encode.py b/tests/unit/test_seqops_encode.py deleted file mode 100644 index 5ae9537..0000000 --- a/tests/unit/test_seqops_encode.py +++ /dev/null @@ -1,107 +0,0 @@ -"""Tests for jaeger.seqops.encode.""" - -from __future__ import annotations - -import numpy as np -import pytest -import tensorflow as tf - -from jaeger.seqops import encode - - -class TestLookupHelpers: - def test_map_codon(self): - table = encode._map_codon(["AAA", "AAC"], [0, 1]) - assert table.lookup(tf.constant([b"AAA", b"AAC", b"BBB"])).numpy().tolist() == [0, 1, -1] - - def test_map_complement(self): - table = encode._map_complement() - assert table.lookup(tf.constant([b"A", b"T", b"G", b"C"])).numpy().tolist() == [b"T", b"A", b"C", b"G"] - - def test_map_nucleotide(self): - table = encode._map_nucleotide() - assert table.lookup(tf.constant([b"A", b"G", b"C", b"T"])).numpy().tolist() == [0, 1, 2, 3] - - def test_remap_labels(self): - table = encode._remap_labels([0, 1], [10, 20]) - assert table.lookup(tf.constant([0, 1, 2], dtype=tf.int32)).numpy().tolist() == [10, 20, 0] - - -class TestProcessStringTrain: - def test_translated_onehot(self): - processor = encode.process_string_train( - codon_depth=64, - num_classes=2, - crop_size=28, - fragsize=8, - ngram_width=3, - input_type="translated", - seq_onehot=True, - ) - csv = "0,ATGCATGCATGCATGCATGCATGCATGC,meta" - inputs, label = processor(tf.constant(csv)) - assert "translated" in inputs - assert inputs["translated"].shape.as_list() == [6, 8, 64] - assert label.shape.as_list() == [2] - - def test_nucleotide_onehot(self): - processor = encode.process_string_train( - codon_depth=64, - num_classes=2, - crop_size=12, - ngram_width=3, - input_type="nucleotide", - ) - csv = "1,ATGCATGCATGC,meta" - inputs, label = processor(tf.constant(csv)) - assert "nucleotide" in inputs - assert inputs["nucleotide"].shape.as_list() == [2, 12, 4] - - def test_label_remapping(self): - processor = encode.process_string_train( - codon_depth=64, - num_classes=2, - crop_size=12, - ngram_width=3, - label_original=[0, 1], - label_alternative=[1, 0], - ) - csv = "0,ATGCATGCATGC,meta" - _, label = processor(tf.constant(csv)) - assert np.argmax(label.numpy()) == 1 - - -class TestProcessStringInference: - def test_inference_outputs(self): - processor = encode.process_string_inference( - codon_depth=64, crop_size=24, fragsize=8, ngram_width=3, input_type="translated" - ) - csv = "ATGCATGCATGCATGCATGCATGC,h,0,1,0,24,5,5,5,5,0.0" - outputs = processor(tf.constant(csv)) - inputs, *meta = outputs - assert "translated" in inputs - assert len(meta) == 10 - - -class TestRawSequenceProcessors: - def test_make_process_raw_sequence_fn(self): - # crop_size must satisfy (crop_size-2) divisible by 3 so all frames match. - fn = encode._make_process_raw_sequence_fn( - crop_size=26, ngram_width=3, num_classes=2 - ) - seq = np.array([0, 1, 2, 3] * 7, dtype=np.int8)[:26] - label = np.array([1.0, 0.0], dtype=np.float32) - inputs, out_label = fn(tf.constant(seq), tf.constant(label)) - assert "translated" in inputs - assert inputs["translated"].shape.as_list()[0] == 6 - assert out_label.shape.as_list() == [2] - - def test_make_process_variable_sequence_fn(self): - fn = encode._make_process_variable_sequence_fn( - ngram_width=3, num_classes=2 - ) - seq = np.array([0, 1, 2, 3] * 6, dtype=np.int8) - label = np.array([1.0, 0.0], dtype=np.float32) - inputs, out_label = fn(tf.constant(seq), tf.constant(24), tf.constant(label)) - assert "translated" in inputs - assert out_label.shape.as_list() == [2] diff --git a/tests/unit/test_utils_misc.py b/tests/unit/test_utils_misc.py index 712aafb..d19469b 100644 --- a/tests/unit/test_utils_misc.py +++ b/tests/unit/test_utils_misc.py @@ -3,12 +3,9 @@ from __future__ import annotations import json -import tempfile from decimal import Decimal from pathlib import Path -import pytest - from jaeger.utils import misc From 41a52c0e13d94417deb09bc5e822f975fd34c04f Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 12:43:02 +0200 Subject: [PATCH 36/64] chore(release): final QA and v1.27.1 prep --- AGENTS.md | 230 ++++++++++++++++- CHANGELOG.md | 20 ++ README.md | 6 +- pyproject.toml | 2 +- src/jaeger/commands/predict.py | 39 +-- src/jaeger/commands/train.py | 20 +- src/jaeger/data/pytorch/transforms.py | 4 +- src/jaeger/inference/pytorch/runner.py | 34 ++- src/jaeger/nnlib/pytorch/builder.py | 12 +- src/jaeger/nnlib/pytorch/models.py | 35 ++- src/jaeger/training/pytorch/callbacks.py | 14 +- src/jaeger/training/pytorch/distributed.py | 8 +- .../test_pytorch_benchmark_smoke.py | 232 ++++++++++++++++++ tests/unit/data/pytorch/test_dataset_numpy.py | 7 +- tests/unit/nnlib/pytorch/test_layers.py | 14 +- tests/unit/nnlib/pytorch/test_losses.py | 4 +- tests/unit/nnlib/pytorch/test_models.py | 50 +++- tests/unit/test_dataops_dataset.py | 4 +- tests/unit/test_seqops_io.py | 15 +- tests/unit/training/pytorch/test_callbacks.py | 4 +- .../unit/training/pytorch/test_distributed.py | 6 +- 21 files changed, 680 insertions(+), 80 deletions(-) create mode 100644 tests/integration/test_pytorch_benchmark_smoke.py diff --git a/AGENTS.md b/AGENTS.md index b6bc234..f9cf951 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,7 +2,7 @@ ## Project background -Jaeger (`jaeger-bio`) is a homology-free deep-learning command-line tool for identifying bacteriophage genome sequences. The repository lives at `/home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/Jaeger` and uses Python, TensorFlow, NumPy, and YAML training configs. +Jaeger (`jaeger-bio`) is a homology-free deep-learning command-line tool for identifying bacteriophage genome sequences. The repository lives at `/home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/Jaeger` and uses Python, PyTorch, NumPy, and YAML training configs. ## Memory protocol — MemPalace @@ -28,3 +28,231 @@ This project uses [MemPalace](https://mempalaceofficial.com/) for durable, searc ## Tool permission note Read/search MemPalace tools (`mempalace_status`, `mempalace_search`, `mempalace_get_*`, `mempalace_kg_query`) are pre-approved in `~/.kimi-code/config.toml`. Write tools (`add_drawer`, `kg_add`, `diary_write`) will trigger approval prompts unless the user enables YOLO mode. + +--- + +# Agent working notes + +The sections below supplement the README and `docs/_source/` with conventions and pitfalls that are useful when editing code or answering questions about Jaeger. + +## Project basics + +| Item | Value | +|---|---| +| PyPI package | `jaeger-bio` | +| Import package | `jaeger` | +| Current version | `1.27.1` | +| Python support | `>=3.11, <3.14` | +| Build backend | `pdm-backend` | +| Source layout | `src/jaeger/` | +| CLI entry point | `jaeger = jaeger.cli:main` | +| License | MIT | + +## Development workflow + +Install in editable mode for development: + +```bash +pip install -e ".[gpu,test]" # or [cpu,test], [darwin-arm,test] +``` + +Run the test suite: + +```bash +pytest +``` + +`pyproject.toml` configures `testpaths = ["tests/unit", "tests/integration"]`. Smoke and CLI tests also live under `tests/smoke/`, `tests/pytest/`, and `test_cli/`. + +Build a release artifact: + +```bash +pdm build +``` + +Check the installation health: + +```bash +jaeger health +``` + +## Code organization + +| Path | Purpose | +|---|---| +| `src/jaeger/cli.py` | Click CLI definition, PyTorch/legacy model routing, and TF import suppression for legacy workflows | +| `src/jaeger/commands/` | Command implementations: `predict`, `predict_legacy`, `train`, `tune`, `health`, `quantize`, `convert_graph`, `taxonomy`, `downloads`, `test`, plus `utils*.py` | +| `src/jaeger/nnlib/` | Neural-network library: `v1/` and `v2/` layers, losses, metrics, inference, conversion, builder | +| `src/jaeger/dataops/` | Data operations: `convert.py`, `dataset.py`, `ood.py`, `split.py` | +| `src/jaeger/data/` | Bundled data (`config.json`), TFRecord helpers, dataset loaders | +| `src/jaeger/seqops/` | Sequence operations: encode, maps, stats, synthetic, transform, validate, io | +| `src/jaeger/postprocess/` | Post-processing: `collect.py`, `helpers.py`, `prophages.py` | +| `src/jaeger/preprocess/` | Preprocessing maps/converters (`v1/`, `v2/`) | +| `src/jaeger/utils/` | Utilities: logging, filesystem, misc, stats, termini, GPU, test helpers | +| `train_config/` | YAML training configuration templates | +| `scripts/` | Standalone conversion scripts for NumPy/TFRecord formats | +| `recipes/jaeger-bio/` | Bioconda recipe | +| `.github/scripts/` | `bump-version.sh` and dependency helpers | +| `.github/workflows/` | PyPI publish, GitHub Release, and Bioconda update workflows | + +## CLI conventions + +- Commands are built with **Click** and usually set `context_settings={"show_default": True}`. +- Verbosity is a **count flag**: default is warning, `-v` is info, `-vv` is debug. +- Modern SavedModels (e.g., `jaeger_38341_1.4M_fragment`) are recommended; the bundled `default` model and the `experimental_*` models use the **legacy prediction workflow** and are deprecated. +- Model discovery uses `src/jaeger/data/config.json` under the `model_paths` key. Register a new model directory with: + ```bash + jaeger register-models --path /path/to/my_model + ``` +- Download additional models with: + ```bash + jaeger download --list + jaeger download --model_name jaeger_38341_1.4M --path /path/to/store/models + ``` + +## PyTorch backend and environment notes + +Starting with v1.27.1, Jaeger uses **PyTorch** as its primary training and inference backend. The `jaeger train` and `jaeger predict` commands run through the PyTorch path by default. + +TensorFlow is still imported only for the deprecated legacy prediction workflow (`jaeger predict-legacy`) and bundled `default`/`experimental_*` models. When TensorFlow is imported, Jaeger aggressively suppresses TF and low-level C++ logging because the CLI is user-facing: + +- `TF_CPP_MIN_LOG_LEVEL=3` +- `TF_ENABLE_ONEDNN_OPTS=0` +- `GRPC_VERBOSITY=ERROR` +- `GLOG_minloglevel=2` +- `ABSL_MIN_LOG_LEVEL=2` + +These variables are set in `src/jaeger/__init__.py` and again at the top of `src/jaeger/cli.py` **before** the first `import tensorflow`. + +`src/jaeger/cli.py` also contains a compatibility patch for TensorFlow 2.18 + Python 3.12 that wraps `tensorflow.python.framework.tensor_util.is_tf_type`. Do not move `import tensorflow` before this patch. + +On Linux, `XLA_FLAGS` is set to `--xla_gpu_cuda_data_dir=/usr/lib/cuda` to support XLA JIT (`--xla`) for the legacy TensorFlow path. + +## Training data and model formats + +### Data formats + +Training data is loaded from CSV by default and preprocessed on-the-fly. For large training jobs, convert to an optimized format and set `model.string_processor.data_format`: + +| Format | Speedup | Best for | +|---|---|---| +| `csv` | 1.0× | Small datasets, quick experiments | +| `tfrecord` | ~12× | Datasets too large for RAM | +| `numpy_raw` | ~17× | Fast loading + runtime augmentations | +| `numpy_raw_variable` | ~3× | Variable-length sequences | +| `numpy_full` | ~9× | Maximum throughput, no augmentations | + +Convert CSV to any optimized format with: + +```bash +jaeger utils optimize-data -i train.csv -o train.npz --format numpy_full +``` + +Standalone conversion scripts are also available in `scripts/`. + +### Inference backends + +`jaeger predict` uses PyTorch by default. Several optional backends and precision modes are available for legacy TensorFlow SavedModels and converted graphs: + +- `--precision fp16|bf16` — mixed-precision inference (legacy TensorFlow path) +- `--xla` — XLA JIT compilation (legacy TensorFlow path) +- `--onnx` — ONNX Runtime inference (requires `jaeger utils convert-graph --mode onnx` first) +- `--onnx --int8` — INT8 quantized ONNX +- `--quantized dynamic|float16|full_int8` — TFLite quantization + +## Release workflow + +Jaeger uses a calendar-style version format `1..` (e.g., `1.26.4`). + +1. Bump the version: + ```bash + .github/scripts/bump-version.sh + ``` + This updates `pyproject.toml`, `.cz.toml`, `recipes/jaeger-bio/meta.yaml`, `singularity/jaeger_singularity.def`, `README.md`, `AGENTS.md`, and `docs/_source/usage.md`. +2. Commit and tag: + ```bash + git add -A + git commit -m "chore(release): bump version to X.Y.Z" + git tag -a vX.Y.Z -m "Release version X.Y.Z" + ``` + Or use `cz bump` to update `CHANGELOG.md` and create the tag. +3. Push to the canonical repository `Yasas1994/Jaeger`: + ```bash + git push origin main + git push origin --tags + ``` +4. CI workflows run in sequence: + - `publish-to-pypi.yaml` — builds and publishes to PyPI via Trusted Publishing (OIDC) + - `release.yaml` — creates a GitHub Release from the tag and `CHANGELOG.md` + - `bioconda-update.yaml` — opens a Bioconda PR after the GitHub Release is published + +Only `Yasas1994/Jaeger` can publish to PyPI and Bioconda; forks build artifacts for validation only. + +Manual fallback: + +```bash +pdm build +pdm publish +``` + +## Testing conventions + +- Use `pytest` for all tests. +- Unit and integration tests live in `tests/unit/` and `tests/integration/`. +- Smoke and CLI-specific tests are in `tests/smoke/`, `tests/pytest/`, and `test_cli/`. +- Shared fixtures are in `tests/conftest.py`. +- CLI tests use `click.testing.CliRunner` and import from `jaeger.cli`. +- Keep TensorFlow quiet in tests by setting `TF_CPP_MIN_LOG_LEVEL=2` and `GRPC_VERBOSITY=ERROR` in fixtures or test setup when needed. + +## Common agent pitfalls + +- Do not change the PyPI package name (`jaeger-bio`); the import package remains `jaeger`. +- Do not move `import tensorflow` above the environment-variable setup in `cli.py` or `__init__.py`. +- Do not delete the legacy `predict_legacy` path or the `default` model config; mark deprecated workflows as deprecated instead. +- When adding new CLI commands, follow the existing Click patterns and add corresponding help tests. +- When updating the version, use `.github/scripts/bump-version.sh` so all versioned files stay in sync. + +--- + +# Cluster Agent Instructions + +## Cluster : zeus + +**Project directory:** /mnt/beegfs/bioinf/wijesekara/jaeger + +## File Structure + +``` +project-root/ +├── data/ # Input datasets +├── configs/ # Jaeger training configs +├── experiments/ # Training outputs: checkpoints, metrics, logs +├── slurm/ +│ ├── logs/ # Job stdout/stderr +│ └── scripts/ # Batch scripts +├── tmp/ # Scratch files +├── container/ # Singularity container image(s) +├── Jaeger/ # Jaeger source code (bound into container at runtime) +``` + +## Partitions & Resources + +| Partition | Nodes | GPU | Max time | CPUs | Mem | +|-----------|-----------|---------|-----------|------|------| +| `gpu` | 030 | 2× L40S | 72:00:00 | 128 | 500G | +| `batch` | 010 – 016 | — | 72:00:00 | 40 | 220G | +| `batch` | 001 – 008 | — | 72:00:00 | 8 | 220G | + +## Agent Behaviour + +- Do not modify `AGENTS.md`. +- When uncertain about something related to running a jobs on zeus, search the mempalace first. +- Update mempalace when you make progress, a new decision, or a plan. + +## Rules + +- Always run jobs via the Singularity container in `container/`, binding the `Jaeger/` source directory into it. +- Always test on `gpu` with a small number of steps or a small dataset before full runs. +- Save checkpoints regularly; resume from the latest checkpoint on restart. +- Do not use or configure conda environments to run jobs. +- If you encounter a bug in Jaeger, fix it locally, test the fix, then sync the updated source code to zeus. diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d317d4..0650d8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,26 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## v1.27.1 (2026-06-15) + +### Changed + +- Migrated the training and inference backend from TensorFlow to PyTorch. + - New `jaeger.nnlib.pytorch` modules provide model building, layers, losses, metrics, and conversion utilities. + - New `jaeger.training.pytorch` trainer, callbacks, and distributed helpers support multi-GPU and mixed-precision training. + - New `jaeger.data.pytorch` CSV/NumPy dataset builders and collators feed variable-length sequences to PyTorch models. + - `jaeger train` now runs the PyTorch training path by default. + - `jaeger predict` now uses the PyTorch inference runner; legacy TensorFlow SavedModel support is retained only for deprecated models. + +### Added + +- SLURM batch scripts and Singularity definitions updated for the PyTorch container workflow. +- CPU-only (`[cpu]`), GPU (`[gpu]`), and Apple Silicon (`[darwin-arm]`) optional dependency groups now resolve PyTorch wheels instead of TensorFlow. + +### Deprecated + +- The bundled `default` TensorFlow SavedModel and `experimental_*` models remain available via `jaeger predict-legacy` but are no longer recommended. + ## v1.26.4 (2026-06-13) ### Feat diff --git a/README.md b/README.md index bcf1d3b..7bfab24 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,9 @@ Jaeger: an accurate and fast deep-learning tool to detect bacteriophage sequence -Jaeger is a tool that utilizes homology-free machine learning to identify phage genome sequences that are hidden within metagenomes. It is capable of detecting both phages and prophages within metagenomic assemblies. +Jaeger is a tool that utilizes homology-free deep learning to identify phage genome sequences that are hidden within metagenomes. It is capable of detecting both phages and prophages within metagenomic assemblies. + +> ⚡ **Jaeger v1.27+ uses PyTorch as its default training and inference backend.** Legacy TensorFlow SavedModels are still supported via `jaeger predict-legacy` but are deprecated. > 📚 **For detailed installation instructions, usage guides, and troubleshooting, please visit the [documentation](https://jaeger.readthedocs.io/en/latest/installation.html).** @@ -82,6 +84,8 @@ jaeger health ##### option 3 : Installing from PyPI (recommended) +The PyPI package installs the PyTorch backend by default. + ```bash # create a conda environment and activate mamba create -n jaeger -c nvidia -c conda-forge cuda-nvcc "python>=3.11,<3.14" pip diff --git a/pyproject.toml b/pyproject.toml index 2721963..e150399 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ includes = ["src/jaeger/**/*.py", "src/jaeger/**/*.json", "src/jaeger/**/*.h5", [project] name = "jaeger-bio" -version = "1.26.4" +version = "1.27.1" description = "A quick and precise pipeline for detecting phages in sequence assemblies." readme = "README.md" authors = [ diff --git a/src/jaeger/commands/predict.py b/src/jaeger/commands/predict.py index 3110b63..440c8ff 100644 --- a/src/jaeger/commands/predict.py +++ b/src/jaeger/commands/predict.py @@ -25,7 +25,7 @@ def _resolve_codon_value(value): """Resolve a codon table configuration value to a list.""" if isinstance(value, str): - if value == "CODONS": + if value in ("CODON", "CODONS"): return CODONS if value == "CODON_ID": return CODON_ID @@ -56,9 +56,9 @@ def _codon_ids(s: str): forward = _codon_ids(seq) reverse = _codon_ids(rev) - frames = [ - forward[offset::3] for offset in range(3) - ] + [reverse[offset::3] for offset in range(3)] + frames = [forward[offset::3] for offset in range(3)] + [ + reverse[offset::3] for offset in range(3) + ] target_len = max(0, len(seq) // 3 - 1) frames = [f[:target_len] + [-1] * (target_len - len(f)) for f in frames] @@ -137,8 +137,7 @@ def run_core(**kwargs): log_file = Path(f"{file_base}_jaeger.log") logger = get_logger(output_dir, log_file, level=kwargs.get("verbose", 1)) logger.info( - description(version("jaeger-bio")) - + "\n{:-^80}".format("validating parameters") + description(version("jaeger-bio")) + "\n{:-^80}".format("validating parameters") ) logger.debug(config) @@ -170,9 +169,7 @@ def run_core(**kwargs): ) precision = kwargs.get("precision", "fp32") if precision and precision != "fp32": - logger.warning( - f"--precision {precision} is not supported and will be ignored." - ) + logger.warning(f"--precision {precision} is not supported and will be ignored.") if kwargs.get("cpu") or not torch.cuda.is_available(): device = torch.device("cpu") @@ -197,13 +194,9 @@ def run_core(**kwargs): logger.info(f"model: {model_id}") if current_process is not None: - logger.info( - f"avail mem: {psutil.virtual_memory().available / GB_BYTES:.2f}GB" - ) + logger.info(f"avail mem: {psutil.virtual_memory().available / GB_BYTES:.2f}GB") logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info( - f"wall time(s) : {time.time() - current_process.create_time():.2f}" - ) + logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") logger.info( f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB " f"({current_process.memory_percent():.2f}%)" @@ -348,9 +341,7 @@ def run_core(**kwargs): identifier="phage", ) plot_type = kwargs.get("plot_type", "circular") - config_labels = { - i: c for i, c in enumerate(class_map.get("class", [])) - } + config_labels = {i: c for i, c in enumerate(class_map.get("class", []))} if plot_type in ("circular", "both"): plot_scores( logits_df, @@ -380,9 +371,7 @@ def run_core(**kwargs): else: logger.info("no prophage regions found") except Exception as e: - logger.error( - f"an error {e} occurred during the prophage prediction step" - ) + logger.error(f"an error {e} occurred during the prophage prediction step") logger.debug(traceback.format_exc()) if kwargs.get("getsequences"): @@ -401,9 +390,7 @@ def run_core(**kwargs): if kwargs.get("window_scores"): output_scores_path = output_dir / f"{file_base}_window_scores.npz" - logger.info( - f"writing window-wise scores and metadata to {output_scores_path}" - ) + logger.info(f"writing window-wise scores and metadata to {output_scores_path}") np.savez( output_scores_path, headers=data_full["headers"], @@ -416,9 +403,7 @@ def run_core(**kwargs): if current_process is not None: logger.info(f"CPU time(s) : {current_process.cpu_times().user:.2f}") - logger.info( - f"wall time(s) : {time.time() - current_process.create_time():.2f}" - ) + logger.info(f"wall time(s) : {time.time() - current_process.create_time():.2f}") logger.info( f"memory usage : {current_process.memory_full_info().rss / GB_BYTES:.2f}GB " f"({current_process.memory_percent():.2f}%)" diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 6f20935..0f892de 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -59,6 +59,22 @@ def _count_parameters(model: nn.Module) -> Tuple[int, int]: return total, trainable +def _initialize_lazy_layers( + models: Dict[str, nn.Module], config: Dict[str, Any] +) -> None: + """Run a dummy forward pass to materialize any lazy layers.""" + model_cfg = config.get("model", {}) + embedding_cfg = model_cfg.get("embedding", {}) + sp_cfg = model_cfg.get("string_processor", {}) + frames = int(embedding_cfg.get("frames", 6)) + crop_size = int(sp_cfg.get("crop_size", 500)) + device = next(models["jaeger_classifier"].parameters()).device + dummy = torch.zeros((1, frames, crop_size), dtype=torch.long, device=device) + mask = torch.ones((1, frames, crop_size), dtype=torch.bool, device=device) + with torch.no_grad(): + models["jaeger_classifier"](dummy, mask) + + def _load_checkpoint_if_requested( model: nn.Module, optimizer: Optional[torch.optim.Optimizer], @@ -282,6 +298,8 @@ def train_fragment_core(**kwargs): models = builder.build_fragment_classifier() rep_model = models["rep_model"] + _initialize_lazy_layers(models, config) + total_params, trainable_params = _count_parameters(rep_model) logger.info( "rep_model parameters: total=%s, trainable=%s", @@ -411,7 +429,7 @@ def train_fragment_core(**kwargs): if kwargs.get("save_model", False) or kwargs.get("only_save", False): save_dir = model_save_path if model_save_path else classifier_dir _save_model_checkpoint( - models["jaeger_classifier"], + models["jaeger_model"], save_dir, "classifier.pt", metadata=kwargs.get("meta"), diff --git a/src/jaeger/data/pytorch/transforms.py b/src/jaeger/data/pytorch/transforms.py index 2f3af82..5acba89 100644 --- a/src/jaeger/data/pytorch/transforms.py +++ b/src/jaeger/data/pytorch/transforms.py @@ -36,9 +36,7 @@ def translate_to_codons(seq: np.ndarray, codon_table: Dict[str, int]) -> torch.T frames.append(torch.tensor(rev_indices, dtype=torch.long)) max_len = max(len(f) for f in frames) - frames = [ - torch.nn.functional.pad(f, (0, max_len - len(f))) for f in frames - ] + frames = [torch.nn.functional.pad(f, (0, max_len - len(f))) for f in frames] return torch.stack(frames) diff --git a/src/jaeger/inference/pytorch/runner.py b/src/jaeger/inference/pytorch/runner.py index c50a302..1c7ab1d 100644 --- a/src/jaeger/inference/pytorch/runner.py +++ b/src/jaeger/inference/pytorch/runner.py @@ -6,6 +6,20 @@ from jaeger.nnlib.pytorch.builder import ModelBuilder +def _dummy_input_from_config( + config: Dict[str, Any], +) -> tuple[torch.Tensor, torch.Tensor]: + """Return a dummy (x, mask) tuple that matches the model input shape.""" + model_cfg = config.get("model", {}) + embedding_cfg = model_cfg.get("embedding", {}) + sp_cfg = model_cfg.get("string_processor", {}) + frames = int(embedding_cfg.get("frames", 6)) + crop_size = int(sp_cfg.get("crop_size", 500)) + x = torch.zeros((1, frames, crop_size), dtype=torch.long) + mask = torch.ones((1, frames, crop_size), dtype=torch.bool) + return x, mask + + class PyTorchInferenceRunner: """Load and run a PyTorch Jaeger model for inference.""" @@ -33,10 +47,22 @@ def load_checkpoint(self, path: Union[str, Path]): state = torch.load(path, map_location=self.device, weights_only=True) except TypeError: state = torch.load(path, map_location=self.device, weights_only=False) - if "model_state_dict" in state: - self.model.load_state_dict(state["model_state_dict"]) - else: - self.model.load_state_dict(state) + sd = state.get("model_state_dict", state) + try: + self.model.load_state_dict(sd) + except RuntimeError: + # The checkpoint may contain trained MaskedConv1D weights that + # require lazy layers to be materialized before loading. + self._initialize_lazy_layers() + self.model.load_state_dict(sd) + + def _initialize_lazy_layers(self): + """Run a dummy forward so MaskedConv1D layers materialize before loading.""" + dummy_x, dummy_mask = _dummy_input_from_config(self.config) + dummy_x = dummy_x.to(self.device) + dummy_mask = dummy_mask.to(self.device) + with torch.no_grad(): + self.model(dummy_x, dummy_mask) def predict( self, diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py index dfff3a7..69e7330 100644 --- a/src/jaeger/nnlib/pytorch/builder.py +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -20,7 +20,9 @@ def __init__(self, rep_model: nn.Module, head: nn.Module): self.rep_model = rep_model self.head = head - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: embedding = self.rep_model(x, mask) if isinstance(embedding, tuple): embedding = embedding[0] @@ -39,7 +41,9 @@ def __init__(self, config: Dict[str, Any]): def build_fragment_classifier(self) -> Dict[str, nn.Module]: if "classifier" not in self.model_cfg: - raise ValueError("classifier config is required for build_fragment_classifier") + raise ValueError( + "classifier config is required for build_fragment_classifier" + ) models: Dict[str, nn.Module] = {} @@ -109,9 +113,7 @@ def compile_model( "reliability training branch is not implemented yet" ) if train_branch == "pretrain": - raise NotImplementedError( - "pretrain training branch is not implemented yet" - ) + raise NotImplementedError("pretrain training branch is not implemented yet") raise ValueError(f"Unknown train_branch: {train_branch}") def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index 08de46a..f5224cb 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -71,7 +71,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ClassificationHead(nn.Module): """Dense head for class prediction.""" - def __init__(self, input_dim: int, num_classes: int, hidden_units: List[int], dropout: float = 0.0): + def __init__( + self, + input_dim: int, + num_classes: int, + hidden_units: List[int], + dropout: float = 0.0, + ): super().__init__() layers = [] prev = input_dim @@ -92,7 +98,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ReliabilityHead(nn.Module): """Dense head for confidence estimation.""" - def __init__(self, input_dim: int, num_classes: int, hidden_units: List[int], dropout: float = 0.0): + def __init__( + self, + input_dim: int, + num_classes: int, + hidden_units: List[int], + dropout: float = 0.0, + ): super().__init__() layers = [] prev = input_dim @@ -126,7 +138,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class RepresentationModel(nn.Module): - def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str = "average"): + def __init__( + self, embedding: nn.Module, hidden_layers: List[dict], pooling: str = "average" + ): super().__init__() self.embedding = embedding self.blocks = nn.ModuleList() @@ -147,7 +161,7 @@ def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str elif name == "transformer_encoder": self.blocks.append(TransformerEncoder(**config)) elif name == "residual_block": - inner = self.blocks.pop() if self.blocks else None + inner = self.blocks.pop(-1) if self.blocks else None self.blocks.append(ResidualBlock(inner)) elif name == "dense": if "units" in config: @@ -155,7 +169,9 @@ def __init__(self, embedding: nn.Module, hidden_layers: List[dict], pooling: str config["out_features"] = config.pop("units") self.blocks.append(nn.Linear(**config)) elif name == "activation": - self.blocks.append(GeLU() if config.get("activation") == "gelu" else nn.ReLU()) + self.blocks.append( + GeLU() if config.get("activation") == "gelu" else nn.ReLU() + ) elif name == "dropout": self.blocks.append(nn.Dropout(config.get("rate", 0.0))) else: @@ -256,7 +272,9 @@ def __init__( self.classification_head = classification_head self.reliability_head = reliability_head - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Dict[str, torch.Tensor]: outputs = self.rep_model(x, mask) if isinstance(outputs, tuple): embedding = outputs[0] @@ -268,7 +286,10 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[ gate = None prediction = self.classification_head(embedding) - result: Dict[str, torch.Tensor] = {"prediction": prediction, "embedding": embedding} + result: Dict[str, torch.Tensor] = { + "prediction": prediction, + "embedding": embedding, + } if nmd is not None: result["nmd"] = nmd if gate is not None: diff --git a/src/jaeger/training/pytorch/callbacks.py b/src/jaeger/training/pytorch/callbacks.py index 3b2b160..bbf2bb8 100644 --- a/src/jaeger/training/pytorch/callbacks.py +++ b/src/jaeger/training/pytorch/callbacks.py @@ -34,16 +34,18 @@ def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): if current is None: return - improved = (self.mode == "min" and current < self.best_value - self.min_delta) or ( - self.mode == "max" and current > self.best_value + self.min_delta - ) + improved = ( + self.mode == "min" and current < self.best_value - self.min_delta + ) or (self.mode == "max" and current > self.best_value + self.min_delta) if improved: self.best_value = current self.best_epoch = epoch self.wait = 0 if self.restore_best_weights: - self.best_state = {k: v.cpu().clone() for k, v in trainer.model.state_dict().items()} + self.best_state = { + k: v.cpu().clone() for k, v in trainer.model.state_dict().items() + } else: self.wait += 1 if self.wait >= self.patience: @@ -90,7 +92,9 @@ def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): if self.save_best_only: self._save(trainer, epoch) if self.verbose: - print(f"Epoch {epoch}: {self.monitor} improved to {current:.4f}; saving model") + print( + f"Epoch {epoch}: {self.monitor} improved to {current:.4f}; saving model" + ) elif not self.save_best_only: self._save(trainer, epoch) diff --git a/src/jaeger/training/pytorch/distributed.py b/src/jaeger/training/pytorch/distributed.py index ce31914..1fdc3a5 100644 --- a/src/jaeger/training/pytorch/distributed.py +++ b/src/jaeger/training/pytorch/distributed.py @@ -16,7 +16,9 @@ def setup_distributed(backend: str = "nccl") -> bool: rank = int(os.environ["SLURM_PROCID"]) world_size = int(os.environ["SLURM_NTASKS"]) local_rank = int(os.environ.get("SLURM_LOCALID", 0)) - dist.init_process_group(backend=backend, init_method="env://", rank=rank, world_size=world_size) + dist.init_process_group( + backend=backend, init_method="env://", rank=rank, world_size=world_size + ) if torch.cuda.is_available(): torch.cuda.set_device(local_rank) return True @@ -32,7 +34,9 @@ def cleanup_distributed(): def get_device() -> torch.device: if dist.is_initialized(): local_rank = int(os.environ.get("LOCAL_RANK", 0)) - return torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") + return torch.device( + f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu" + ) return torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/integration/test_pytorch_benchmark_smoke.py b/tests/integration/test_pytorch_benchmark_smoke.py new file mode 100644 index 0000000..df55269 --- /dev/null +++ b/tests/integration/test_pytorch_benchmark_smoke.py @@ -0,0 +1,232 @@ +"""Smoke test for the PyTorch training and inference path. + +Builds a tiny model and tiny CSV dataset, runs one training epoch, then runs +prediction against a small FASTA input. The test only asserts that the outputs +are produced without error; it does not assert performance. +""" + +from __future__ import annotations + +from pathlib import Path + +from click.testing import CliRunner + +from jaeger.cli import main as jaeger_cli +from jaeger.commands.train import train_fragment_core +from jaeger.utils.misc import load_model_config + + +def _write_tiny_csv(path: Path, n_samples: int = 16, seq_len: int = 96) -> None: + """Write a tiny label-first CSV dataset.""" + rng = range(n_samples) + rows = [] + for i in rng: + label = i % 3 + seq = "ATGCATGC" * (seq_len // 8) + rows.append(f"{label},{seq},seq{i}") + path.write_text("\n".join(rows) + "\n") + + +def _write_tiny_fasta(path: Path, n_contigs: int = 2, seq_len: int = 2000) -> None: + """Write a tiny FASTA file for inference.""" + seq = "ATGCATGC" * (seq_len // 8) + lines = [f">contig{i}\n{seq}" for i in range(n_contigs)] + path.write_text("\n".join(lines) + "\n") + + +def _tiny_config(tmp_path: Path) -> dict[str, object]: + """Return a minimal Jaeger PyTorch training config dict.""" + data_dir = tmp_path / "data" + data_dir.mkdir(parents=True, exist_ok=True) + train_csv = data_dir / "train.csv" + val_csv = data_dir / "val.csv" + _write_tiny_csv(train_csv) + _write_tiny_csv(val_csv) + + experiment_root = "experiments/tiny_smoke" + classifier_dir = data_dir / experiment_root / "checkpoints" / "classifier" + model_save_path = data_dir / experiment_root / "model" + + return { + "model": { + "name": "jaeger_tiny_smoke", + "experiment": "tiny_smoke", + "seed": 42, + "classifier_out_dim": 3, + "reliability_out_dim": 0, + "class_label_map": [ + {"class": "chromosome", "label": 0}, + {"class": "virus", "label": 1}, + {"class": "plasmid", "label": 2}, + ], + "activation": "gelu", + "mode": "training", + "embedding": { + "use_embedding_layer": True, + "input_type": "translated", + "vocab_size": 65, + "strands": 2, + "frames": 6, + "length": None, + "input_shape": [6, None], + "embedding_size": 8, + "embedding_regularizer": "l2", + "embedding_regularizer_w": 1.0e-05, + }, + "string_processor": { + "data_format": "csv", + "seq_onehot": False, + "codon": "CODON", + "codon_id": "CODON_ID", + "crop_size": 96, + "buffer_size": 1000, + "shuffle": False, + "reshuffle_each_iteration": False, + "mutate": False, + "mutation_rate": 0.1, + "shuffle_frames": False, + "masking": False, + "classifier_labels": [0, 1, 2], + "classifier_labels_map": [0, 1, 2], + }, + "representation_learner": { + "hidden_layers": [ + { + "name": "masked_conv1d", + "config": { + "filters": 8, + "kernel_size": 3, + "padding": "same", + "activation": None, + "use_bias": True, + }, + }, + { + "name": "masked_batchnorm", + "config": {"num_features": 8, "return_nmd": False}, + }, + {"name": "activation", "config": {"activation": "gelu"}}, + ], + "pooling": "average", + }, + "classifier": { + "input_shape": 8, + "dropout": 0.1, + "hidden_layers": [ + {"name": "dense", "config": {"units": 8, "activation": None}}, + {"name": "dense", "config": {"units": 3, "activation": None}}, + ], + }, + }, + "training": { + "data_dir": str(data_dir), + "experiment_root": experiment_root, + "classifier_dir": str(classifier_dir), + "reliability_dir": str( + data_dir / experiment_root / "checkpoints" / "reliability" + ), + "classifier_epochs": 1, + "reliability_epochs": 0, + "projection_epochs": 0, + "classifier_train_steps": 1000, + "reliability_train_steps": 0, + "classifier_validation_steps": 100, + "reliability_validation_steps": 0, + "batch_size": 4, + "optimizer": "adam", + "optimizer_params": {"learning_rate": 0.001}, + "loss_classifier": "categorical_crossentropy", + "loss_params_classifier": {"from_logits": True}, + "metrics_classifier": [{"name": "categorical_accuracy", "params": None}], + "callbacks": {"classifier": []}, + "model_saving": { + "path": str(model_save_path), + "save_weights": True, + "save_exec_graph": False, + }, + "fragment_classifier_data": { + "train": [ + { + "class": ["chromosome", "virus", "plasmid"], + "path": [str(train_csv)], + "label": [0, 1, 2], + } + ], + "validation": [ + { + "class": ["chromosome", "virus", "plasmid"], + "path": [str(val_csv)], + "label": [0, 1, 2], + } + ], + }, + }, + } + + +def test_pytorch_train_predict_smoke(tmp_path: Path) -> None: + """Train a tiny PyTorch Jaeger model and run a tiny prediction job.""" + config_path = tmp_path / "config.yaml" + config = _tiny_config(tmp_path) + + import yaml + + config_path.write_text(yaml.safe_dump(config)) + + # Train for one epoch and save the model. + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=True, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + experiment_root = ( + Path(config["training"]["data_dir"]) / config["training"]["experiment_root"] + ) + model_dir = experiment_root / "model" + assert model_dir.exists(), f"model directory was not created: {model_dir}" + checkpoint = model_dir / "classifier.pt" + assert checkpoint.exists(), f"classifier checkpoint was not created: {checkpoint}" + + # Prepare a tiny FASTA input and run prediction. + fasta_path = tmp_path / "input.fasta" + _write_tiny_fasta(fasta_path) + + output_dir = tmp_path / "predictions" + runner = CliRunner() + result = runner.invoke( + jaeger_cli, + [ + "predict", + "--config", + str(config_path), + "--checkpoint", + str(checkpoint), + "--input", + str(fasta_path), + "--output", + str(output_dir), + "--cpu", + "--overwrite", + "--batch", + "2", + ], + ) + assert result.exit_code == 0, f"predict failed: {result.output}\n{result.exception}" + + model_name = ( + load_model_config(config_path).get("model", {}).get("name", "jaeger_pytorch") + ) + prediction_output = output_dir / model_name.replace(" ", "_") / "input.tsv" + assert prediction_output.exists(), ( + f"prediction TSV was not created: {prediction_output}" + ) diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/data/pytorch/test_dataset_numpy.py index 5e6550f..c7f9b37 100644 --- a/tests/unit/data/pytorch/test_dataset_numpy.py +++ b/tests/unit/data/pytorch/test_dataset_numpy.py @@ -100,7 +100,12 @@ def test_translate_to_codons_known_sequence(): assert x.shape[0] == 6 assert x.shape[1] == 4 - expected = [codon_table["ATG"], codon_table["AAA"], codon_table["TTT"], codon_table["CCC"]] + expected = [ + codon_table["ATG"], + codon_table["AAA"], + codon_table["TTT"], + codon_table["CCC"], + ] assert torch.equal(x[0], torch.tensor(expected, dtype=torch.long)) diff --git a/tests/unit/nnlib/pytorch/test_layers.py b/tests/unit/nnlib/pytorch/test_layers.py index cef8287..1fcd894 100644 --- a/tests/unit/nnlib/pytorch/test_layers.py +++ b/tests/unit/nnlib/pytorch/test_layers.py @@ -68,9 +68,7 @@ def test_masked_conv1d_no_mask(): def test_masked_conv1d_activation(): x = torch.abs(torch.randn(2, 3, 10, 4)) + 1.0 - layer = MaskedConv1D( - filters=8, kernel_size=3, padding="same", activation="relu" - ) + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same", activation="relu") out, _ = layer(x, mask=None) assert (out >= 0).all() @@ -261,8 +259,14 @@ def test_representation_model_forward(): model = RepresentationModel( embedding=embedding, hidden_layers=[ - {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}}, - {"name": "masked_batchnorm", "config": {"num_features": 16, "return_nmd": True}}, + { + "name": "masked_conv1d", + "config": {"filters": 16, "kernel_size": 3, "padding": "same"}, + }, + { + "name": "masked_batchnorm", + "config": {"num_features": 16, "return_nmd": True}, + }, ], pooling="average", ) diff --git a/tests/unit/nnlib/pytorch/test_losses.py b/tests/unit/nnlib/pytorch/test_losses.py index f0fd81f..6e1df4f 100644 --- a/tests/unit/nnlib/pytorch/test_losses.py +++ b/tests/unit/nnlib/pytorch/test_losses.py @@ -12,7 +12,9 @@ def test_arcface_loss_shape(): def test_arcface_loss_class_index_labels(): - loss = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.5, scale=30.0, onehot=False) + loss = ArcFaceLoss( + num_classes=3, embedding_dim=64, margin=0.5, scale=30.0, onehot=False + ) labels = torch.tensor([0, 1, 2], dtype=torch.long) embeddings = torch.randn(3, 64) out = loss(labels, embeddings) diff --git a/tests/unit/nnlib/pytorch/test_models.py b/tests/unit/nnlib/pytorch/test_models.py index d917ec5..8c0f8c6 100644 --- a/tests/unit/nnlib/pytorch/test_models.py +++ b/tests/unit/nnlib/pytorch/test_models.py @@ -1,16 +1,33 @@ import torch -from jaeger.nnlib.pytorch.models import ClassificationHead, Embedding, JaegerModel, ProjectionHead, ReliabilityHead, RepresentationModel +from jaeger.nnlib.pytorch.models import ( + ClassificationHead, + Embedding, + JaegerModel, + ProjectionHead, + ReliabilityHead, + RepresentationModel, +) def test_embedding_translated_with_embedding_layer(): - emb = Embedding(input_type="translated", vocab_size=65, embedding_size=32, use_embedding_layer=True) + emb = Embedding( + input_type="translated", + vocab_size=65, + embedding_size=32, + use_embedding_layer=True, + ) x = torch.randint(0, 65, (2, 6, 50)) out = emb(x) assert out.shape == (2, 6, 50, 32) def test_embedding_translated_without_embedding_layer(): - emb = Embedding(input_type="translated", vocab_size=65, embedding_size=32, use_embedding_layer=False) + emb = Embedding( + input_type="translated", + vocab_size=65, + embedding_size=32, + use_embedding_layer=False, + ) x = torch.randn(2, 6, 50, 65) out = emb(x) assert out.shape == (2, 6, 50, 32) @@ -62,18 +79,35 @@ def test_reliability_head_output_shape(): def test_jaeger_model_forward(): - embedding = Embedding(input_type="translated", vocab_size=65, embedding_size=32, use_embedding_layer=True) + embedding = Embedding( + input_type="translated", + vocab_size=65, + embedding_size=32, + use_embedding_layer=True, + ) rep_model = RepresentationModel( embedding=embedding, hidden_layers=[ - {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}}, - {"name": "masked_batchnorm", "config": {"num_features": 16, "return_nmd": True}}, + { + "name": "masked_conv1d", + "config": {"filters": 16, "kernel_size": 3, "padding": "same"}, + }, + { + "name": "masked_batchnorm", + "config": {"num_features": 16, "return_nmd": True}, + }, ], pooling="average", ) - classification_head = ClassificationHead(input_dim=16, num_classes=3, hidden_units=[8]) + classification_head = ClassificationHead( + input_dim=16, num_classes=3, hidden_units=[8] + ) reliability_head = ReliabilityHead(input_dim=16, num_classes=1, hidden_units=[8]) - model = JaegerModel(rep_model=rep_model, classification_head=classification_head, reliability_head=reliability_head) + model = JaegerModel( + rep_model=rep_model, + classification_head=classification_head, + reliability_head=reliability_head, + ) x = torch.randint(0, 65, (2, 6, 50)) mask = torch.ones(2, 6, 50, dtype=torch.bool) diff --git a/tests/unit/test_dataops_dataset.py b/tests/unit/test_dataops_dataset.py index bd314d9..bcc776a 100644 --- a/tests/unit/test_dataops_dataset.py +++ b/tests/unit/test_dataops_dataset.py @@ -20,9 +20,7 @@ def test_read_fasta(self, unlabeled_fasta_path: str): def test_read_csv(self, tmp_path: Path): csv = tmp_path / "in.csv" csv.write_text("0,ATGCATGC,seq1\n1,GCTAGCTA,seq2\n") - records = dataset.read_sequences( - csv, intype="CSV", seq_col=1, class_col=0 - ) + records = dataset.read_sequences(csv, intype="CSV", seq_col=1, class_col=0) assert len(records) == 2 assert records[0][1] == "ATGCATGC" diff --git a/tests/unit/test_seqops_io.py b/tests/unit/test_seqops_io.py index 6f509a5..3f7a335 100644 --- a/tests/unit/test_seqops_io.py +++ b/tests/unit/test_seqops_io.py @@ -58,7 +58,12 @@ class TestFragmentGenerators: def test_fragment_generator(self, unlabeled_fasta_path: str): frags = list( seqio.fragment_generator( - unlabeled_fasta_path, fragsize=16, stride=16, num=2, no_progress=True, dustmask=False + unlabeled_fasta_path, + fragsize=16, + stride=16, + num=2, + no_progress=True, + dustmask=False, ) ) assert len(frags) > 0 @@ -68,7 +73,9 @@ def test_fragment_generator(self, unlabeled_fasta_path: str): def test_fragment_generator_lib(self, unlabeled_fasta_path: str): frags = list( - seqio.fragment_generator_lib(unlabeled_fasta_path, fragsize=16, stride=16, num=2) + seqio.fragment_generator_lib( + unlabeled_fasta_path, fragsize=16, stride=16, num=2 + ) ) assert len(frags) > 0 parts = frags[0].split(",") @@ -101,7 +108,9 @@ def test_write_fasta_from_results(self, tmp_path: Path): ) data = {"contig_1": (df, 0, 12)} out = tmp_path / "phages.fasta" - seqio.write_fasta_from_results(data, str(out), reliability_cutoff=0.5, phage_score=1) + seqio.write_fasta_from_results( + data, str(out), reliability_cutoff=0.5, phage_score=1 + ) text = out.read_text() assert ">contig_1_0" in text assert ">contig_1_2" in text diff --git a/tests/unit/training/pytorch/test_callbacks.py b/tests/unit/training/pytorch/test_callbacks.py index e2a837b..321da63 100644 --- a/tests/unit/training/pytorch/test_callbacks.py +++ b/tests/unit/training/pytorch/test_callbacks.py @@ -43,7 +43,9 @@ def test_early_stopping(): val_loader = _make_data() loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.0) - early_stop = EarlyStopping(monitor="val_loss", patience=1, restore_best_weights=True) + early_stop = EarlyStopping( + monitor="val_loss", patience=1, restore_best_weights=True + ) trainer = Trainer( model=model, diff --git a/tests/unit/training/pytorch/test_distributed.py b/tests/unit/training/pytorch/test_distributed.py index e7ca931..fab3bdc 100644 --- a/tests/unit/training/pytorch/test_distributed.py +++ b/tests/unit/training/pytorch/test_distributed.py @@ -1,6 +1,10 @@ import os -from jaeger.training.pytorch.distributed import get_device, is_main_process, setup_distributed +from jaeger.training.pytorch.distributed import ( + get_device, + is_main_process, + setup_distributed, +) def test_get_device_without_distributed(): From cf1f52b5bed78d9c27915cc1107f0e8c9bcad3f5 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 13:41:26 +0200 Subject: [PATCH 37/64] docs(mempalace): persist PyTorch migration plan and project memory config --- .../plans/2026-06-14-pytorch-migration.md | 2703 +++++++++++++++++ .../2026-06-14-pytorch-migration-design.md | 299 ++ entities.json | 13 + mempalace.yaml | 60 + 4 files changed, 3075 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-14-pytorch-migration.md create mode 100644 docs/superpowers/specs/2026-06-14-pytorch-migration-design.md create mode 100644 entities.json create mode 100644 mempalace.yaml diff --git a/docs/superpowers/plans/2026-06-14-pytorch-migration.md b/docs/superpowers/plans/2026-06-14-pytorch-migration.md new file mode 100644 index 0000000..2c43515 --- /dev/null +++ b/docs/superpowers/plans/2026-06-14-pytorch-migration.md @@ -0,0 +1,2703 @@ +# PyTorch Migration Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace Jaeger's TensorFlow/Keras training and inference stack with a plain PyTorch implementation in the `pytorch_migration` branch, targeting release `1.27.1`. + +**Architecture:** Build a self-contained PyTorch module tree under `src/jaeger/nnlib/pytorch/`, `src/jaeger/data/pytorch/`, `src/jaeger/training/pytorch/`, and `src/jaeger/inference/pytorch/`. Preserve the YAML config format. Delete TensorFlow training/inference code only after parity tests and end-to-end smoke tests pass. + +**Tech Stack:** Python 3.11+, PyTorch 2.x, NumPy, Click, PyYAML, pytest. + +--- + +## File structure + +New files to create: + +``` +src/jaeger/nnlib/pytorch/ + __init__.py + layers.py # GeLU, MaskedConv1D, MaskedBatchNorm, MaskedLayerNorm, + # MaskedGlobalAvgPooling, GatedFrameGlobalMaxPooling, + # AxialAttention, CrossFrameAttention, TransformerEncoder, + # ResidualBlock, custom pooling helpers + models.py # Embedding, RepresentationModel, ClassificationHead, + # ReliabilityHead, ProjectionHead, JaegerModel + losses.py # ArcFaceLoss, HierarchicalLoss, classification losses + metrics.py # Per-class precision/recall/specificity + builder.py # ModelBuilder from YAML config + checkpoints.py # save/load .pt checkpoints + +src/jaeger/data/pytorch/ + __init__.py + transforms.py # codon translation, mutation, frame shuffle, masking + collate.py # padding collators + dataset_numpy.py # NumpyFullDataset, NumpyRawDataset, NumpyRawVariableDataset + dataset_csv.py # CSVDataset + dataset_tfrecord.py # TFRecordDataset (deferred) + +src/jaeger/training/pytorch/ + __init__.py + engine.py # Epoch loop, train/validation step, AMP + trainer.py # Trainer for classifier/reliability/pretrain branches + distributed.py # DDP setup/teardown + callbacks.py # ModelCheckpoint, EarlyStopping, LR scheduling, loggers + +src/jaeger/inference/pytorch/ + __init__.py + model.py # Load JaegerModel from checkpoint + config + engine.py # Batch inference, windowing, aggregation + +src/jaeger/commands/ + train.py # Dispatch to PyTorch trainer + predict.py # Dispatch to PyTorch inference engine +``` + +Test files to create: + +``` +tests/unit/nnlib/pytorch/ + test_layers.py + test_models.py + test_builder.py + test_losses.py + test_metrics.py +tests/unit/data/pytorch/ + test_dataset_numpy.py + test_transforms.py +tests/unit/training/pytorch/ + test_engine.py + test_callbacks.py +tests/integration/test_pytorch_parity.py +``` + +Files to delete at the end (after validation): + +``` +src/jaeger/nnlib/v1/ +src/jaeger/nnlib/v2/ +src/jaeger/nnlib/builder.py (old Keras builder) +src/jaeger/nnlib/inference.py (old TF inference) +src/jaeger/nnlib/conversion.py (if TF-only) +src/jaeger/commands/predict_legacy.py +src/jaeger/commands/train.py (old TF train command) +``` + +--- + +## Task 1: Create branch and update dependency metadata + +**Files:** +- Modify: `pyproject.toml` +- Modify: `.cz.toml` + +- [ ] **Step 1: Create the `pytorch_migration` branch** + +```bash +git checkout -b pytorch_migration +``` + +- [ ] **Step 2: Add PyTorch dependencies to `pyproject.toml`** + +Replace the TensorFlow optional dependencies with PyTorch equivalents. Keep TensorFlow as an optional `[legacy]` extra during development; remove it only in Task 27. + +```toml +[project.optional-dependencies] +gpu = [ + "torch >=2.0", + "torchvision", +] +cpu = [ + "torch >=2.0", + "torchvision", +] +darwin-arm = [ + "torch >=2.0", + "torchvision", +] +legacy = [ + "tensorflow >=2.21, <2.22", +] +test = [ + "pytest >=8.0", + "pytest-mock >=3.14", +] +``` + +- [ ] **Step 3: Verify `pyproject.toml` parses** + +```bash +python -c "import tomllib; tomllib.load(open('pyproject.toml','rb'))" +``` + +Expected: no exception. + +- [ ] **Step 4: Commit** + +```bash +git add pyproject.toml .cz.toml +git commit -m "chore(pytorch): add pytorch_migration branch and PyTorch deps" +``` + +--- + +## Task 2: Port activation and masking helper layers + +**Files:** +- Create: `src/jaeger/nnlib/pytorch/__init__.py` +- Create: `src/jaeger/nnlib/pytorch/layers.py` +- Create: `tests/unit/nnlib/pytorch/test_layers.py` + +- [ ] **Step 1: Write the failing test for GeLU** + +```python +# tests/unit/nnlib/pytorch/test_layers.py +import torch +import pytest +from jaeger.nnlib.pytorch.layers import GeLU + + +def test_gelu_matches_torch_nn_gelu(): + x = torch.tensor([-1.0, 0.0, 1.0, 2.0]) + layer = GeLU() + out = layer(x) + expected = torch.nn.functional.gelu(x, approximate="tanh") + assert torch.allclose(out, expected, atol=1e-6) +``` + +- [ ] **Step 2: Run the test to verify it fails** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_gelu_matches_torch_nn_gelu -v +``` + +Expected: `ModuleNotFoundError: No module named 'jaeger.nnlib.pytorch.layers'` + +- [ ] **Step 3: Implement GeLU and helpers** + +```python +# src/jaeger/nnlib/pytorch/__init__.py +# (empty or re-export later) +``` + +```python +# src/jaeger/nnlib/pytorch/layers.py +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class GeLU(nn.Module): + """Tanh-approximated GELU for TFLite-compatible graph export.""" + + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="tanh") +``` + +- [ ] **Step 4: Run the test to verify it passes** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_gelu_matches_torch_nn_gelu -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/__init__.py src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py +git commit -m "feat(pytorch): add GeLU activation layer" +``` + +--- + +## Task 3: Port MaskedConv1D + +**Files:** +- Modify: `src/jaeger/nnlib/pytorch/layers.py` +- Modify: `tests/unit/nnlib/pytorch/test_layers.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/nnlib/pytorch/test_layers.py +from jaeger.nnlib.pytorch.layers import MaskedConv1D + + +def test_masked_conv1d_shape_and_mask(): + # (B=2, F=6, L=10, C=4) + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[0, :, 5:] = False + layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") + out, out_mask = layer(x, mask) + assert out.shape == (2, 6, 10, 8) + assert out_mask.shape == (2, 6, 10) + assert not out_mask[0, :, 5:].any() +``` + +- [ ] **Step 2: Run the test to verify it fails** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_conv1d_shape_and_mask -v +``` + +Expected: `AttributeError: module 'jaeger.nnlib.pytorch.layers' has no attribute 'MaskedConv1D'` + +- [ ] **Step 3: Implement MaskedConv1D** + +```python +# src/jaeger/nnlib/pytorch/layers.py +class MaskedConv1D(nn.Module): + """1D convolution over the sequence axis of (B, F, L, C) inputs. + + Masked positions are zeroed before convolution and the output mask is + propagated based on whether the full kernel window contained valid values. + """ + + def __init__( + self, + filters: int, + kernel_size: int, + strides: int = 1, + padding: str = "valid", + dilation_rate: int = 1, + activation: Optional[str] = None, + use_bias: bool = True, + kernel_initializer: str = "glorot_uniform", + bias_initializer: str = "zeros", + kernel_regularizer: Optional[float] = None, + ): + super().__init__() + self.filters = filters + self.kernel_size = kernel_size + self.strides = strides + self.padding = padding.lower() + self.dilation_rate = dilation_rate + self.use_bias = use_bias + self.activation = activation + + self.conv = nn.Conv1d( + in_channels=filters, # placeholder; set in build + out_channels=filters, + kernel_size=kernel_size, + stride=strides, + padding=0, # handled manually + dilation=dilation_rate, + bias=use_bias, + ) + # Actual in_channels will be determined at first forward; override then. + + def _resolve_padding(self, length: int) -> int: + if self.padding == "same": + dilated = self.dilation_rate * (self.kernel_size - 1) + return (length + self.strides - 1) // self.strides * self.strides - length + dilated + return 0 + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + b, f, l, c = x.shape + # Lazy set conv in_channels on first call + if self.conv.in_channels != c: + self.conv = nn.Conv1d( + in_channels=c, + out_channels=self.filters, + kernel_size=self.kernel_size, + stride=self.strides, + padding=0, + dilation=self.dilation_rate, + bias=self.use_bias, + ).to(x.device) + + if mask is not None: + x = x * mask.unsqueeze(-1).to(x.dtype) + + # Merge batch and frame dims: (B*F, C, L) + x_2d = x.reshape(b * f, c, l) + if self.padding == "same": + pad_total = self._resolve_padding(l) + pad_left = pad_total // 2 + pad_right = pad_total - pad_left + x_2d = F.pad(x_2d, (pad_left, pad_right)) + + out = self.conv(x_2d) + _, _, l_out = out.shape + out = out.reshape(b, f, l_out, self.filters) + + if self.activation: + out = getattr(F, self.activation)(out) + + out_mask = None + if mask is not None: + mask_f = mask.reshape(b * f, 1, l).to(x.dtype) + if self.padding == "same": + mask_f = F.pad(mask_f, (pad_left, pad_right)) + with torch.no_grad(): + kernel = torch.ones( + (1, 1, self.kernel_size), dtype=mask_f.dtype, device=mask_f.device + ) + out_mask = F.conv1d( + mask_f, + kernel, + stride=self.strides, + dilation=self.dilation_rate, + ) + out_mask = (out_mask >= self.kernel_size).squeeze(1) + out_mask = out_mask.reshape(b, f, l_out) + + return out, out_mask +``` + +- [ ] **Step 4: Run the test to verify it passes** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_conv1d_shape_and_mask -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py +git commit -m "feat(pytorch): add MaskedConv1D with mask propagation" +``` + +--- + +## Task 4: Port MaskedBatchNorm and MaskedLayerNorm + +**Files:** +- Modify: `src/jaeger/nnlib/pytorch/layers.py` +- Modify: `tests/unit/nnlib/pytorch/test_layers.py` + +- [ ] **Step 1: Write the failing test** + +```python +def test_masked_batchnorm_output_shape_and_nmd(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[0, :, 5:] = False + layer = MaskedBatchNorm(num_features=4, return_nmd=True) + out, nmd = layer(x, mask) + assert out.shape == (2, 6, 10, 4) + assert nmd.shape == (2, 4) + + +def test_masked_layer_norm_shape(): + x = torch.randn(2, 6, 10, 4) + mask = torch.ones(2, 6, 10, dtype=torch.bool) + mask[0, :, 5:] = False + layer = MaskedLayerNorm(num_features=4) + out = layer(x, mask) + assert out.shape == (2, 6, 10, 4) +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_batchnorm_output_shape_and_nmd tests/unit/nnlib/pytorch/test_layers.py::test_masked_layer_norm_shape -v +``` + +Expected: `AttributeError` for missing classes. + +- [ ] **Step 3: Implement MaskedBatchNorm and MaskedLayerNorm** + +```python +# src/jaeger/nnlib/pytorch/layers.py +class MaskedBatchNorm(nn.Module): + """Batch normalization that excludes masked positions from statistics. + + Can optionally return normalized mean difference (nmd) vectors. + """ + + def __init__( + self, + num_features: int, + eps: float = 1e-5, + momentum: float = 0.9, + return_nmd: bool = False, + ): + super().__init__() + self.num_features = num_features + self.eps = eps + self.momentum = momentum + self.return_nmd = return_nmd + + self.gamma = nn.Parameter(torch.ones(num_features)) + self.beta = nn.Parameter(torch.zeros(num_features)) + self.register_buffer("running_mean", torch.zeros(num_features)) + self.register_buffer("running_var", torch.ones(num_features)) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + xf = x.to(torch.float32) + b, f, l, c = xf.shape + + if mask is not None: + mask_f = mask.unsqueeze(-1).to(torch.float32) + masked_x = xf * mask_f + valid_count = mask_f.sum(dim=(0, 1, 2)) + self.eps + mean = masked_x.sum(dim=(0, 1, 2)) / valid_count + var = ((masked_x - mean) * mask_f).pow(2).sum(dim=(0, 1, 2)) / valid_count + else: + mean = xf.mean(dim=(0, 1, 2)) + var = xf.var(dim=(0, 1, 2), unbiased=False) + + if self.training: + self.running_mean = self.momentum * self.running_mean + (1 - self.momentum) * mean.detach() + self.running_var = self.momentum * self.running_var + (1 - self.momentum) * var.detach() + mean_use, var_use = mean, var + else: + mean_use, var_use = self.running_mean, self.running_var + + mean_use = mean_use.view(1, 1, 1, -1) + var_use = var_use.view(1, 1, 1, -1) + normalized = (xf - mean_use) / torch.sqrt(var_use + self.eps) + out = normalized * self.gamma.view(1, 1, 1, -1) + self.beta.view(1, 1, 1, -1) + out = out.to(x.dtype) + + if self.return_nmd: + if mask is not None: + per_ex_sum = masked_x.sum(dim=(1, 2)) + per_ex_count = mask_f.sum(dim=(1, 2)) + self.eps + mean_channel = per_ex_sum / per_ex_count + else: + mean_channel = xf.mean(dim=(1, 2)) + nmd = (mean_channel - mean).to(x.dtype) + return out, nmd + + return out, None + + +class MaskedLayerNorm(nn.Module): + """Layer normalization that excludes masked positions.""" + + def __init__(self, num_features: int, eps: float = 1e-3): + super().__init__() + self.num_features = num_features + self.eps = eps + self.gamma = nn.Parameter(torch.ones(num_features)) + self.beta = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + xf = x.to(torch.float32) + if mask is not None: + mask_f = mask.unsqueeze(-1).to(torch.float32) + masked_x = xf * mask_f + count = mask_f.sum(dim=-1, keepdim=True) + self.eps + mean = masked_x.sum(dim=-1, keepdim=True) / count + var = ((masked_x - mean) * mask_f).pow(2).sum(dim=-1, keepdim=True) / count + else: + mean = xf.mean(dim=-1, keepdim=True) + var = xf.var(dim=-1, keepdim=True, unbiased=False) + + normalized = (xf - mean) / torch.sqrt(var + self.eps) + out = normalized * self.gamma.view(1, 1, 1, -1) + self.beta.view(1, 1, 1, -1) + if mask is not None: + out = out * mask_f + return out.to(x.dtype) +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_batchnorm_output_shape_and_nmd tests/unit/nnlib/pytorch/test_layers.py::test_masked_layer_norm_shape -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py +git commit -m "feat(pytorch): add MaskedBatchNorm and MaskedLayerNorm" +``` + +--- + +## Task 5: Port attention and pooling layers + +**Files:** +- Modify: `src/jaeger/nnlib/pytorch/layers.py` +- Modify: `tests/unit/nnlib/pytorch/test_layers.py` + +- [ ] **Step 1: Write the failing tests** + +```python +def test_gated_frame_pooling_shape(): + x = torch.randn(2, 6, 10, 4) + layer = GatedFrameGlobalMaxPooling(return_gate=False) + out = layer(x) + assert out.shape == (2, 4) + + +def test_axial_attention_shape(): + x = torch.randn(2, 6, 10, 4) + layer = AxialAttention(embed_dim=4, num_heads=2) + out, mask = layer(x) + assert out.shape == (2, 6, 10, 4) +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_gated_frame_pooling_shape tests/unit/nnlib/pytorch/test_layers.py::test_axial_attention_shape -v +``` + +Expected: `AttributeError` for missing classes. + +- [ ] **Step 3: Implement the layers** + +```python +# src/jaeger/nnlib/pytorch/layers.py +class GatedFrameGlobalMaxPooling(nn.Module): + """Frame-aware global max pooling. Input (B,F,L,D) -> output (B,D).""" + + def __init__(self, return_gate: bool = False): + super().__init__() + self.return_gate = return_gate + self.score_dense = nn.Linear(1, 1) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]: + # x: (B, F, L, D) + per_frame = x.max(dim=2)[0] # (B, F, D) + logits = self.score_dense(per_frame).squeeze(-1) # (B, F) + gates = torch.softmax(logits, dim=1) + pooled = (per_frame * gates.unsqueeze(-1)).sum(dim=1) # (B, D) + if self.return_gate: + return pooled, gates + return pooled + + +class AxialAttention(nn.Module): + """Axial attention over the sequence axis.""" + + def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0): + super().__init__() + self.attn = nn.MultiheadAttention( + embed_dim=embed_dim, + num_heads=num_heads, + dropout=dropout, + batch_first=True, + ) + self.norm = nn.LayerNorm(embed_dim) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + b, f, l, d = x.shape + # Reshape to (B*F, L, D) + x_2d = x.reshape(b * f, l, d) + key_padding_mask = None + if mask is not None: + key_padding_mask = ~mask.reshape(b * f, l) + attn_out, _ = self.attn( + x_2d, x_2d, x_2d, key_padding_mask=key_padding_mask, need_weights=False + ) + out = self.norm(x_2d + attn_out).reshape(b, f, l, d) + out_mask = mask + return out, out_mask +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py::test_gated_frame_pooling_shape tests/unit/nnlib/pytorch/test_layers.py::test_axial_attention_shape -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py +git commit -m "feat(pytorch): add attention and pooling layers" +``` + +--- + +## Task 5b: Implement remaining custom layers and RepresentationModel + +**Files:** +- Modify: `src/jaeger/nnlib/pytorch/layers.py` +- Modify: `src/jaeger/nnlib/pytorch/models.py` +- Modify: `tests/unit/nnlib/pytorch/test_layers.py` + +- [ ] **Step 1: Add SinusoidalPositionEmbedding, ResidualBlock, TransformerEncoder, CrossFrameAttention** + +```python +# src/jaeger/nnlib/pytorch/layers.py +class SinusoidalPositionEmbedding(nn.Module): + def __init__(self, max_wavelength: int = 10000): + super().__init__() + self.max_wavelength = max_wavelength + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, F, L, D) + b, f, l, d = x.shape + position = torch.arange(l, device=x.device).unsqueeze(1).float() + div_term = torch.exp( + torch.arange(0, d, 2, device=x.device).float() + * (-math.log(self.max_wavelength) / d) + ) + pe = torch.zeros(l, d, device=x.device) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + return x + pe.view(1, 1, l, d) + + +class ResidualBlock(nn.Module): + def __init__(self, layer: nn.Module): + super().__init__() + self.layer = layer + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + if hasattr(self.layer, "forward") and "mask" in self.layer.forward.__code__.co_varnames: + out, out_mask = self.layer(x, mask) + return out + x, out_mask + out = self.layer(x) + return out + x, mask + + +class TransformerEncoder(nn.Module): + def __init__(self, embed_dim: int, num_heads: int, num_layers: int = 1, dropout: float = 0.0): + super().__init__() + layer = nn.TransformerEncoderLayer( + d_model=embed_dim, + nhead=num_heads, + dim_feedforward=embed_dim * 4, + dropout=dropout, + batch_first=True, + ) + self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + b, f, l, d = x.shape + x_2d = x.reshape(b * f, l, d) + key_mask = None + if mask is not None: + key_mask = ~mask.reshape(b * f, l) + out = self.encoder(x_2d, src_key_padding_mask=key_mask) + return out.reshape(b, f, l, d), mask + + +class CrossFrameAttention(nn.Module): + def __init__(self, embed_dim: int, num_heads: int): + super().__init__() + self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) + self.norm = nn.LayerNorm(embed_dim) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): + b, f, l, d = x.shape + # Treat frames as sequence: (B*L, F, D) + x_t = x.permute(0, 2, 1, 3).reshape(b * l, f, d) + key_mask = None + if mask is not None: + key_mask = ~mask.permute(0, 2, 1).reshape(b * l, f) + out, _ = self.attn(x_t, x_t, x_t, key_padding_mask=key_mask, need_weights=False) + out = self.norm(x_t + out).reshape(b, l, f, d).permute(0, 2, 1, 3) + return out, mask +``` + +- [ ] **Step 2: Add RepresentationModel** + +```python +# src/jaeger/nnlib/pytorch/models.py +class RepresentationModel(nn.Module): + def __init__(self, embedding: nn.Module, hidden_layers: list, pooling: str = "average"): + super().__init__() + self.embedding = embedding + self.blocks = nn.ModuleList() + self.return_nmd = False + self.output_dim = None + self.nmd_dim = None + # Build blocks from config; simplified here + for cfg in hidden_layers: + name = cfg["name"].lower() + config = cfg.get("config", {}) + if name == "masked_conv1d": + self.blocks.append(MaskedConv1D(**config)) + elif name == "masked_batchnorm": + self.blocks.append(MaskedBatchNorm(**config)) + elif name == "masked_layer_norm": + self.blocks.append(MaskedLayerNorm(**config)) + elif name == "axial_attention": + self.blocks.append(AxialAttention(**config)) + elif name == "cross_frame_attention": + self.blocks.append(CrossFrameAttention(**config)) + elif name == "transformer_encoder": + self.blocks.append(TransformerEncoder(**config)) + elif name == "residual_block": + inner = self.blocks.pop() if self.blocks else None + self.blocks.append(ResidualBlock(inner)) + elif name == "dense": + self.blocks.append(nn.Linear(**config)) + elif name == "activation": + self.blocks.append(GeLU() if config.get("activation") == "gelu" else nn.ReLU()) + elif name == "dropout": + self.blocks.append(nn.Dropout(config.get("rate", 0.0))) + else: + raise ValueError(f"Unknown layer type: {name}") + self.pooler = self._build_pooler(pooling) + + def _build_pooler(self, pooling: str): + from jaeger.nnlib.pytorch.layers import GatedFrameGlobalMaxPooling, MaskedGlobalAvgPooling + if pooling == "average": + return MaskedGlobalAvgPooling() + elif pooling == "gatedframe": + return GatedFrameGlobalMaxPooling(return_gate=False) + raise ValueError(f"Unknown pooling: {pooling}") + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor | tuple: + x = self.embedding(x) + nmd = None + for block in self.blocks: + if hasattr(block, "return_nmd") and block.return_nmd: + x, nmd = block(x, mask) + else: + x, mask = block(x, mask) + pooled = self.pooler(x, mask) + if nmd is not None: + return pooled, nmd + return pooled +``` + +- [ ] **Step 3: Run tests** + +```bash +pytest tests/unit/nnlib/pytorch/test_layers.py -v +``` + +Expected: PASS + +- [ ] **Step 4: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/layers.py src/jaeger/nnlib/pytorch/models.py tests/unit/nnlib/pytorch/test_layers.py +git commit -m "feat(pytorch): add remaining custom layers and RepresentationModel" +``` + +--- + +## Task 6: Build PyTorch model modules + +**Files:** +- Create: `src/jaeger/nnlib/pytorch/models.py` +- Create: `tests/unit/nnlib/pytorch/test_models.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/nnlib/pytorch/test_models.py +import torch +from jaeger.nnlib.pytorch.models import ClassificationHead, RepresentationModel + + +def test_classification_head_output_shape(): + head = ClassificationHead(input_dim=64, num_classes=3, hidden_units=[128, 64]) + x = torch.randn(2, 64) + out = head(x) + assert out.shape == (2, 3) + + +def test_representation_model_output_shape(): + # Minimal config matching nn_config_500bp_baseline structure + model_cfg = { + "embedding": { + "input_type": "translated", + "use_embedding_layer": False, + "embedding_size": 64, + "input_shape": (6, None), + "vocab_size": 65, + "codon_depth": 1, + }, + "representation_learner": { + "hidden_layers": [ + {"name": "masked_conv1d", "config": {"filters": 32, "kernel_size": 7, "padding": "same"}} + ], + "pooling": "average", + }, + "classifier": {"input_shape": 32, "hidden_layers": [{"name": "dense", "config": {"units": 3}}]}, + } + # Build via builder in Task 9; for now test a small standalone module + assert True +``` + +- [ ] **Step 2: Run the test to verify it fails** + +```bash +pytest tests/unit/nnlib/pytorch/test_models.py -v +``` + +Expected: `ModuleNotFoundError` + +- [ ] **Step 3: Implement model modules** + +```python +# src/jaeger/nnlib/pytorch/models.py +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from jaeger.nnlib.pytorch.layers import GeLU + + +class Embedding(nn.Module): + """Translates DNA into 6-frame codon embeddings or nucleotide one-hot.""" + + def __init__( + self, + input_type: str, + vocab_size: Optional[int], + embedding_size: int, + use_embedding_layer: bool, + use_positional_embeddings: bool = False, + positional_embedding_length: Optional[int] = None, + ): + super().__init__() + self.input_type = input_type + self.use_embedding_layer = use_embedding_layer + self.use_positional_embeddings = use_positional_embeddings + + if input_type == "translated": + if use_embedding_layer: + self.embed = nn.Embedding(vocab_size, embedding_size, padding_idx=0) + nn.init.orthogonal_(self.embed.weight) + else: + self.embed = nn.Linear(vocab_size, embedding_size, bias=False) + nn.init.orthogonal_(self.embed.weight) + elif input_type == "nucleotide": + self.embed = None + else: + raise ValueError(f"Invalid input_type: {input_type}") + + if use_positional_embeddings: + from jaeger.nnlib.pytorch.layers import SinusoidalPositionEmbedding + + self.positional = SinusoidalPositionEmbedding(positional_embedding_length) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x shape depends on input_type; caller must pass correct shape + if self.input_type == "translated": + if self.use_embedding_layer: + return self.embed(x) + else: + return self.embed(x) + return x + + +class ClassificationHead(nn.Module): + """Dense head for class prediction.""" + + def __init__(self, input_dim: int, num_classes: int, hidden_units: list[int], dropout: float = 0.0): + super().__init__() + layers = [] + prev = input_dim + for units in hidden_units: + layers.append(nn.Linear(prev, units)) + layers.append(nn.LayerNorm(units)) + layers.append(GeLU()) + if dropout > 0: + layers.append(nn.Dropout(dropout)) + prev = units + layers.append(nn.Linear(prev, num_classes)) + self.net = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class ReliabilityHead(nn.Module): + """Dense head for confidence estimation.""" + + def __init__(self, input_dim: int, num_classes: int, hidden_units: list[int], dropout: float = 0.0): + super().__init__() + layers = [] + prev = input_dim + for units in hidden_units: + layers.append(nn.Linear(prev, units)) + layers.append(nn.LayerNorm(units)) + layers.append(GeLU()) + if dropout > 0: + layers.append(nn.Dropout(dropout)) + prev = units + layers.append(nn.Linear(prev, num_classes)) + self.net = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class ProjectionHead(nn.Module): + """Projection head for ArcFace self-supervised pretraining.""" + + def __init__(self, input_dim: int, projection_dim: int): + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, projection_dim), + nn.LayerNorm(projection_dim), + GeLU(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class JaegerModel(nn.Module): + """Combined model exposing all outputs for inference.""" + + def __init__( + self, + rep_model: nn.Module, + classification_head: nn.Module, + reliability_head: Optional[nn.Module] = None, + ): + super().__init__() + self.rep_model = rep_model + self.classification_head = classification_head + self.reliability_head = reliability_head + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + outputs = self.rep_model(x, mask) + if isinstance(outputs, tuple): + embedding, nmd = outputs[0], outputs[1] + gate = outputs[2] if len(outputs) > 2 else None + else: + embedding = outputs + nmd = None + gate = None + + prediction = self.classification_head(embedding) + result: Dict[str, torch.Tensor] = {"prediction": prediction, "embedding": embedding} + if nmd is not None: + result["nmd"] = nmd + if gate is not None: + result["gate"] = gate + if self.reliability_head is not None and nmd is not None: + result["reliability"] = self.reliability_head(nmd) + return result +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +pytest tests/unit/nnlib/pytorch/test_models.py -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/models.py tests/unit/nnlib/pytorch/test_models.py +git commit -m "feat(pytorch): add model modules (embedding, heads, JaegerModel)" +``` + +--- + +## Task 7: Port losses + +**Files:** +- Create: `src/jaeger/nnlib/pytorch/losses.py` +- Create: `tests/unit/nnlib/pytorch/test_losses.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/nnlib/pytorch/test_losses.py +import torch +from jaeger.nnlib.pytorch.losses import ArcFaceLoss + + +def test_arcface_loss_shape(): + loss = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.5, scale=30.0) + labels = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float32) + embeddings = torch.randn(3, 64) + out = loss(labels, embeddings) + assert out.ndim == 0 +``` + +- [ ] **Step 2: Run the test to verify it fails** + +```bash +pytest tests/unit/nnlib/pytorch/test_losses.py -v +``` + +Expected: `ModuleNotFoundError` + +- [ ] **Step 3: Implement losses** + +```python +# src/jaeger/nnlib/pytorch/losses.py +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ArcFaceLoss(nn.Module): + """ArcFace additive angular margin loss.""" + + def __init__( + self, + num_classes: int, + embedding_dim: int, + margin: float = 0.5, + scale: float = 30.0, + onehot: bool = True, + ): + super().__init__() + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.margin = margin + self.scale = scale + self.onehot = onehot + self.weight = nn.Parameter(torch.Tensor(num_classes, embedding_dim)) + nn.init.xavier_uniform_(self.weight) + + def forward(self, labels: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor: + # embeddings: (B, D), labels: (B, C) one-hot or (B,) long + embeddings_norm = F.normalize(embeddings, p=2, dim=1) + weight_norm = F.normalize(self.weight, p=2, dim=1) + cos_t = torch.matmul(embeddings_norm, weight_norm.t()) + + if self.onehot: + target = labels.argmax(dim=1) + else: + target = labels + + # Additive angular margin + cos_m = math.cos(self.margin) + sin_m = math.sin(self.margin) + sin_t = torch.sqrt(1.0 - cos_t.pow(2) + 1e-6) + cos_t_plus_m = cos_t * cos_m - sin_t * sin_m + one_hot = F.one_hot(target, num_classes=self.num_classes).to(cos_t.dtype) + logits = one_hot * cos_t_plus_m + (1.0 - one_hot) * cos_t + logits = logits * self.scale + return F.cross_entropy(logits, target) +``` + +- [ ] **Step 4: Run the test to verify it passes** + +```bash +pytest tests/unit/nnlib/pytorch/test_losses.py -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/losses.py tests/unit/nnlib/pytorch/test_losses.py +git commit -m "feat(pytorch): add ArcFace loss" +``` + +--- + +## Task 8: Port metrics + +**Files:** +- Create: `src/jaeger/nnlib/pytorch/metrics.py` +- Create: `tests/unit/nnlib/pytorch/test_metrics.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/nnlib/pytorch/test_metrics.py +import torch +from jaeger.nnlib.pytorch.metrics import PrecisionForClass + + +def test_precision_for_class(): + metric = PrecisionForClass(class_id=1) + preds = torch.tensor([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]) + labels = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.float32) + metric.update(preds, labels) + result = metric.compute() + assert 0.0 <= result <= 1.0 +``` + +- [ ] **Step 2: Run the test to verify it fails** + +```bash +pytest tests/unit/nnlib/pytorch/test_metrics.py -v +``` + +Expected: `ModuleNotFoundError` + +- [ ] **Step 3: Implement metrics** + +```python +# src/jaeger/nnlib/pytorch/metrics.py +import torch +import torch.nn.functional as F + + +class PrecisionForClass: + def __init__(self, class_id: int): + self.class_id = class_id + self.tp = 0 + self.fp = 0 + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + pred_labels = y_pred.argmax(dim=-1) + true_labels = y_true.argmax(dim=-1) + cls = self.class_id + self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) + self.fp += int(((pred_labels == cls) & (true_labels != cls)).sum().item()) + + def compute(self) -> float: + if self.tp + self.fp == 0: + return 0.0 + return self.tp / (self.tp + self.fp) + + +class RecallForClass(PrecisionForClass): + def __init__(self, class_id: int): + super().__init__(class_id) + self.fn = 0 + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + pred_labels = y_pred.argmax(dim=-1) + true_labels = y_true.argmax(dim=-1) + cls = self.class_id + self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) + self.fn += int(((pred_labels != cls) & (true_labels == cls)).sum().item()) + + def compute(self) -> float: + if self.tp + self.fn == 0: + return 0.0 + return self.tp / (self.tp + self.fn) +``` + +- [ ] **Step 4: Run the test to verify it passes** + +```bash +pytest tests/unit/nnlib/pytorch/test_metrics.py -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/metrics.py tests/unit/nnlib/pytorch/test_metrics.py +git commit -m "feat(pytorch): add per-class precision/recall metrics" +``` + +--- + +## Task 9: Implement ModelBuilder + +**Files:** +- Create: `src/jaeger/nnlib/pytorch/builder.py` +- Create: `tests/unit/nnlib/pytorch/test_builder.py` +- Modify: `src/jaeger/nnlib/pytorch/models.py` (add RepresentationModel) + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/nnlib/pytorch/test_builder.py +import torch +from jaeger.nnlib.pytorch.builder import ModelBuilder + + +def test_builder_creates_jaeger_model(): + config = { + "model": { + "name": "test_model", + "classifier_out_dim": 3, + "reliability_out_dim": 1, + "class_label_map": [{"class": "phage", "label": 1}, {"class": "bacteria", "label": 0}], + "embedding": { + "input_type": "translated", + "use_embedding_layer": False, + "embedding_size": 32, + "input_shape": [6, None], + "vocab_size": 65, + "codon_depth": 1, + }, + "string_processor": {"codon": "CODON", "codon_id": "CODON_ID"}, + "representation_learner": { + "hidden_layers": [ + {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}} + ], + "pooling": "average", + }, + "classifier": {"input_shape": 16, "hidden_layers": [{"name": "dense", "config": {"units": 3}}]}, + }, + "training": {"batch_size": 2, "optimizer": "adam", "optimizer_params": {"lr": 1e-3}}, + } + builder = ModelBuilder(config) + models = builder.build_fragment_classifier() + x = torch.randint(0, 65, (2, 6, 50)) + mask = torch.ones(2, 6, 50, dtype=torch.bool) + out = models["jaeger_model"](x, mask) + assert out["prediction"].shape == (2, 3) +``` + +- [ ] **Step 2: Run the test to verify it fails** + +```bash +pytest tests/unit/nnlib/pytorch/test_builder.py -v +``` + +Expected: `ModuleNotFoundError` + +- [ ] **Step 3: Implement ModelBuilder** + +```python +# src/jaeger/nnlib/pytorch/builder.py +from pathlib import Path +from typing import Any, Dict, Optional + +import torch +import torch.nn as nn + +from jaeger.nnlib.pytorch.layers import ( + AxialAttention, + CrossFrameAttention, + GatedFrameGlobalMaxPooling, + GeLU, + MaskedBatchNorm, + MaskedConv1D, + MaskedLayerNorm, + ResidualBlock, + TransformerEncoder, +) +from jaeger.nnlib.pytorch.losses import ArcFaceLoss +from jaeger.nnlib.pytorch.models import ( + ClassificationHead, + Embedding, + JaegerModel, + ProjectionHead, + ReliabilityHead, + RepresentationModel, +) + + +class ModelBuilder: + """Build PyTorch Jaeger models from YAML config.""" + + def __init__(self, config: Dict[str, Any]): + self.cfg = config + self.model_cfg = config.get("model", {}) + self.train_cfg = config.get("training", {}) + self.classifier_out_dim = int(self.model_cfg.get("classifier_out_dim", 0)) + self.reliability_out_dim = int(self.model_cfg.get("reliability_out_dim", 0)) + + def build_fragment_classifier(self) -> Dict[str, nn.Module]: + models: Dict[str, nn.Module] = {} + + embedding_cfg = self.model_cfg.get("embedding", {}) + embedding = Embedding( + input_type=embedding_cfg.get("input_type"), + vocab_size=embedding_cfg.get("vocab_size"), + embedding_size=embedding_cfg.get("embedding_size", 4), + use_embedding_layer=embedding_cfg.get("use_embedding_layer", False), + ) + + rep_cfg = self.model_cfg.get("representation_learner", {}) + rep_model = RepresentationModel( + embedding=embedding, + hidden_layers=rep_cfg.get("hidden_layers", []), + pooling=rep_cfg.get("pooling", "average"), + ) + models["rep_model"] = rep_model + + if "classifier" in self.model_cfg: + cls_cfg = self.model_cfg["classifier"] + head = ClassificationHead( + input_dim=cls_cfg.get("input_shape", rep_model.output_dim), + num_classes=self.classifier_out_dim, + hidden_units=[layer["config"]["units"] for layer in cls_cfg.get("hidden_layers", [])[:-1]], + ) + models["classification_head"] = head + models["jaeger_classifier"] = nn.Sequential(rep_model, head) + + if "reliability_model" in self.model_cfg: + rel_cfg = self.model_cfg["reliability_model"] + rel_head = ReliabilityHead( + input_dim=rel_cfg.get("input_shape", rep_model.nmd_dim), + num_classes=self.reliability_out_dim, + hidden_units=[layer["config"]["units"] for layer in rel_cfg.get("hidden_layers", [])[:-1]], + ) + models["reliability_head"] = rel_head + + models["jaeger_model"] = JaegerModel( + rep_model=rep_model, + classification_head=models["classification_head"], + reliability_head=models.get("reliability_head"), + ) + return models + + def compile_model(self, models: Dict[str, nn.Module], train_branch: str = "classifier") -> tuple: + opt_name = self.train_cfg.get("optimizer", "adam").lower() + opt_params = self.train_cfg.get("optimizer_params", {}) + opt_class = self._get_optimizer(opt_name) + + if train_branch == "classifier": + model = models["jaeger_classifier"] + optimizer = opt_class(model.parameters(), **opt_params) + loss = nn.CrossEntropyLoss() + return model, optimizer, loss + elif train_branch == "reliability": + model = models["jaeger_reliability"] + optimizer = opt_class(model.parameters(), **opt_params) + loss = nn.BCEWithLogitsLoss() + return model, optimizer, loss + elif train_branch == "pretrain": + model = models["jaeger_projection"] + optimizer = opt_class(model.parameters(), **opt_params) + loss = models["arcface_loss"] + return model, optimizer, loss + else: + raise ValueError(f"Unknown train_branch: {train_branch}") + + def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: + from jaeger.nnlib.pytorch.metrics import PrecisionForClass, RecallForClass + metrics = {} + out_dim = self.classifier_out_dim if branch == "classifier" else self.reliability_out_dim + for cls in range(out_dim): + metrics[f"precision_class_{cls}"] = PrecisionForClass(class_id=cls) + metrics[f"recall_class_{cls}"] = RecallForClass(class_id=cls) + return metrics + + def _get_optimizer(self, name: str, kwargs: Dict[str, Any]) -> torch.optim.Optimizer: + optimizers = { + "adam": torch.optim.Adam, + "sgd": torch.optim.SGD, + "rmsprop": torch.optim.RMSprop, + } + return optimizers[name] +``` + +Also update `models.py` to add `RepresentationModel` with proper output tracking. + +- [ ] **Step 4: Run the test to verify it passes** + +```bash +pytest tests/unit/nnlib/pytorch/test_builder.py -v +``` + +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/nnlib/pytorch/builder.py src/jaeger/nnlib/pytorch/models.py tests/unit/nnlib/pytorch/test_builder.py +git commit -m "feat(pytorch): add ModelBuilder" +``` + +--- + +## Task 10: Add parity tests against TensorFlow baseline + +**Files:** +- Create: `tests/integration/test_pytorch_parity.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/integration/test_pytorch_parity.py +import numpy as np +import pytest +import torch + +from jaeger.nnlib.builder import DynamicModelBuilder as TFBuilder +from jaeger.nnlib.pytorch.builder import ModelBuilder as PTBuilder + + +CONFIG = { + "model": { + "name": "parity_test", + "classifier_out_dim": 3, + "reliability_out_dim": 1, + "class_label_map": [{"class": "phage", "label": 1}, {"class": "bacteria", "label": 0}, {"class": "eukarya", "label": 2}], + "embedding": { + "input_type": "translated", + "use_embedding_layer": False, + "embedding_size": 32, + "input_shape": [6, None], + "vocab_size": 65, + "codon_depth": 1, + }, + "string_processor": { + "codon": "CODON", + "codon_id": "CODON_ID", + "crop_size": 50, + "input_type": "translated", + "use_embedding_layer": False, + }, + "representation_learner": { + "hidden_layers": [ + {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}} + ], + "pooling": "average", + }, + "classifier": {"input_shape": 16, "hidden_layers": [{"name": "dense", "config": {"units": 3}}]}, + }, + "training": {"batch_size": 2, "optimizer": "adam", "optimizer_params": {"lr": 1e-3}}, +} + + +def test_forward_parity(): + tf_builder = TFBuilder(CONFIG) + tf_models = tf_builder.build_fragment_classifier() + tf_model = tf_models["jaeger_model"] + + pt_builder = PTBuilder(CONFIG) + pt_models = pt_builder.build_fragment_classifier() + pt_model = pt_models["jaeger_model"] + pt_model.eval() + + _copy_tf_weights_to_pt(tf_model, pt_model) + + x = np.random.randint(0, 65, size=(2, 6, 50)).astype(np.int64) + tf_out = tf_model.predict({"translated": x}) + with torch.no_grad(): + pt_out = pt_model(torch.from_numpy(x), mask=torch.ones(2, 6, 50, dtype=torch.bool)) + + np.testing.assert_allclose( + tf_out["prediction"], + pt_out["prediction"].numpy(), + atol=1e-4, + rtol=1e-4, + ) + + +def _copy_tf_weights_to_pt(tf_model, pt_model): + """Map Keras layer weights to PyTorch state_dict by name and shape.""" + tf_layers = {layer.name: layer for layer in tf_model.layers} + pt_state = pt_model.state_dict() + new_state = {} + for pt_key, pt_tensor in pt_state.items(): + # Example mapping: "rep_model.blocks.0.conv.weight" -> "rep_conv1d_0/kernel:0" + # Implement name-based mapping here; raise ValueError if no match. + # This is architecture-specific and must be updated as layers are added. + raise NotImplementedError(f"Implement mapping for {pt_key}") + pt_model.load_state_dict(new_state) +``` + +- [ ] **Step 2: Implement weight-copy helper** + +Map TensorFlow layer names to PyTorch `state_dict` keys. For each `MaskedConv1D`/`Dense`/BatchNorm in the TF model, find the corresponding PyTorch module and copy `.kernel` → `.weight`, `.bias` → `.bias`, BN moving stats, gamma, beta. + +- [ ] **Step 3: Run the test** + +```bash +pytest tests/integration/test_pytorch_parity.py::test_forward_parity -v +``` + +Expected: PASS after weight mapping is correct. + +- [ ] **Step 4: Commit** + +```bash +git add tests/integration/test_pytorch_parity.py +git commit -m "test(pytorch): add TF vs PyTorch forward parity test" +``` + +--- + +## Task 11: Implement numpy_full PyTorch Dataset + +**Files:** +- Create: `src/jaeger/data/pytorch/dataset_numpy.py` +- Create: `src/jaeger/data/pytorch/collate.py` +- Create: `tests/unit/data/pytorch/test_dataset_numpy.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/data/pytorch/test_dataset_numpy.py +import numpy as np +import torch +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset + + +def test_numpy_full_dataset(tmp_path): + data = np.random.randint(0, 65, size=(10, 6, 50)).astype(np.int32) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] + path = tmp_path / "data.npz" + np.savez(path, translated=data, label=labels) + + ds = NumpyFullDataset(path, input_key="translated") + assert len(ds) == 10 + x, y = ds[0] + assert x.shape == (6, 50) + assert y.shape == (3,) + + loader = torch.utils.data.DataLoader(ds, batch_size=2) + batch_x, batch_y = next(iter(loader)) + assert batch_x.shape == (2, 6, 50) + assert batch_y.shape == (2, 3) +``` + +- [ ] **Step 2: Implement NumpyFullDataset** + +```python +# src/jaeger/data/pytorch/dataset_numpy.py +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +from torch.utils.data import Dataset + + +class NumpyFullDataset(Dataset): + """Loads a fully-preprocessed .npz file.""" + + def __init__(self, path: str | Path, input_key: str = "translated", label_key: str = "label"): + data = np.load(path, allow_pickle=False) + self.inputs = torch.from_numpy(data[input_key]) + self.labels = torch.from_numpy(data[label_key]) + + def __len__(self): + return len(self.labels) + + def __getitem__(self, idx): + x = self.inputs[idx] + mask = (x != 0).any(dim=0) if x.dim() == 3 else torch.ones(x.shape[-1], dtype=torch.bool) + return x, self.labels[idx], mask +``` + +- [ ] **Step 3: Run the test** + +```bash +pytest tests/unit/data/pytorch/test_dataset_numpy.py -v +``` + +Expected: PASS + +- [ ] **Step 4: Add dataset builder helper** + +```python +# src/jaeger/data/pytorch/builders.py +from typing import Any, Dict + +from torch.utils.data import DataLoader + +from jaeger.data.pytorch.dataset_csv import CSVDataset +from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset + + +def build_datasets(config: Dict[str, Any], branch: str = "classifier"): + model_cfg = config.get("model", {}) + train_cfg = config.get("training", {}) + string_cfg = model_cfg.get("string_processor", {}) + embedding_cfg = model_cfg.get("embedding", {}) + data_format = string_cfg.get("data_format", "csv") + batch_size = train_cfg.get("batch_size", 32) + crop_size = string_cfg.get("crop_size", 500) + num_classes = model_cfg.get("classifier_out_dim" if branch == "classifier" else "reliability_out_dim", 3) + + data_key = "fragment_classifier_data" if branch == "classifier" else "fragment_reliability_data" + data_cfg = train_cfg.get(data_key, {}) + + from jaeger.seqops.maps import CODON_ID + codon_table = {c: i for i, c in enumerate(CODON_ID)} + + datasets = {} + for split in ["train", "validation"]: + entries = data_cfg.get(split, []) + paths = [p for entry in entries for p in entry.get("path", [])] + if not paths: + raise ValueError(f"No paths found for {branch}/{split}") + + if data_format == "numpy_full": + datasets[split] = NumpyFullDataset(paths[0]) + elif data_format == "numpy_raw": + datasets[split] = NumpyRawDataset( + paths[0], + crop_size=crop_size, + num_classes=num_classes, + codon_table=codon_table, + shuffle=(split == "train"), + mutate=string_cfg.get("mutate", False), + mutation_rate=string_cfg.get("mutation_rate", 0.1), + shuffle_frames=string_cfg.get("shuffle_frames", False), + ) + elif data_format == "csv": + datasets[split] = CSVDataset( + paths[0], + seq_col=string_cfg.get("seq_col", 1), + class_col=string_cfg.get("class_col", 0), + crop_size=crop_size, + num_classes=num_classes, + codon_table=codon_table, + shuffle=(split == "train"), + mutate=string_cfg.get("mutate", False), + mutation_rate=string_cfg.get("mutation_rate", 0.1), + shuffle_frames=string_cfg.get("shuffle_frames", False), + ) + else: + raise ValueError(f"Unsupported data_format: {data_format}") + + return { + "train": DataLoader(datasets["train"], batch_size=batch_size, shuffle=True, num_workers=4), + "validation": DataLoader(datasets["validation"], batch_size=batch_size, shuffle=False, num_workers=4), + } +``` + +- [ ] **Step 5: Commit** + +```bash +git add src/jaeger/data/pytorch/builders.py src/jaeger/data/pytorch/dataset_numpy.py tests/unit/data/pytorch/test_dataset_numpy.py +git commit -m "feat(pytorch): add NumpyFullDataset and dataset builder" +``` + +--- + +## Task 12: Implement numpy_raw PyTorch Dataset + +**Files:** +- Modify: `src/jaeger/data/pytorch/dataset_numpy.py` +- Modify: `src/jaeger/data/pytorch/transforms.py` +- Modify: `tests/unit/data/pytorch/test_dataset_numpy.py` + +- [ ] **Step 1: Write the failing test** + +```python +def test_numpy_raw_dataset(tmp_path): + seqs = np.random.randint(0, 4, size=(10, 500)).astype(np.int8) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] + path = tmp_path / "raw.npz" + np.savez(path, sequences=seqs, labels=labels) + + from jaeger.seqops.maps import CODON_ID + codon_table = {c: i for i, c in enumerate(CODON_ID)} + ds = NumpyRawDataset( + path, + seq_key="sequences", + label_key="labels", + crop_size=50, + num_classes=3, + codon_table=codon_table, + ) + x, y, mask = ds[0] + assert x.shape == (6, 50) + assert mask.shape == (6, 50) +``` + +- [ ] **Step 2: Implement transforms and dataset** + +```python +# src/jaeger/data/pytorch/transforms.py +import random +from typing import Dict, Optional + +import numpy as np +import torch + + +def _reverse_complement(seq_str: str) -> str: + complement = {"A": "T", "T": "A", "G": "C", "C": "G"} + return "".join(complement.get(b, "N") for b in reversed(seq_str.upper())) + + +def translate_to_codons(seq: np.ndarray, codon_table: Dict[str, int]) -> torch.Tensor: + # seq: (L,) int8 nucleotide indices + # returns (6, L//3) translated codon indices for 6 reading frames + nucleotides = ["A", "T", "G", "C"] + seq_str = "".join(nucleotides[i] for i in seq) + frames = [] + for offset in range(3): + codons = [seq_str[i : i + 3] for i in range(offset, len(seq_str) - 2, 3)] + indices = [codon_table.get(c, 0) for c in codons] + rev = _reverse_complement(seq_str) + rev_codons = [rev[i : i + 3] for i in range(offset, len(rev) - 2, 3)] + rev_indices = [codon_table.get(c, 0) for c in rev_codons] + frames.append(torch.tensor(indices, dtype=torch.long)) + frames.append(torch.tensor(rev_indices, dtype=torch.long)) + return torch.stack(frames) + + +def dna_to_indices(seq: str) -> np.ndarray: + mapping = {"A": 0, "T": 1, "G": 2, "C": 3} + return np.array([mapping.get(base.upper(), 0) for base in seq], dtype=np.int8) + + +def apply_mutation(seq: np.ndarray, rate: float) -> np.ndarray: + if rate == 0: + return seq + mask = np.random.rand(len(seq)) < rate + mutated = seq.copy() + mutated[mask] = np.random.randint(0, 4, size=mask.sum()) + return mutated + + +def shuffle_frames(x: torch.Tensor) -> torch.Tensor: + # x: (6, L) + frames = list(range(6)) + random.shuffle(frames) + return x[frames] +``` + +```python +# src/jaeger/data/pytorch/dataset_numpy.py +class NumpyRawDataset(Dataset): + """Loads raw int8 sequences and applies runtime preprocessing.""" + + def __init__( + self, + path: str | Path, + seq_key: str = "sequences", + label_key: str = "labels", + crop_size: int = 500, + num_classes: int = 3, + codon_table: Optional[Dict[str, int]] = None, + shuffle: bool = True, + mutate: bool = False, + mutation_rate: float = 0.1, + shuffle_frames: bool = False, + ): + data = np.load(path, allow_pickle=False) + self.seqs = data[seq_key] + self.labels = torch.from_numpy(data[label_key]) + self.crop_size = crop_size + self.num_classes = num_classes + self.codon_table = codon_table + self.shuffle = shuffle + self.mutate = mutate + self.mutation_rate = mutation_rate + self.shuffle_frames = shuffle_frames + + def __len__(self): + return len(self.labels) + + def __getitem__(self, idx): + seq = self.seqs[idx] + if self.mutate: + seq = apply_mutation(seq, self.mutation_rate) + x = translate_to_codons(seq, self.codon_table) + # crop / pad to crop_size + l = x.shape[1] + if l > self.crop_size: + start = random.randint(0, l - self.crop_size) + x = x[:, start : start + self.crop_size] + elif l < self.crop_size: + pad = self.crop_size - l + x = torch.nn.functional.pad(x, (0, pad)) + mask = (x != 0) + if self.shuffle_frames: + x = shuffle_frames(x) + return x, self.labels[idx], mask +``` + +- [ ] **Step 3: Run the test** + +```bash +pytest tests/unit/data/pytorch/test_dataset_numpy.py::test_numpy_raw_dataset -v +``` + +Expected: PASS + +- [ ] **Step 4: Commit** + +```bash +git add src/jaeger/data/pytorch/transforms.py src/jaeger/data/pytorch/dataset_numpy.py tests/unit/data/pytorch/test_dataset_numpy.py +git commit -m "feat(pytorch): add NumpyRawDataset and transforms" +``` + +--- + +## Task 13: Implement CSV PyTorch Dataset + +**Files:** +- Create: `src/jaeger/data/pytorch/dataset_csv.py` +- Create: `tests/unit/data/pytorch/test_dataset_csv.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/data/pytorch/test_dataset_csv.py +import csv +import torch +from jaeger.data.pytorch.dataset_csv import CSVDataset + + +def test_csv_dataset(tmp_path): + path = tmp_path / "data.csv" + with open(path, "w", newline="") as fh: + writer = csv.writer(fh) + for i in range(10): + writer.writerow([i % 3, "ATGC" * 20, f"seq{i}"]) + + ds = CSVDataset(path, seq_col=1, class_col=0, crop_size=50, num_classes=3) + x, y, mask = ds[0] + assert x.shape[0] == 6 + assert y.shape == (3,) +``` + +- [ ] **Step 2: Implement CSVDataset** + +```python +# src/jaeger/data/pytorch/dataset_csv.py +import csv +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons + + +class CSVDataset(Dataset): + def __init__( + self, + path: str | Path, + seq_col: int, + class_col: int, + crop_size: int, + num_classes: int, + codon_table: dict, + shuffle: bool = True, + mutate: bool = False, + mutation_rate: float = 0.1, + shuffle_frames: bool = False, + ): + self.samples = [] + with open(path) as fh: + reader = csv.reader(fh) + for row in reader: + label = int(row[class_col]) + seq = row[seq_col] + self.samples.append((seq, label)) + self.crop_size = crop_size + self.num_classes = num_classes + self.codon_table = codon_table + self.shuffle = shuffle + self.mutate = mutate + self.mutation_rate = mutation_rate + self.shuffle_frames = shuffle_frames + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + seq, label = self.samples[idx] + indices = dna_to_indices(seq) + if self.mutate: + indices = apply_mutation(indices, self.mutation_rate) + x = translate_to_codons(indices, self.codon_table) + # crop/pad + l = x.shape[1] + if l > self.crop_size: + start = random.randint(0, l - self.crop_size) + x = x[:, start : start + self.crop_size] + elif l < self.crop_size: + x = torch.nn.functional.pad(x, (0, self.crop_size - l)) + mask = (x != 0) + y = torch.nn.functional.one_hot(torch.tensor(label), num_classes=self.num_classes).float() + return x, y, mask +``` + +- [ ] **Step 3: Run the test** + +```bash +pytest tests/unit/data/pytorch/test_dataset_csv.py -v +``` + +Expected: PASS + +- [ ] **Step 4: Commit** + +```bash +git add src/jaeger/data/pytorch/dataset_csv.py tests/unit/data/pytorch/test_dataset_csv.py +git commit -m "feat(pytorch): add CSVDataset" +``` + +--- + +## Task 14: Implement training engine + +**Files:** +- Create: `src/jaeger/training/pytorch/engine.py` +- Create: `tests/unit/training/pytorch/test_engine.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/unit/training/pytorch/test_engine.py +import torch +from jaeger.training.pytorch.engine import train_one_epoch, validate_one_epoch + + +def test_train_one_epoch_runs(): + model = torch.nn.Linear(10, 2) + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + loss_fn = torch.nn.CrossEntropyLoss() + loader = [(torch.randn(2, 10), torch.tensor([0, 1]))] + metrics = {} + result = train_one_epoch(model, loader, optimizer, loss_fn, device="cpu", use_amp=False, metrics=metrics) + assert "loss" in result +``` + +- [ ] **Step 2: Implement engine** + +```python +# src/jaeger/training/pytorch/engine.py +from typing import Any, Callable, Dict, Optional + +import torch +from torch.cuda.amp import GradScaler, autocast + + +def train_one_epoch( + model: torch.nn.Module, + dataloader, + optimizer: torch.optim.Optimizer, + loss_fn: Callable, + device: torch.device, + use_amp: bool = False, + metrics: Optional[Dict[str, Callable]] = None, + max_steps: Optional[int] = None, +) -> Dict[str, float]: + model.train() + scaler = GradScaler() if use_amp else None + total_loss = 0.0 + num_batches = 0 + + for step, batch in enumerate(dataloader): + if max_steps is not None and step >= max_steps: + break + optimizer.zero_grad() + batch = _to_device(batch, device) + if len(batch) == 3: + x, y, mask = batch + else: + x, y = batch + mask = None + + with autocast(device_type=device.type, enabled=use_amp): + outputs = model(x, mask) + # If y is one-hot and loss_fn is CrossEntropyLoss, convert to class indices + if y.dim() > 1 and y.shape[-1] > 1 and isinstance(loss_fn, torch.nn.CrossEntropyLoss): + y_idx = y.argmax(dim=-1) + loss = loss_fn(outputs["prediction"], y_idx) + else: + loss = loss_fn(outputs, y) + + if use_amp: + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + loss.backward() + optimizer.step() + + total_loss += loss.item() + num_batches += 1 + if metrics: + for name, metric in metrics.items(): + metric.update(outputs["prediction"].detach(), y) + + result = {"loss": total_loss / max(num_batches, 1)} + if metrics: + for name, metric in metrics.items(): + result[name] = metric.compute() + return result + + +def validate_one_epoch( + model: torch.nn.Module, + dataloader, + loss_fn: Callable, + device: torch.device, + use_amp: bool = False, + metrics: Optional[Dict[str, Callable]] = None, + max_steps: Optional[int] = None, +) -> Dict[str, float]: + model.eval() + total_loss = 0.0 + num_batches = 0 + with torch.no_grad(): + for step, batch in enumerate(dataloader): + if max_steps is not None and step >= max_steps: + break + batch = _to_device(batch, device) + if len(batch) == 3: + x, y, mask = batch + else: + x, y = batch + mask = None + with autocast(device_type=device.type, enabled=use_amp): + outputs = model(x, mask) + if y.dim() > 1 and y.shape[-1] > 1 and isinstance(loss_fn, torch.nn.CrossEntropyLoss): + y_idx = y.argmax(dim=-1) + loss = loss_fn(outputs["prediction"], y_idx) + else: + loss = loss_fn(outputs, y) + total_loss += loss.item() + num_batches += 1 + if metrics: + for name, metric in metrics.items(): + metric.update(outputs["prediction"], y) + + result = {"loss": total_loss / max(num_batches, 1)} + if metrics: + for name, metric in metrics.items(): + result[name] = metric.compute() + return result + + +def _to_device(batch, device): + if isinstance(batch, (tuple, list)): + return tuple(t.to(device) if isinstance(t, torch.Tensor) else t for t in batch) + return batch.to(device) +``` + +- [ ] **Step 3: Run the test** + +```bash +pytest tests/unit/training/pytorch/test_engine.py -v +``` + +Expected: PASS + +- [ ] **Step 4: Commit** + +```bash +git add src/jaeger/training/pytorch/engine.py tests/unit/training/pytorch/test_engine.py +git commit -m "feat(pytorch): add training engine" +``` + +--- + +## Task 15: Implement Trainer + +**Files:** +- Create: `src/jaeger/training/pytorch/trainer.py` +- Modify: `src/jaeger/nnlib/pytorch/checkpoints.py` + +- [ ] **Step 1: Implement checkpoint helper** + +```python +# src/jaeger/nnlib/pytorch/checkpoints.py +import torch + + +def save_checkpoint(path, model, optimizer, epoch, metrics, branch, config): + torch.save( + { + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "epoch": epoch, + "metrics": metrics, + "branch": branch, + "config": config, + }, + path, + ) + + +def load_checkpoint(path, model, optimizer=None): + ckpt = torch.load(path, map_location="cpu") + model.load_state_dict(ckpt["model_state_dict"]) + if optimizer is not None and "optimizer_state_dict" in ckpt: + optimizer.load_state_dict(ckpt["optimizer_state_dict"]) + return ckpt +``` + +- [ ] **Step 2: Implement Trainer** + +```python +# src/jaeger/training/pytorch/trainer.py +import json +from pathlib import Path +from typing import Any, Dict, Optional + +import torch +import yaml +from torch.utils.data import DataLoader + +from jaeger.nnlib.pytorch.builder import ModelBuilder +from jaeger.nnlib.pytorch.checkpoints import load_checkpoint, save_checkpoint +from jaeger.training.pytorch.callbacks import CallbackList +from jaeger.training.pytorch.engine import train_one_epoch, validate_one_epoch +from jaeger.utils.logging import get_logger + + +logger = get_logger(log_file=None, log_path=None, level=3) + + +class Trainer: + """Orchestrates classifier/reliability/pretrain training.""" + + def __init__( + self, + config: Dict[str, Any], + device: torch.device, + use_amp: bool = False, + checkpoint_dir: Optional[Path] = None, + ): + self.config = config + self.device = device + self.use_amp = use_amp + self.checkpoint_dir = checkpoint_dir + self.builder = ModelBuilder(config) + self.models = self.builder.build_fragment_classifier() + self.callbacks = CallbackList() + + def fit_classifier( + self, + train_loader: DataLoader, + val_loader: DataLoader, + epochs: int, + initial_epoch: int = 0, + ): + model, optimizer, loss_fn = self.builder.compile_model(self.models, train_branch="classifier") + model = model.to(self.device) + metrics = self.builder.get_metrics(branch="classifier") + + for epoch in range(initial_epoch, epochs): + train_metrics = train_one_epoch( + model, train_loader, optimizer, loss_fn, self.device, self.use_amp, metrics + ) + val_metrics = validate_one_epoch( + model, val_loader, loss_fn, self.device, self.use_amp, metrics + ) + self.callbacks.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) + + def fit_reliability(self, train_loader, val_loader, epochs, initial_epoch=0): + model, optimizer, loss_fn = self.builder.compile_model(self.models, train_branch="reliability") + model = model.to(self.device) + metrics = self.builder.get_metrics(branch="reliability") + for epoch in range(initial_epoch, epochs): + train_metrics = train_one_epoch(model, train_loader, optimizer, loss_fn, self.device, self.use_amp, metrics) + val_metrics = validate_one_epoch(model, val_loader, loss_fn, self.device, self.use_amp, metrics) + self.callbacks.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) + + def fit_projection(self, train_loader, val_loader, epochs, initial_epoch=0): + model, optimizer, loss_fn = self.builder.compile_model(self.models, train_branch="pretrain") + model = model.to(self.device) + for epoch in range(initial_epoch, epochs): + train_metrics = train_one_epoch(model, train_loader, optimizer, loss_fn, self.device, self.use_amp) + val_metrics = validate_one_epoch(model, val_loader, loss_fn, self.device, self.use_amp) + self.callbacks.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) + + def save_model(self, suffix: str = "fragment", metadata: Optional[str] = None): + from jaeger.nnlib.pytorch.checkpoints import save_checkpoint + import yaml + + model_name = self.config["model"]["name"] + experiment = self.config["model"].get("experiment", "1") + base_dir = Path(self.config["training"].get("base_dir", ".")) + save_dir = base_dir / f"{model_name}_exp{experiment}" + save_dir.mkdir(parents=True, exist_ok=True) + + torch.save(self.models["jaeger_model"].state_dict(), save_dir / "model.pt") + with open(save_dir / "model_config.yaml", "w") as fh: + yaml.safe_dump(self.config, fh) + with open(save_dir / "class_label_map.yaml", "w") as fh: + yaml.safe_dump({"classes": self.config["model"].get("class_label_map", [])}, fh) + + if metadata: + Path(metadata).write_text(json.dumps({"model_path": str(save_dir)}, indent=2)) + + logger.info(f"Saved PyTorch model to {save_dir}") + return save_dir +``` + +- [ ] **Step 3: Smoke test** + +```bash +python -c "from jaeger.training.pytorch.trainer import Trainer; print('OK')" +``` + +Expected: `OK` + +- [ ] **Step 4: Commit** + +```bash +git add src/jaeger/training/pytorch/trainer.py src/jaeger/nnlib/pytorch/checkpoints.py +git commit -m "feat(pytorch): add Trainer and checkpoint helpers" +``` + +--- + +## Task 16: Implement callbacks + +**Files:** +- Create: `src/jaeger/training/pytorch/callbacks.py` +- Create: `tests/unit/training/pytorch/test_callbacks.py` + +- [ ] **Step 1: Implement callbacks** + +```python +# src/jaeger/training/pytorch/callbacks.py +from pathlib import Path + +import torch + + +class CallbackList: + def __init__(self, callbacks=None): + self.callbacks = callbacks or [] + + def on_epoch_end(self, epoch, train_metrics, val_metrics, model, optimizer): + for cb in self.callbacks: + cb.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) + + +class ModelCheckpoint: + def __init__(self, checkpoint_dir: Path, monitor: str = "val_loss", mode: str = "min"): + self.checkpoint_dir = Path(checkpoint_dir) + self.checkpoint_dir.mkdir(parents=True, exist_ok=True) + self.monitor = monitor + self.mode = mode + self.best = float("inf") if mode == "min" else float("-inf") + + def on_epoch_end(self, epoch, train_metrics, val_metrics, model, optimizer): + value = val_metrics.get(self.monitor) + if value is None: + return + improved = (self.mode == "min" and value < self.best) or (self.mode == "max" and value > self.best) + if improved: + self.best = value + path = self.checkpoint_dir / f"epoch:{epoch}-loss:{value:.4f}.pt" + torch.save( + { + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "epoch": epoch, + "metrics": val_metrics, + }, + path, + ) +``` + +- [ ] **Step 2: Smoke test** + +```bash +pytest tests/unit/training/pytorch/test_callbacks.py -v +``` + +Expected: PASS + +- [ ] **Step 3: Commit** + +```bash +git add src/jaeger/training/pytorch/callbacks.py tests/unit/training/pytorch/test_callbacks.py +git commit -m "feat(pytorch): add training callbacks" +``` + +--- + +## Task 17: Implement distributed training helper + +**Files:** +- Create: `src/jaeger/training/pytorch/distributed.py` + +- [ ] **Step 1: Implement DDP setup** + +```python +# src/jaeger/training/pytorch/distributed.py +import os + +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + + +def setup_distributed(): + rank = int(os.environ.get("RANK", 0)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + dist.init_process_group("nccl") + torch.cuda.set_device(local_rank) + return rank, local_rank, world_size + + +def cleanup_distributed(): + dist.destroy_process_group() + + +def wrap_model_ddp(model, device_ids): + return DDP(model, device_ids=device_ids) +``` + +- [ ] **Step 2: Smoke test** + +```bash +python -c "from jaeger.training.pytorch.distributed import setup_distributed; print('OK')" +``` + +Expected: `OK` (does nothing if not launched with torchrun) + +- [ ] **Step 3: Commit** + +```bash +git add src/jaeger/training/pytorch/distributed.py +git commit -m "feat(pytorch): add distributed training helpers" +``` + +--- + +## Task 18: Update `jaeger train` command + +**Files:** +- Modify: `src/jaeger/commands/train.py` + +- [ ] **Step 1: Replace TF training orchestration with PyTorch** + +Keep the Click options. Replace `train_fragment_core` body with: + +```python +# src/jaeger/commands/train.py +def train_fragment_core(**kwargs): + import torch + from jaeger.training.pytorch.trainer import Trainer + from jaeger.data.pytorch.builders import build_datasets + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + config = load_model_config(Path(kwargs["config"])) + config["mix_precision"] = kwargs.get("mixed_precision", False) + config["from_last_checkpoint"] = kwargs.get("from_last_checkpoint", False) + config["force"] = kwargs.get("force", False) + + trainer = Trainer( + config=config, + device=device, + use_amp=kwargs.get("mixed_precision", False), + ) + + classifier_loaders = build_datasets(config, branch="classifier") + reliability_loaders = build_datasets(config, branch="reliability") + + train_cfg = config.get("training", {}) + + if kwargs.get("self_supervised_pretraining", False): + trainer.fit_projection( + classifier_loaders["train"], + classifier_loaders["validation"], + epochs=train_cfg.get("projection_epochs", 10), + ) + + if not kwargs.get("only_reliability_head", False): + trainer.fit_classifier( + classifier_loaders["train"], + classifier_loaders["validation"], + epochs=train_cfg.get("classifier_epochs", 100), + ) + + if not kwargs.get("only_classification_head", False): + trainer.fit_reliability( + reliability_loaders["train"], + reliability_loaders["validation"], + epochs=train_cfg.get("reliability_epochs", 100), + ) + + if kwargs.get("save_model", False) or kwargs.get("only_save", False): + trainer.save_model(suffix="fragment", metadata=kwargs.get("meta")) +``` + +- [ ] **Step 2: Smoke test** + +```bash +jaeger train --help +``` + +Expected: help text shows PyTorch-specific options. + +- [ ] **Step 3: Commit** + +```bash +git add src/jaeger/commands/train.py +git commit -m "feat(pytorch): wire jaeger train to PyTorch trainer" +``` + +--- + +## Task 19: Implement PyTorch inference backend + +**Files:** +- Create: `src/jaeger/inference/pytorch/model.py` +- Create: `src/jaeger/inference/pytorch/engine.py` + +- [ ] **Step 1: Implement model loader** + +```python +# src/jaeger/inference/pytorch/model.py +from pathlib import Path + +import torch +import yaml + +from jaeger.nnlib.pytorch.builder import ModelBuilder + + +class PyTorchJaegerModel: + def __init__(self, model_dir: Path, device: torch.device): + self.model_dir = Path(model_dir) + self.device = device + with open(self.model_dir / "model_config.yaml") as fh: + config = yaml.safe_load(fh) + builder = ModelBuilder(config) + self.models = builder.build_fragment_classifier() + self.jaeger_model = self.models["jaeger_model"].to(device) + self.jaeger_model.load_state_dict(torch.load(self.model_dir / "model.pt", map_location=device)) + self.jaeger_model.eval() + + from jaeger.seqops.maps import CODON_ID + self.codon_table = {c: i for i, c in enumerate(CODON_ID)} + + def predict(self, x, mask=None): + with torch.no_grad(): + return self.jaeger_model(x.to(self.device), mask.to(self.device) if mask is not None else None) +``` + +- [ ] **Step 2: Implement inference engine** + +```python +# src/jaeger/inference/pytorch/engine.py +import numpy as np +import torch +from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons +from jaeger.inference.pytorch.model import PyTorchJaegerModel + + +def run_inference(model_dir, fasta_path, batch_size=96, device="cuda", fsize=2000, stride=2000, **kwargs): + import pyfastx + from jaeger.inference.pytorch.model import PyTorchJaegerModel + + device = torch.device(device if torch.cuda.is_available() else "cpu") + model = PyTorchJaegerModel(model_dir, device) + + results = [] + for name, seq in pyfastx.Fasta(fasta_path, build_index=False): + windows = [] + for i in range(0, max(1, len(seq) - fsize + 1), stride): + window = seq[i : i + fsize] + if len(window) < 100: + continue + indices = dna_to_indices(window) + windows.append((name, len(seq), i, indices)) + + batch_preds = [] + for start in range(0, len(windows), batch_size): + batch = windows[start : start + batch_size] + lengths = [len(w[3]) for w in batch] + max_len = max(lengths) + padded = torch.zeros(len(batch), 6, max_len, dtype=torch.long) + mask = torch.zeros(len(batch), 6, max_len, dtype=torch.bool) + for j, (_, _, _, indices) in enumerate(batch): + translated = translate_to_codons(indices, model.codon_table) + l = translated.shape[1] + padded[j, :, :l] = translated + mask[j, :, :l] = True + with torch.no_grad(): + out = model.predict(padded, mask) + batch_preds.append(out["prediction"].cpu().numpy()) + + results.append({ + "name": name, + "length": len(seq), + "predictions": np.concatenate(batch_preds) if batch_preds else np.array([]), + }) + + return results +``` + +- [ ] **Step 3: Smoke test** + +```bash +python -c "from jaeger.inference.pytorch.model import PyTorchJaegerModel; print('OK')" +``` + +Expected: `OK` + +- [ ] **Step 4: Commit** + +```bash +git add src/jaeger/inference/pytorch/model.py src/jaeger/inference/pytorch/engine.py +git commit -m "feat(pytorch): add PyTorch inference backend" +``` + +--- + +## Task 20: Update `jaeger predict` command + +**Files:** +- Modify: `src/jaeger/commands/predict.py` + +- [ ] **Step 1: Add PyTorch model detection and routing** + +```python +# src/jaeger/commands/predict.py +def run_core(**kwargs): + from jaeger.inference.pytorch.engine import run_inference + return run_inference(**kwargs) +``` + +- [ ] **Step 2: Smoke test** + +```bash +jaeger predict --help +``` + +Expected: help text unchanged. + +- [ ] **Step 3: Commit** + +```bash +git add src/jaeger/commands/predict.py +git commit -m "feat(pytorch): route jaeger predict to PyTorch backend when model.pt present" +``` + +--- + +## Task 21: Update SLURM scripts and container + +**Files:** +- Modify: `slurm/*.slurm` +- Modify: `singularity/jaeger_singularity.def` + +- [ ] **Step 1: Update a training SLURM script** + +Edit `slurm/baseline_500bp.slurm`. Replace the training invocation line: + +```bash +# Before: +jaeger train -c train_config/nn_config_500bp_baseline.yaml + +# After: +torchrun --nproc_per_node=2 --nnodes=1 jaeger train -c train_config/nn_config_500bp_baseline.yaml --mixed_precision +``` + +- [ ] **Step 2: Update Singularity definition** + +Edit `singularity/jaeger_singularity.def`. Replace TensorFlow install lines with: + +```singularity +pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 +``` + +- [ ] **Step 3: Commit** + +```bash +git add slurm/ singularity/ +git commit -m "feat(pytorch): update SLURM scripts and Singularity container for PyTorch" +``` + +--- + +## Task 22: Delete TensorFlow training and inference code + +**Files:** +- Delete: `src/jaeger/nnlib/v1/` +- Delete: `src/jaeger/nnlib/v2/` +- Delete: `src/jaeger/nnlib/builder.py` +- Delete: `src/jaeger/nnlib/inference.py` +- Delete: `src/jaeger/nnlib/conversion.py` +- Delete: `src/jaeger/commands/predict_legacy.py` + +- [ ] **Step 1: Remove directories and files** + +```bash +git rm -rf src/jaeger/nnlib/v1/ src/jaeger/nnlib/v2/ +git rm src/jaeger/nnlib/builder.py src/jaeger/nnlib/inference.py src/jaeger/nnlib/conversion.py +git rm src/jaeger/commands/predict_legacy.py +``` + +- [ ] **Step 2: Verify imports still work** + +```bash +python -c "import jaeger; print('OK')" +``` + +Expected: `OK` + +- [ ] **Step 3: Commit** + +```bash +git commit -m "refactor(pytorch): remove TensorFlow training and inference code" +``` + +--- + +## Task 23: Final QA, benchmarks, and release prep + +**Files:** +- Modify: `docs/_source/train.md` +- Modify: `docs/_source/optimizations.md` +- Modify: `recipes/jaeger-bio/meta.yaml` +- Modify: `install.sh` + +- [ ] **Step 1: Run full test suite** + +```bash +pytest tests/unit tests/integration -v +``` + +Expected: All tests pass. + +- [ ] **Step 2: Run smoke training on Zeus** + +```bash +sbatch slurm/baseline_500bp.slurm +``` + +Expected: Job completes, checkpoints saved, logs show faster batches/sec than TF baseline. + +- [ ] **Step 3: Run inference smoke test** + +```bash +jaeger predict -i tests/fixtures/small.fasta -o /tmp/predict_out --model_path path/to/pytorch/model +``` + +Expected: Predictions written without errors. + +- [ ] **Step 4: Update documentation** + +- `docs/_source/train.md`: replace TensorFlow training instructions with PyTorch examples. +- `docs/_source/optimizations.md`: update inference backend table; remove TF-specific optimizations. + +- [ ] **Step 5: Update Bioconda recipe and install script** + +```yaml +# recipes/jaeger-bio/meta.yaml +requirements: + run: + - pytorch + - torchvision +``` + +```bash +# install.sh +# Replace TensorFlow install lines with PyTorch install lines. +``` + +- [ ] **Step 6: Bump version to 1.27.1** + +```bash +.github/scripts/bump-version.sh 1 27 1 +``` + +- [ ] **Step 7: Final commit and push** + +```bash +git add -A +git commit -m "chore(release): bump version to 1.27.1" +git push origin pytorch_migration +``` + +--- diff --git a/docs/superpowers/specs/2026-06-14-pytorch-migration-design.md b/docs/superpowers/specs/2026-06-14-pytorch-migration-design.md new file mode 100644 index 0000000..48164e7 --- /dev/null +++ b/docs/superpowers/specs/2026-06-14-pytorch-migration-design.md @@ -0,0 +1,299 @@ +# PyTorch Migration Design + +## Context + +Jaeger's training stack currently uses TensorFlow/Keras. Empirical testing shows TensorFlow is too slow for architectures that use attention mechanisms and energy-based training regimes. This design migrates the training pipeline — and the inference path that consumes trained models — to PyTorch in a single release. + +## Goals + +- Replace TensorFlow training with a plain PyTorch training stack. +- Migrate `jaeger predict` to support PyTorch-trained models natively. +- Maintain the existing YAML config format so current `train_config/` files remain valid. +- Validate parity with the TensorFlow baseline before merging. +- Run on the Zeus cluster using the existing Singularity + SLURM workflow. + +## Non-goals + +- Loading existing TensorFlow/Keras checkpoints into PyTorch. This is a clean break; models are trained from scratch. +- Supporting TensorFlow and PyTorch training backends simultaneously. The `pytorch_migration` branch becomes PyTorch-only for training. +- Loading old TensorFlow SavedModels in the PyTorch release. The release is a clean break for inference as well; users who need legacy models can stay on the previous Jaeger version. + +## Clarifying decisions + +| Question | Decision | +|---|---| +| Load existing TF checkpoints? | No — train from scratch. | +| PyTorch framework | Plain PyTorch (not Lightning or HF). | +| Data loading | Native `torch.utils.data.Dataset` + `DataLoader`. | +| Inference | Migrate to PyTorch in the same release. | +| Next version | `1.27.1` after merge. | + +## Section 1 — High-level architecture & module layout + +The migration creates a new, self-contained PyTorch stack under `src/jaeger/` while leaving the existing TensorFlow code untouched on the `pytorch_migration` branch. The TF code is deleted only at the very end of the branch, once PyTorch training and inference are validated. + +``` +src/jaeger/ +├── nnlib/ +│ └── pytorch/ # NEW: PyTorch model library +│ ├── __init__.py +│ ├── layers.py # PyTorch equivalents of custom Keras layers +│ ├── models.py # RepModel, ClassifierHead, ReliabilityHead, JaegerModel +│ ├── losses.py # ArcFace, hierarchical loss, classification losses +│ ├── metrics.py # Per-class precision/recall/specificity +│ ├── builder.py # Build nn.Modules from YAML config +│ └── checkpoints.py # Save/load .pt checkpoints + metadata +├── data/ +│ └── pytorch/ # NEW: PyTorch datasets & loaders +│ ├── __init__.py +│ ├── dataset_csv.py # CSV Dataset with padding & augmentations +│ ├── dataset_numpy.py # numpy_full / numpy_raw / numpy_raw_variable +│ ├── dataset_tfrecord.py # TFRecord Dataset (deferred) +│ ├── transforms.py # codon translation, mutation, frame shuffle +│ └── collate.py # padding collate functions +├── training/ +│ └── pytorch/ # NEW: training loops +│ ├── __init__.py +│ ├── trainer.py # Classification / reliability / pretrain loops +│ ├── engine.py # epoch loop, validation, mixed precision +│ ├── distributed.py # DDP setup/teardown +│ └── callbacks.py # ModelCheckpoint, EarlyStopping, LR scheduling +├── inference/ +│ └── pytorch/ # NEW: PyTorch inference backend +│ ├── __init__.py +│ ├── model.py # Load PyTorch checkpoint + config +│ └── engine.py # Batch inference, windowing, aggregation +├── commands/ +│ ├── train.py # MODIFIED: dispatch to PyTorch trainer +│ └── predict.py # MODIFIED: add PyTorch backend option +``` + +Key principles: +- The YAML config format stays the same so existing `train_config/` files remain valid. +- `DynamicModelBuilder` is replaced by `jaeger.nnlib.pytorch.builder.ModelBuilder`, but it reads the same config. +- Training entry point `jaeger train` detects the backend from config or CLI flag; on the `pytorch_migration` branch it defaults to PyTorch. +- Existing TF training and inference files (`nnlib/v1/`, `nnlib/v2/`, current `commands/train.py`, current `commands/predict.py`, etc.) remain during development for reference and are removed before merge. The PyTorch release does not load legacy TF models. + +## Section 2 — Model architecture + +The current Keras builder assembles four pieces: + +1. **Representation learner** — embedding + stack of residual/attention blocks → `(embedding, nmd, [gate])` +2. **Classification head** — dense layers → class logits +3. **Reliability head** — dense layers → confidence score +4. **Projection head** (optional) — for ArcFace self-supervised pretraining + +In PyTorch, each becomes an `nn.Module`: + +- `Embedding` — maps codon indices or one-hot nucleotides to `(B, 6, L, D)` or `(B, 6, L, 3, D)` depending on config. +- `ResidualBlock`, `AxialAttention`, `CrossFrameAttention`, `TransformerEncoder` — replicate the Keras custom layers in `nnlib/pytorch/layers.py`. +- `RepresentationModel` — composes embedding + blocks + pooling → returns `(embedding, nmd, gate)` tuple. +- `ClassificationHead` / `ReliabilityHead` / `ProjectionHead` — small MLPs. +- `JaegerModel` — combines all heads and returns the same output dict as the Keras model: `{"prediction": ..., "embedding": ..., "nmd": ..., "gate": ..., "reliability": ...}`. + +Key design choices: +- **Masking**: Keras uses `Masking` layers and `supports_masking`. In PyTorch we pass a `mask` tensor explicitly and use it in `MaskedBatchNorm`, `MaskedConv1D`, pooling, and attention. +- **Variable length**: PyTorch handles padded batches naturally; we use a collate function to pad sequences and pass the mask. +- **Mixed precision**: Use `torch.cuda.amp.autocast` + `GradScaler` (or `torch.amp` if using newer PyTorch). +- **Weight init**: Keep the same initializers where possible (orthogonal for embeddings/dense, glorot for conv) to make parity tests fair. + +The builder reads the same YAML config and constructs the module graph, preserving hyperparameters (filters, kernel sizes, attention heads, dropout, regularization). + +## Section 3 — Data pipeline + +Current `tf.data` supports five formats: `csv`, `tfrecord`, `numpy_raw`, `numpy_raw_variable`, `numpy_full`. The PyTorch pipeline replaces this with `torch.utils.data.Dataset` + `DataLoader`. + +Priority order: + +1. **Phase B1 — `numpy_full`** (highest priority) + - Fastest loading, no runtime augmentations. + - `Dataset` loads the `.npz` once and returns preprocessed tensors. + - `DataLoader` handles shuffling and batching; no padding needed because all sequences are already cropped. + +2. **Phase B2 — `numpy_raw`** + - Returns int8 sequences + labels. + - Apply codon translation, n-gram extraction, frame shuffle, mutation in a `transform` callable. + - Collate function pads to the batch max length and produces the mask. + +3. **Phase B3 — `csv`** + - Parse CSV lines, apply the same transforms as `numpy_raw`. + - Slower, kept for backward compatibility and small experiments. + +4. **Phase B4 — `tfrecord` / `numpy_raw_variable`** (deferred if time-constrained) + - TFRecord is lower priority now that NumPy formats exist. + - Variable-length can be handled with a custom collate similar to `numpy_raw`. + +Augmentations (from `process_string_train`) move to `src/jaeger/data/pytorch/transforms.py`: +- codon translation +- frame shuffling +- random mutation +- masking +- n-gram encoding + +Each `Dataset` returns a tuple `(input_tensor, label_tensor, mask_tensor)`. The collate function stacks inputs and labels and combines masks. + +For multi-GPU, use `DistributedSampler` with `DataLoader`. + +## Section 4 — Training loop + +The current training flow in `commands/train.py` is: + +1. Build models inside a distribution strategy scope. +2. Compile classifier, reliability, and optionally pretrain branches. +3. Call `.fit()` on each branch sequentially. +4. Save weights + SavedModel graphs. + +In PyTorch, this becomes an explicit training engine: + +- `Trainer` class in `src/jaeger/training/pytorch/trainer.py` owns: + - model + - optimizer(s) + - loss function(s) + - metric trackers + - device / DDP wrapper + - AMP gradient scaler + +- Three training modes, matching the Keras flow: + 1. **Self-supervised pretraining** — train `rep_model` + `projection_head` with ArcFace loss. + 2. **Classifier training** — train `rep_model` + `classification_head`. + 3. **Reliability training** — freeze `rep_model`, train `reliability_head`. + +- `Engine` in `src/jaeger/training/pytorch/engine.py` implements one epoch: + - training step with `autocast`, loss backward, gradient clipping, optimizer step + - validation step + - metric aggregation + - logging + +- `callbacks.py` implements: + - `ModelCheckpoint` — save best/last `.pt` checkpoints + - `EarlyStopping` + - `ReduceLROnPlateau` + - `TensorBoardLogger` / CSV logger + +- Multi-GPU on Zeus: + - `torchrun --nproc_per_node=2 ... jaeger train ...` + - `DistributedDataParallel` wrapper + - `DistributedSampler` for each DataLoader + +- Resume logic: + - Checkpoint file contains model state_dict, optimizer state_dict, epoch, and config. + - `--from_last_checkpoint` loads the latest checkpoint and continues. + +- CLI flags preserved: + - `--mixed_precision`, `--xla` (replaced with AMP), `--from_last_checkpoint`, `--force`, `--only_classification_head`, `--only_reliability_head`, `--only_heads`, `--save_model`, `--only_save`, `--self_supervised_pretraining`. + +## Section 5 — Inference migration + +Current `jaeger predict` supports multiple backends (TF SavedModel, ONNX, TFLite, TensorRT). For the PyTorch release, we add a first-class PyTorch backend. + +- Save format: + - `model.pt` — model state_dict + - `model_config.yaml` — architecture config (same as training config) + - `class_label_map.yaml` — class mapping + - Optional: `model.pt` can be a TorchScript module if tracing works for the variable-length inputs. + +- `src/jaeger/inference/pytorch/model.py`: + - Load config, instantiate `JaegerModel`, load state_dict. + - Move to device (CUDA if available). + - Set to eval mode. + +- `src/jaeger/inference/pytorch/engine.py`: + - Replicate the windowing + batching logic from the current `predict.py`. + - Run forward pass, collect `prediction`, `embedding`, `reliability`. + - Support `--precision fp16` via `autocast`. + - Support `--batch` and `--workers`. + +- CLI integration: + - `jaeger predict` expects PyTorch models (look for `model.pt` + `model_config.yaml`). + - Legacy TensorFlow SavedModels are not loaded by this release. + +- Quantization/optimization: + - Not in the first PyTorch release unless explicitly required. + - Can be added later via `torch.compile` or ONNX export from PyTorch. + +## Section 6 — Checkpointing, AMP, distributed training, and dependencies + +- **Checkpoints**: PyTorch checkpoints are dictionaries containing: + - `model_state_dict` + - `optimizer_state_dict` + - `epoch` + - `config` + - `metrics` + - `branch` (classifier / reliability / projection) + - File naming: `epoch:{epoch}-loss:{val_loss:.4f}.pt` + +- **Mixed precision**: Use `torch.amp.autocast(device_type="cuda")` with `GradScaler`. The Keras `mixed_float16` policy maps directly. + +- **Distributed training**: + - Launch via `torchrun` / `torch.distributed.launch`. + - `DistributedDataParallel` wraps the model. + - `DistributedSampler` for each DataLoader. + - Only the rank-0 process saves checkpoints and logs. + +- **Dependencies**: + - Replace `tensorflow[and-cuda]` with `torch` and `torchvision`. + - Keep `numpy`, `pyyaml`, `click`, `rich`, `scipy`, `pandas`, `polars`, `pyfastx`. + - Add `tensorboard` if using TensorBoard logging. + - Remove TF-only dependencies from the PyTorch branch's training path. + +- **Zeus cluster**: + - Update SLURM scripts to use `torchrun --nproc_per_node=2`. + - Ensure Singularity container has PyTorch + CUDA installed. + - Bind `Jaeger/` source directory at runtime as per cluster rules. + +## Section 7 — Testing & validation strategy + +To make sure the PyTorch models match or exceed TF quality: + +1. **Unit tests** for each new PyTorch module: + - `nnlib/pytorch/layers.py` — forward pass shape tests, mask behavior. + - `nnlib/pytorch/builder.py` — build models from existing configs. + - `data/pytorch/dataset_*.py` — dataset length, batch shapes, augmentation correctness. + - `training/pytorch/engine.py` — one training step, checkpoint save/load. + +2. **Parity tests**: + - Build the same architecture in TF and PyTorch with fixed random seeds. + - Compare forward-pass outputs on identical synthetic input to within tolerance. + - Compare a few training steps on a tiny dataset to verify gradients and weight updates are similar. + +3. **Smoke tests**: + - `jaeger train -c train_config/nn_config_500bp_baseline.yaml` runs for a few epochs. + - `jaeger predict` with a PyTorch-trained model on a small FASTA. + +4. **Benchmarks**: + - Measure batches/sec on Zeus `gpu` partition for attention-based configs. + - Compare against TF baseline on the same hardware. + +5. **CI**: + - Add a GitHub Actions job that installs the PyTorch branch and runs unit + smoke tests. + - Keep the TF tests running until the branch is merged. + +## Section 8 — Branch, release, and migration plan + +- **Branch**: `pytorch_migration` created from `main`. +- **Version policy**: Bump to `1.27.1` when the branch is merged. During development, keep the current version. +- **Phase milestones**: + 1. Model architecture + parity tests + 2. `numpy_full` data loader + classifier training loop + 3. Reliability + pretraining loops + 4. CSV / `numpy_raw` loaders + 5. PyTorch inference backend + 6. SLURM script updates + Zeus benchmarks + 7. Delete TF training/inference code, update docs, final QA +- **Merge criteria**: + - All unit and smoke tests pass. + - Parity tests show acceptable agreement with TF on the same config. + - Zeus benchmark shows faster training for attention/energy-based configs. + - `jaeger predict` works end-to-end with PyTorch-trained models. + - Bioconda recipe and install script updated to PyTorch dependencies. +- **Release**: After merge, tag and release as normal using `.github/scripts/bump-version.sh`. + +## Open questions / risks + +| Risk | Mitigation | +|---|---| +| Custom Keras layers (masked conv, batch norm, attention) are complex to port exactly. | Write parity tests layer-by-layer; fix mismatches before end-to-end training. | +| Mixed precision behavior differs between TF and PyTorch. | Test AMP separately; make it opt-in at first. | +| Zeus Singularity container may need a new PyTorch image. | Build and test the container early in Phase 6. | +| Variable-length padding/collate may be slower than `tf.data`. | Benchmark `numpy_full` first; optimize collate if needed. | +| Removing TF may break other commands (e.g., `quantize`, `convert_graph`, `taxonomy`). | Audit all `commands/` before deleting TF; keep TF-only commands if they have no PyTorch equivalent yet. | diff --git a/entities.json b/entities.json new file mode 100644 index 0000000..1a247de --- /dev/null +++ b/entities.json @@ -0,0 +1,13 @@ +{ + "people": [ + "Yasas Wijesekara" + ], + "projects": [ + "jaeger-bio", + "PyPI", + "TensorRT", + "Conda", + "Jaeger" + ], + "topics": [] +} \ No newline at end of file diff --git a/mempalace.yaml b/mempalace.yaml new file mode 100644 index 0000000..b92b69f --- /dev/null +++ b/mempalace.yaml @@ -0,0 +1,60 @@ +wing: jaeger +rooms: +- name: slurm + description: Files from slurm/ + keywords: + - slurm + - slurm +- name: testing + description: Files from tests/ + keywords: + - testing + - tests +- name: scripts + description: Files from scripts/ + keywords: + - scripts + - scripts +- name: src + description: Files from src/ + keywords: + - src + - src +- name: singularity + description: Files from singularity/ + keywords: + - singularity + - singularity +- name: train_config + description: Files from train_config/ + keywords: + - train_config + - train_config +- name: documentation + description: Files from docs/ + keywords: + - documentation + - docs +- name: memory + description: Files from memory/ + keywords: + - memory + - memory +- name: recipes + description: Files from recipes/ + keywords: + - recipes + - recipes +- name: test_log + description: Files from test_log/ + keywords: + - test_log + - test_log +- name: test_cli + description: Files from test_cli/ + keywords: + - test_cli + - test_cli +- name: general + description: Files that don't fit other rooms + keywords: [] From 2757cf0b0d4be76f13653bce2611b56424859e41 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 13:41:57 +0200 Subject: [PATCH 38/64] revert: keep MemPalace/planning files untracked --- .../plans/2026-06-14-pytorch-migration.md | 2703 ----------------- .../2026-06-14-pytorch-migration-design.md | 299 -- entities.json | 13 - mempalace.yaml | 60 - 4 files changed, 3075 deletions(-) delete mode 100644 docs/superpowers/plans/2026-06-14-pytorch-migration.md delete mode 100644 docs/superpowers/specs/2026-06-14-pytorch-migration-design.md delete mode 100644 entities.json delete mode 100644 mempalace.yaml diff --git a/docs/superpowers/plans/2026-06-14-pytorch-migration.md b/docs/superpowers/plans/2026-06-14-pytorch-migration.md deleted file mode 100644 index 2c43515..0000000 --- a/docs/superpowers/plans/2026-06-14-pytorch-migration.md +++ /dev/null @@ -1,2703 +0,0 @@ -# PyTorch Migration Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Replace Jaeger's TensorFlow/Keras training and inference stack with a plain PyTorch implementation in the `pytorch_migration` branch, targeting release `1.27.1`. - -**Architecture:** Build a self-contained PyTorch module tree under `src/jaeger/nnlib/pytorch/`, `src/jaeger/data/pytorch/`, `src/jaeger/training/pytorch/`, and `src/jaeger/inference/pytorch/`. Preserve the YAML config format. Delete TensorFlow training/inference code only after parity tests and end-to-end smoke tests pass. - -**Tech Stack:** Python 3.11+, PyTorch 2.x, NumPy, Click, PyYAML, pytest. - ---- - -## File structure - -New files to create: - -``` -src/jaeger/nnlib/pytorch/ - __init__.py - layers.py # GeLU, MaskedConv1D, MaskedBatchNorm, MaskedLayerNorm, - # MaskedGlobalAvgPooling, GatedFrameGlobalMaxPooling, - # AxialAttention, CrossFrameAttention, TransformerEncoder, - # ResidualBlock, custom pooling helpers - models.py # Embedding, RepresentationModel, ClassificationHead, - # ReliabilityHead, ProjectionHead, JaegerModel - losses.py # ArcFaceLoss, HierarchicalLoss, classification losses - metrics.py # Per-class precision/recall/specificity - builder.py # ModelBuilder from YAML config - checkpoints.py # save/load .pt checkpoints - -src/jaeger/data/pytorch/ - __init__.py - transforms.py # codon translation, mutation, frame shuffle, masking - collate.py # padding collators - dataset_numpy.py # NumpyFullDataset, NumpyRawDataset, NumpyRawVariableDataset - dataset_csv.py # CSVDataset - dataset_tfrecord.py # TFRecordDataset (deferred) - -src/jaeger/training/pytorch/ - __init__.py - engine.py # Epoch loop, train/validation step, AMP - trainer.py # Trainer for classifier/reliability/pretrain branches - distributed.py # DDP setup/teardown - callbacks.py # ModelCheckpoint, EarlyStopping, LR scheduling, loggers - -src/jaeger/inference/pytorch/ - __init__.py - model.py # Load JaegerModel from checkpoint + config - engine.py # Batch inference, windowing, aggregation - -src/jaeger/commands/ - train.py # Dispatch to PyTorch trainer - predict.py # Dispatch to PyTorch inference engine -``` - -Test files to create: - -``` -tests/unit/nnlib/pytorch/ - test_layers.py - test_models.py - test_builder.py - test_losses.py - test_metrics.py -tests/unit/data/pytorch/ - test_dataset_numpy.py - test_transforms.py -tests/unit/training/pytorch/ - test_engine.py - test_callbacks.py -tests/integration/test_pytorch_parity.py -``` - -Files to delete at the end (after validation): - -``` -src/jaeger/nnlib/v1/ -src/jaeger/nnlib/v2/ -src/jaeger/nnlib/builder.py (old Keras builder) -src/jaeger/nnlib/inference.py (old TF inference) -src/jaeger/nnlib/conversion.py (if TF-only) -src/jaeger/commands/predict_legacy.py -src/jaeger/commands/train.py (old TF train command) -``` - ---- - -## Task 1: Create branch and update dependency metadata - -**Files:** -- Modify: `pyproject.toml` -- Modify: `.cz.toml` - -- [ ] **Step 1: Create the `pytorch_migration` branch** - -```bash -git checkout -b pytorch_migration -``` - -- [ ] **Step 2: Add PyTorch dependencies to `pyproject.toml`** - -Replace the TensorFlow optional dependencies with PyTorch equivalents. Keep TensorFlow as an optional `[legacy]` extra during development; remove it only in Task 27. - -```toml -[project.optional-dependencies] -gpu = [ - "torch >=2.0", - "torchvision", -] -cpu = [ - "torch >=2.0", - "torchvision", -] -darwin-arm = [ - "torch >=2.0", - "torchvision", -] -legacy = [ - "tensorflow >=2.21, <2.22", -] -test = [ - "pytest >=8.0", - "pytest-mock >=3.14", -] -``` - -- [ ] **Step 3: Verify `pyproject.toml` parses** - -```bash -python -c "import tomllib; tomllib.load(open('pyproject.toml','rb'))" -``` - -Expected: no exception. - -- [ ] **Step 4: Commit** - -```bash -git add pyproject.toml .cz.toml -git commit -m "chore(pytorch): add pytorch_migration branch and PyTorch deps" -``` - ---- - -## Task 2: Port activation and masking helper layers - -**Files:** -- Create: `src/jaeger/nnlib/pytorch/__init__.py` -- Create: `src/jaeger/nnlib/pytorch/layers.py` -- Create: `tests/unit/nnlib/pytorch/test_layers.py` - -- [ ] **Step 1: Write the failing test for GeLU** - -```python -# tests/unit/nnlib/pytorch/test_layers.py -import torch -import pytest -from jaeger.nnlib.pytorch.layers import GeLU - - -def test_gelu_matches_torch_nn_gelu(): - x = torch.tensor([-1.0, 0.0, 1.0, 2.0]) - layer = GeLU() - out = layer(x) - expected = torch.nn.functional.gelu(x, approximate="tanh") - assert torch.allclose(out, expected, atol=1e-6) -``` - -- [ ] **Step 2: Run the test to verify it fails** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_gelu_matches_torch_nn_gelu -v -``` - -Expected: `ModuleNotFoundError: No module named 'jaeger.nnlib.pytorch.layers'` - -- [ ] **Step 3: Implement GeLU and helpers** - -```python -# src/jaeger/nnlib/pytorch/__init__.py -# (empty or re-export later) -``` - -```python -# src/jaeger/nnlib/pytorch/layers.py -import math -from typing import Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class GeLU(nn.Module): - """Tanh-approximated GELU for TFLite-compatible graph export.""" - - def __init__(self): - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return F.gelu(x, approximate="tanh") -``` - -- [ ] **Step 4: Run the test to verify it passes** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_gelu_matches_torch_nn_gelu -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/__init__.py src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py -git commit -m "feat(pytorch): add GeLU activation layer" -``` - ---- - -## Task 3: Port MaskedConv1D - -**Files:** -- Modify: `src/jaeger/nnlib/pytorch/layers.py` -- Modify: `tests/unit/nnlib/pytorch/test_layers.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/nnlib/pytorch/test_layers.py -from jaeger.nnlib.pytorch.layers import MaskedConv1D - - -def test_masked_conv1d_shape_and_mask(): - # (B=2, F=6, L=10, C=4) - x = torch.randn(2, 6, 10, 4) - mask = torch.ones(2, 6, 10, dtype=torch.bool) - mask[0, :, 5:] = False - layer = MaskedConv1D(filters=8, kernel_size=3, padding="same") - out, out_mask = layer(x, mask) - assert out.shape == (2, 6, 10, 8) - assert out_mask.shape == (2, 6, 10) - assert not out_mask[0, :, 5:].any() -``` - -- [ ] **Step 2: Run the test to verify it fails** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_conv1d_shape_and_mask -v -``` - -Expected: `AttributeError: module 'jaeger.nnlib.pytorch.layers' has no attribute 'MaskedConv1D'` - -- [ ] **Step 3: Implement MaskedConv1D** - -```python -# src/jaeger/nnlib/pytorch/layers.py -class MaskedConv1D(nn.Module): - """1D convolution over the sequence axis of (B, F, L, C) inputs. - - Masked positions are zeroed before convolution and the output mask is - propagated based on whether the full kernel window contained valid values. - """ - - def __init__( - self, - filters: int, - kernel_size: int, - strides: int = 1, - padding: str = "valid", - dilation_rate: int = 1, - activation: Optional[str] = None, - use_bias: bool = True, - kernel_initializer: str = "glorot_uniform", - bias_initializer: str = "zeros", - kernel_regularizer: Optional[float] = None, - ): - super().__init__() - self.filters = filters - self.kernel_size = kernel_size - self.strides = strides - self.padding = padding.lower() - self.dilation_rate = dilation_rate - self.use_bias = use_bias - self.activation = activation - - self.conv = nn.Conv1d( - in_channels=filters, # placeholder; set in build - out_channels=filters, - kernel_size=kernel_size, - stride=strides, - padding=0, # handled manually - dilation=dilation_rate, - bias=use_bias, - ) - # Actual in_channels will be determined at first forward; override then. - - def _resolve_padding(self, length: int) -> int: - if self.padding == "same": - dilated = self.dilation_rate * (self.kernel_size - 1) - return (length + self.strides - 1) // self.strides * self.strides - length + dilated - return 0 - - def forward( - self, x: torch.Tensor, mask: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - b, f, l, c = x.shape - # Lazy set conv in_channels on first call - if self.conv.in_channels != c: - self.conv = nn.Conv1d( - in_channels=c, - out_channels=self.filters, - kernel_size=self.kernel_size, - stride=self.strides, - padding=0, - dilation=self.dilation_rate, - bias=self.use_bias, - ).to(x.device) - - if mask is not None: - x = x * mask.unsqueeze(-1).to(x.dtype) - - # Merge batch and frame dims: (B*F, C, L) - x_2d = x.reshape(b * f, c, l) - if self.padding == "same": - pad_total = self._resolve_padding(l) - pad_left = pad_total // 2 - pad_right = pad_total - pad_left - x_2d = F.pad(x_2d, (pad_left, pad_right)) - - out = self.conv(x_2d) - _, _, l_out = out.shape - out = out.reshape(b, f, l_out, self.filters) - - if self.activation: - out = getattr(F, self.activation)(out) - - out_mask = None - if mask is not None: - mask_f = mask.reshape(b * f, 1, l).to(x.dtype) - if self.padding == "same": - mask_f = F.pad(mask_f, (pad_left, pad_right)) - with torch.no_grad(): - kernel = torch.ones( - (1, 1, self.kernel_size), dtype=mask_f.dtype, device=mask_f.device - ) - out_mask = F.conv1d( - mask_f, - kernel, - stride=self.strides, - dilation=self.dilation_rate, - ) - out_mask = (out_mask >= self.kernel_size).squeeze(1) - out_mask = out_mask.reshape(b, f, l_out) - - return out, out_mask -``` - -- [ ] **Step 4: Run the test to verify it passes** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_conv1d_shape_and_mask -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py -git commit -m "feat(pytorch): add MaskedConv1D with mask propagation" -``` - ---- - -## Task 4: Port MaskedBatchNorm and MaskedLayerNorm - -**Files:** -- Modify: `src/jaeger/nnlib/pytorch/layers.py` -- Modify: `tests/unit/nnlib/pytorch/test_layers.py` - -- [ ] **Step 1: Write the failing test** - -```python -def test_masked_batchnorm_output_shape_and_nmd(): - x = torch.randn(2, 6, 10, 4) - mask = torch.ones(2, 6, 10, dtype=torch.bool) - mask[0, :, 5:] = False - layer = MaskedBatchNorm(num_features=4, return_nmd=True) - out, nmd = layer(x, mask) - assert out.shape == (2, 6, 10, 4) - assert nmd.shape == (2, 4) - - -def test_masked_layer_norm_shape(): - x = torch.randn(2, 6, 10, 4) - mask = torch.ones(2, 6, 10, dtype=torch.bool) - mask[0, :, 5:] = False - layer = MaskedLayerNorm(num_features=4) - out = layer(x, mask) - assert out.shape == (2, 6, 10, 4) -``` - -- [ ] **Step 2: Run the tests to verify they fail** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_batchnorm_output_shape_and_nmd tests/unit/nnlib/pytorch/test_layers.py::test_masked_layer_norm_shape -v -``` - -Expected: `AttributeError` for missing classes. - -- [ ] **Step 3: Implement MaskedBatchNorm and MaskedLayerNorm** - -```python -# src/jaeger/nnlib/pytorch/layers.py -class MaskedBatchNorm(nn.Module): - """Batch normalization that excludes masked positions from statistics. - - Can optionally return normalized mean difference (nmd) vectors. - """ - - def __init__( - self, - num_features: int, - eps: float = 1e-5, - momentum: float = 0.9, - return_nmd: bool = False, - ): - super().__init__() - self.num_features = num_features - self.eps = eps - self.momentum = momentum - self.return_nmd = return_nmd - - self.gamma = nn.Parameter(torch.ones(num_features)) - self.beta = nn.Parameter(torch.zeros(num_features)) - self.register_buffer("running_mean", torch.zeros(num_features)) - self.register_buffer("running_var", torch.ones(num_features)) - - def forward( - self, x: torch.Tensor, mask: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - xf = x.to(torch.float32) - b, f, l, c = xf.shape - - if mask is not None: - mask_f = mask.unsqueeze(-1).to(torch.float32) - masked_x = xf * mask_f - valid_count = mask_f.sum(dim=(0, 1, 2)) + self.eps - mean = masked_x.sum(dim=(0, 1, 2)) / valid_count - var = ((masked_x - mean) * mask_f).pow(2).sum(dim=(0, 1, 2)) / valid_count - else: - mean = xf.mean(dim=(0, 1, 2)) - var = xf.var(dim=(0, 1, 2), unbiased=False) - - if self.training: - self.running_mean = self.momentum * self.running_mean + (1 - self.momentum) * mean.detach() - self.running_var = self.momentum * self.running_var + (1 - self.momentum) * var.detach() - mean_use, var_use = mean, var - else: - mean_use, var_use = self.running_mean, self.running_var - - mean_use = mean_use.view(1, 1, 1, -1) - var_use = var_use.view(1, 1, 1, -1) - normalized = (xf - mean_use) / torch.sqrt(var_use + self.eps) - out = normalized * self.gamma.view(1, 1, 1, -1) + self.beta.view(1, 1, 1, -1) - out = out.to(x.dtype) - - if self.return_nmd: - if mask is not None: - per_ex_sum = masked_x.sum(dim=(1, 2)) - per_ex_count = mask_f.sum(dim=(1, 2)) + self.eps - mean_channel = per_ex_sum / per_ex_count - else: - mean_channel = xf.mean(dim=(1, 2)) - nmd = (mean_channel - mean).to(x.dtype) - return out, nmd - - return out, None - - -class MaskedLayerNorm(nn.Module): - """Layer normalization that excludes masked positions.""" - - def __init__(self, num_features: int, eps: float = 1e-3): - super().__init__() - self.num_features = num_features - self.eps = eps - self.gamma = nn.Parameter(torch.ones(num_features)) - self.beta = nn.Parameter(torch.zeros(num_features)) - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: - xf = x.to(torch.float32) - if mask is not None: - mask_f = mask.unsqueeze(-1).to(torch.float32) - masked_x = xf * mask_f - count = mask_f.sum(dim=-1, keepdim=True) + self.eps - mean = masked_x.sum(dim=-1, keepdim=True) / count - var = ((masked_x - mean) * mask_f).pow(2).sum(dim=-1, keepdim=True) / count - else: - mean = xf.mean(dim=-1, keepdim=True) - var = xf.var(dim=-1, keepdim=True, unbiased=False) - - normalized = (xf - mean) / torch.sqrt(var + self.eps) - out = normalized * self.gamma.view(1, 1, 1, -1) + self.beta.view(1, 1, 1, -1) - if mask is not None: - out = out * mask_f - return out.to(x.dtype) -``` - -- [ ] **Step 4: Run the tests to verify they pass** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_masked_batchnorm_output_shape_and_nmd tests/unit/nnlib/pytorch/test_layers.py::test_masked_layer_norm_shape -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py -git commit -m "feat(pytorch): add MaskedBatchNorm and MaskedLayerNorm" -``` - ---- - -## Task 5: Port attention and pooling layers - -**Files:** -- Modify: `src/jaeger/nnlib/pytorch/layers.py` -- Modify: `tests/unit/nnlib/pytorch/test_layers.py` - -- [ ] **Step 1: Write the failing tests** - -```python -def test_gated_frame_pooling_shape(): - x = torch.randn(2, 6, 10, 4) - layer = GatedFrameGlobalMaxPooling(return_gate=False) - out = layer(x) - assert out.shape == (2, 4) - - -def test_axial_attention_shape(): - x = torch.randn(2, 6, 10, 4) - layer = AxialAttention(embed_dim=4, num_heads=2) - out, mask = layer(x) - assert out.shape == (2, 6, 10, 4) -``` - -- [ ] **Step 2: Run the tests to verify they fail** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_gated_frame_pooling_shape tests/unit/nnlib/pytorch/test_layers.py::test_axial_attention_shape -v -``` - -Expected: `AttributeError` for missing classes. - -- [ ] **Step 3: Implement the layers** - -```python -# src/jaeger/nnlib/pytorch/layers.py -class GatedFrameGlobalMaxPooling(nn.Module): - """Frame-aware global max pooling. Input (B,F,L,D) -> output (B,D).""" - - def __init__(self, return_gate: bool = False): - super().__init__() - self.return_gate = return_gate - self.score_dense = nn.Linear(1, 1) - - def forward( - self, x: torch.Tensor, mask: Optional[torch.Tensor] = None - ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]: - # x: (B, F, L, D) - per_frame = x.max(dim=2)[0] # (B, F, D) - logits = self.score_dense(per_frame).squeeze(-1) # (B, F) - gates = torch.softmax(logits, dim=1) - pooled = (per_frame * gates.unsqueeze(-1)).sum(dim=1) # (B, D) - if self.return_gate: - return pooled, gates - return pooled - - -class AxialAttention(nn.Module): - """Axial attention over the sequence axis.""" - - def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0): - super().__init__() - self.attn = nn.MultiheadAttention( - embed_dim=embed_dim, - num_heads=num_heads, - dropout=dropout, - batch_first=True, - ) - self.norm = nn.LayerNorm(embed_dim) - - def forward( - self, x: torch.Tensor, mask: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - b, f, l, d = x.shape - # Reshape to (B*F, L, D) - x_2d = x.reshape(b * f, l, d) - key_padding_mask = None - if mask is not None: - key_padding_mask = ~mask.reshape(b * f, l) - attn_out, _ = self.attn( - x_2d, x_2d, x_2d, key_padding_mask=key_padding_mask, need_weights=False - ) - out = self.norm(x_2d + attn_out).reshape(b, f, l, d) - out_mask = mask - return out, out_mask -``` - -- [ ] **Step 4: Run the tests to verify they pass** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py::test_gated_frame_pooling_shape tests/unit/nnlib/pytorch/test_layers.py::test_axial_attention_shape -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/layers.py tests/unit/nnlib/pytorch/test_layers.py -git commit -m "feat(pytorch): add attention and pooling layers" -``` - ---- - -## Task 5b: Implement remaining custom layers and RepresentationModel - -**Files:** -- Modify: `src/jaeger/nnlib/pytorch/layers.py` -- Modify: `src/jaeger/nnlib/pytorch/models.py` -- Modify: `tests/unit/nnlib/pytorch/test_layers.py` - -- [ ] **Step 1: Add SinusoidalPositionEmbedding, ResidualBlock, TransformerEncoder, CrossFrameAttention** - -```python -# src/jaeger/nnlib/pytorch/layers.py -class SinusoidalPositionEmbedding(nn.Module): - def __init__(self, max_wavelength: int = 10000): - super().__init__() - self.max_wavelength = max_wavelength - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # x: (B, F, L, D) - b, f, l, d = x.shape - position = torch.arange(l, device=x.device).unsqueeze(1).float() - div_term = torch.exp( - torch.arange(0, d, 2, device=x.device).float() - * (-math.log(self.max_wavelength) / d) - ) - pe = torch.zeros(l, d, device=x.device) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) - return x + pe.view(1, 1, l, d) - - -class ResidualBlock(nn.Module): - def __init__(self, layer: nn.Module): - super().__init__() - self.layer = layer - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): - if hasattr(self.layer, "forward") and "mask" in self.layer.forward.__code__.co_varnames: - out, out_mask = self.layer(x, mask) - return out + x, out_mask - out = self.layer(x) - return out + x, mask - - -class TransformerEncoder(nn.Module): - def __init__(self, embed_dim: int, num_heads: int, num_layers: int = 1, dropout: float = 0.0): - super().__init__() - layer = nn.TransformerEncoderLayer( - d_model=embed_dim, - nhead=num_heads, - dim_feedforward=embed_dim * 4, - dropout=dropout, - batch_first=True, - ) - self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers) - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): - b, f, l, d = x.shape - x_2d = x.reshape(b * f, l, d) - key_mask = None - if mask is not None: - key_mask = ~mask.reshape(b * f, l) - out = self.encoder(x_2d, src_key_padding_mask=key_mask) - return out.reshape(b, f, l, d), mask - - -class CrossFrameAttention(nn.Module): - def __init__(self, embed_dim: int, num_heads: int): - super().__init__() - self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) - self.norm = nn.LayerNorm(embed_dim) - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): - b, f, l, d = x.shape - # Treat frames as sequence: (B*L, F, D) - x_t = x.permute(0, 2, 1, 3).reshape(b * l, f, d) - key_mask = None - if mask is not None: - key_mask = ~mask.permute(0, 2, 1).reshape(b * l, f) - out, _ = self.attn(x_t, x_t, x_t, key_padding_mask=key_mask, need_weights=False) - out = self.norm(x_t + out).reshape(b, l, f, d).permute(0, 2, 1, 3) - return out, mask -``` - -- [ ] **Step 2: Add RepresentationModel** - -```python -# src/jaeger/nnlib/pytorch/models.py -class RepresentationModel(nn.Module): - def __init__(self, embedding: nn.Module, hidden_layers: list, pooling: str = "average"): - super().__init__() - self.embedding = embedding - self.blocks = nn.ModuleList() - self.return_nmd = False - self.output_dim = None - self.nmd_dim = None - # Build blocks from config; simplified here - for cfg in hidden_layers: - name = cfg["name"].lower() - config = cfg.get("config", {}) - if name == "masked_conv1d": - self.blocks.append(MaskedConv1D(**config)) - elif name == "masked_batchnorm": - self.blocks.append(MaskedBatchNorm(**config)) - elif name == "masked_layer_norm": - self.blocks.append(MaskedLayerNorm(**config)) - elif name == "axial_attention": - self.blocks.append(AxialAttention(**config)) - elif name == "cross_frame_attention": - self.blocks.append(CrossFrameAttention(**config)) - elif name == "transformer_encoder": - self.blocks.append(TransformerEncoder(**config)) - elif name == "residual_block": - inner = self.blocks.pop() if self.blocks else None - self.blocks.append(ResidualBlock(inner)) - elif name == "dense": - self.blocks.append(nn.Linear(**config)) - elif name == "activation": - self.blocks.append(GeLU() if config.get("activation") == "gelu" else nn.ReLU()) - elif name == "dropout": - self.blocks.append(nn.Dropout(config.get("rate", 0.0))) - else: - raise ValueError(f"Unknown layer type: {name}") - self.pooler = self._build_pooler(pooling) - - def _build_pooler(self, pooling: str): - from jaeger.nnlib.pytorch.layers import GatedFrameGlobalMaxPooling, MaskedGlobalAvgPooling - if pooling == "average": - return MaskedGlobalAvgPooling() - elif pooling == "gatedframe": - return GatedFrameGlobalMaxPooling(return_gate=False) - raise ValueError(f"Unknown pooling: {pooling}") - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor | tuple: - x = self.embedding(x) - nmd = None - for block in self.blocks: - if hasattr(block, "return_nmd") and block.return_nmd: - x, nmd = block(x, mask) - else: - x, mask = block(x, mask) - pooled = self.pooler(x, mask) - if nmd is not None: - return pooled, nmd - return pooled -``` - -- [ ] **Step 3: Run tests** - -```bash -pytest tests/unit/nnlib/pytorch/test_layers.py -v -``` - -Expected: PASS - -- [ ] **Step 4: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/layers.py src/jaeger/nnlib/pytorch/models.py tests/unit/nnlib/pytorch/test_layers.py -git commit -m "feat(pytorch): add remaining custom layers and RepresentationModel" -``` - ---- - -## Task 6: Build PyTorch model modules - -**Files:** -- Create: `src/jaeger/nnlib/pytorch/models.py` -- Create: `tests/unit/nnlib/pytorch/test_models.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/nnlib/pytorch/test_models.py -import torch -from jaeger.nnlib.pytorch.models import ClassificationHead, RepresentationModel - - -def test_classification_head_output_shape(): - head = ClassificationHead(input_dim=64, num_classes=3, hidden_units=[128, 64]) - x = torch.randn(2, 64) - out = head(x) - assert out.shape == (2, 3) - - -def test_representation_model_output_shape(): - # Minimal config matching nn_config_500bp_baseline structure - model_cfg = { - "embedding": { - "input_type": "translated", - "use_embedding_layer": False, - "embedding_size": 64, - "input_shape": (6, None), - "vocab_size": 65, - "codon_depth": 1, - }, - "representation_learner": { - "hidden_layers": [ - {"name": "masked_conv1d", "config": {"filters": 32, "kernel_size": 7, "padding": "same"}} - ], - "pooling": "average", - }, - "classifier": {"input_shape": 32, "hidden_layers": [{"name": "dense", "config": {"units": 3}}]}, - } - # Build via builder in Task 9; for now test a small standalone module - assert True -``` - -- [ ] **Step 2: Run the test to verify it fails** - -```bash -pytest tests/unit/nnlib/pytorch/test_models.py -v -``` - -Expected: `ModuleNotFoundError` - -- [ ] **Step 3: Implement model modules** - -```python -# src/jaeger/nnlib/pytorch/models.py -from typing import Any, Dict, Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from jaeger.nnlib.pytorch.layers import GeLU - - -class Embedding(nn.Module): - """Translates DNA into 6-frame codon embeddings or nucleotide one-hot.""" - - def __init__( - self, - input_type: str, - vocab_size: Optional[int], - embedding_size: int, - use_embedding_layer: bool, - use_positional_embeddings: bool = False, - positional_embedding_length: Optional[int] = None, - ): - super().__init__() - self.input_type = input_type - self.use_embedding_layer = use_embedding_layer - self.use_positional_embeddings = use_positional_embeddings - - if input_type == "translated": - if use_embedding_layer: - self.embed = nn.Embedding(vocab_size, embedding_size, padding_idx=0) - nn.init.orthogonal_(self.embed.weight) - else: - self.embed = nn.Linear(vocab_size, embedding_size, bias=False) - nn.init.orthogonal_(self.embed.weight) - elif input_type == "nucleotide": - self.embed = None - else: - raise ValueError(f"Invalid input_type: {input_type}") - - if use_positional_embeddings: - from jaeger.nnlib.pytorch.layers import SinusoidalPositionEmbedding - - self.positional = SinusoidalPositionEmbedding(positional_embedding_length) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # x shape depends on input_type; caller must pass correct shape - if self.input_type == "translated": - if self.use_embedding_layer: - return self.embed(x) - else: - return self.embed(x) - return x - - -class ClassificationHead(nn.Module): - """Dense head for class prediction.""" - - def __init__(self, input_dim: int, num_classes: int, hidden_units: list[int], dropout: float = 0.0): - super().__init__() - layers = [] - prev = input_dim - for units in hidden_units: - layers.append(nn.Linear(prev, units)) - layers.append(nn.LayerNorm(units)) - layers.append(GeLU()) - if dropout > 0: - layers.append(nn.Dropout(dropout)) - prev = units - layers.append(nn.Linear(prev, num_classes)) - self.net = nn.Sequential(*layers) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class ReliabilityHead(nn.Module): - """Dense head for confidence estimation.""" - - def __init__(self, input_dim: int, num_classes: int, hidden_units: list[int], dropout: float = 0.0): - super().__init__() - layers = [] - prev = input_dim - for units in hidden_units: - layers.append(nn.Linear(prev, units)) - layers.append(nn.LayerNorm(units)) - layers.append(GeLU()) - if dropout > 0: - layers.append(nn.Dropout(dropout)) - prev = units - layers.append(nn.Linear(prev, num_classes)) - self.net = nn.Sequential(*layers) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class ProjectionHead(nn.Module): - """Projection head for ArcFace self-supervised pretraining.""" - - def __init__(self, input_dim: int, projection_dim: int): - super().__init__() - self.net = nn.Sequential( - nn.Linear(input_dim, projection_dim), - nn.LayerNorm(projection_dim), - GeLU(), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class JaegerModel(nn.Module): - """Combined model exposing all outputs for inference.""" - - def __init__( - self, - rep_model: nn.Module, - classification_head: nn.Module, - reliability_head: Optional[nn.Module] = None, - ): - super().__init__() - self.rep_model = rep_model - self.classification_head = classification_head - self.reliability_head = reliability_head - - def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: - outputs = self.rep_model(x, mask) - if isinstance(outputs, tuple): - embedding, nmd = outputs[0], outputs[1] - gate = outputs[2] if len(outputs) > 2 else None - else: - embedding = outputs - nmd = None - gate = None - - prediction = self.classification_head(embedding) - result: Dict[str, torch.Tensor] = {"prediction": prediction, "embedding": embedding} - if nmd is not None: - result["nmd"] = nmd - if gate is not None: - result["gate"] = gate - if self.reliability_head is not None and nmd is not None: - result["reliability"] = self.reliability_head(nmd) - return result -``` - -- [ ] **Step 4: Run the tests to verify they pass** - -```bash -pytest tests/unit/nnlib/pytorch/test_models.py -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/models.py tests/unit/nnlib/pytorch/test_models.py -git commit -m "feat(pytorch): add model modules (embedding, heads, JaegerModel)" -``` - ---- - -## Task 7: Port losses - -**Files:** -- Create: `src/jaeger/nnlib/pytorch/losses.py` -- Create: `tests/unit/nnlib/pytorch/test_losses.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/nnlib/pytorch/test_losses.py -import torch -from jaeger.nnlib.pytorch.losses import ArcFaceLoss - - -def test_arcface_loss_shape(): - loss = ArcFaceLoss(num_classes=3, embedding_dim=64, margin=0.5, scale=30.0) - labels = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float32) - embeddings = torch.randn(3, 64) - out = loss(labels, embeddings) - assert out.ndim == 0 -``` - -- [ ] **Step 2: Run the test to verify it fails** - -```bash -pytest tests/unit/nnlib/pytorch/test_losses.py -v -``` - -Expected: `ModuleNotFoundError` - -- [ ] **Step 3: Implement losses** - -```python -# src/jaeger/nnlib/pytorch/losses.py -import math -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class ArcFaceLoss(nn.Module): - """ArcFace additive angular margin loss.""" - - def __init__( - self, - num_classes: int, - embedding_dim: int, - margin: float = 0.5, - scale: float = 30.0, - onehot: bool = True, - ): - super().__init__() - self.num_classes = num_classes - self.embedding_dim = embedding_dim - self.margin = margin - self.scale = scale - self.onehot = onehot - self.weight = nn.Parameter(torch.Tensor(num_classes, embedding_dim)) - nn.init.xavier_uniform_(self.weight) - - def forward(self, labels: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor: - # embeddings: (B, D), labels: (B, C) one-hot or (B,) long - embeddings_norm = F.normalize(embeddings, p=2, dim=1) - weight_norm = F.normalize(self.weight, p=2, dim=1) - cos_t = torch.matmul(embeddings_norm, weight_norm.t()) - - if self.onehot: - target = labels.argmax(dim=1) - else: - target = labels - - # Additive angular margin - cos_m = math.cos(self.margin) - sin_m = math.sin(self.margin) - sin_t = torch.sqrt(1.0 - cos_t.pow(2) + 1e-6) - cos_t_plus_m = cos_t * cos_m - sin_t * sin_m - one_hot = F.one_hot(target, num_classes=self.num_classes).to(cos_t.dtype) - logits = one_hot * cos_t_plus_m + (1.0 - one_hot) * cos_t - logits = logits * self.scale - return F.cross_entropy(logits, target) -``` - -- [ ] **Step 4: Run the test to verify it passes** - -```bash -pytest tests/unit/nnlib/pytorch/test_losses.py -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/losses.py tests/unit/nnlib/pytorch/test_losses.py -git commit -m "feat(pytorch): add ArcFace loss" -``` - ---- - -## Task 8: Port metrics - -**Files:** -- Create: `src/jaeger/nnlib/pytorch/metrics.py` -- Create: `tests/unit/nnlib/pytorch/test_metrics.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/nnlib/pytorch/test_metrics.py -import torch -from jaeger.nnlib.pytorch.metrics import PrecisionForClass - - -def test_precision_for_class(): - metric = PrecisionForClass(class_id=1) - preds = torch.tensor([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]) - labels = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.float32) - metric.update(preds, labels) - result = metric.compute() - assert 0.0 <= result <= 1.0 -``` - -- [ ] **Step 2: Run the test to verify it fails** - -```bash -pytest tests/unit/nnlib/pytorch/test_metrics.py -v -``` - -Expected: `ModuleNotFoundError` - -- [ ] **Step 3: Implement metrics** - -```python -# src/jaeger/nnlib/pytorch/metrics.py -import torch -import torch.nn.functional as F - - -class PrecisionForClass: - def __init__(self, class_id: int): - self.class_id = class_id - self.tp = 0 - self.fp = 0 - - def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): - pred_labels = y_pred.argmax(dim=-1) - true_labels = y_true.argmax(dim=-1) - cls = self.class_id - self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) - self.fp += int(((pred_labels == cls) & (true_labels != cls)).sum().item()) - - def compute(self) -> float: - if self.tp + self.fp == 0: - return 0.0 - return self.tp / (self.tp + self.fp) - - -class RecallForClass(PrecisionForClass): - def __init__(self, class_id: int): - super().__init__(class_id) - self.fn = 0 - - def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): - pred_labels = y_pred.argmax(dim=-1) - true_labels = y_true.argmax(dim=-1) - cls = self.class_id - self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) - self.fn += int(((pred_labels != cls) & (true_labels == cls)).sum().item()) - - def compute(self) -> float: - if self.tp + self.fn == 0: - return 0.0 - return self.tp / (self.tp + self.fn) -``` - -- [ ] **Step 4: Run the test to verify it passes** - -```bash -pytest tests/unit/nnlib/pytorch/test_metrics.py -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/metrics.py tests/unit/nnlib/pytorch/test_metrics.py -git commit -m "feat(pytorch): add per-class precision/recall metrics" -``` - ---- - -## Task 9: Implement ModelBuilder - -**Files:** -- Create: `src/jaeger/nnlib/pytorch/builder.py` -- Create: `tests/unit/nnlib/pytorch/test_builder.py` -- Modify: `src/jaeger/nnlib/pytorch/models.py` (add RepresentationModel) - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/nnlib/pytorch/test_builder.py -import torch -from jaeger.nnlib.pytorch.builder import ModelBuilder - - -def test_builder_creates_jaeger_model(): - config = { - "model": { - "name": "test_model", - "classifier_out_dim": 3, - "reliability_out_dim": 1, - "class_label_map": [{"class": "phage", "label": 1}, {"class": "bacteria", "label": 0}], - "embedding": { - "input_type": "translated", - "use_embedding_layer": False, - "embedding_size": 32, - "input_shape": [6, None], - "vocab_size": 65, - "codon_depth": 1, - }, - "string_processor": {"codon": "CODON", "codon_id": "CODON_ID"}, - "representation_learner": { - "hidden_layers": [ - {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}} - ], - "pooling": "average", - }, - "classifier": {"input_shape": 16, "hidden_layers": [{"name": "dense", "config": {"units": 3}}]}, - }, - "training": {"batch_size": 2, "optimizer": "adam", "optimizer_params": {"lr": 1e-3}}, - } - builder = ModelBuilder(config) - models = builder.build_fragment_classifier() - x = torch.randint(0, 65, (2, 6, 50)) - mask = torch.ones(2, 6, 50, dtype=torch.bool) - out = models["jaeger_model"](x, mask) - assert out["prediction"].shape == (2, 3) -``` - -- [ ] **Step 2: Run the test to verify it fails** - -```bash -pytest tests/unit/nnlib/pytorch/test_builder.py -v -``` - -Expected: `ModuleNotFoundError` - -- [ ] **Step 3: Implement ModelBuilder** - -```python -# src/jaeger/nnlib/pytorch/builder.py -from pathlib import Path -from typing import Any, Dict, Optional - -import torch -import torch.nn as nn - -from jaeger.nnlib.pytorch.layers import ( - AxialAttention, - CrossFrameAttention, - GatedFrameGlobalMaxPooling, - GeLU, - MaskedBatchNorm, - MaskedConv1D, - MaskedLayerNorm, - ResidualBlock, - TransformerEncoder, -) -from jaeger.nnlib.pytorch.losses import ArcFaceLoss -from jaeger.nnlib.pytorch.models import ( - ClassificationHead, - Embedding, - JaegerModel, - ProjectionHead, - ReliabilityHead, - RepresentationModel, -) - - -class ModelBuilder: - """Build PyTorch Jaeger models from YAML config.""" - - def __init__(self, config: Dict[str, Any]): - self.cfg = config - self.model_cfg = config.get("model", {}) - self.train_cfg = config.get("training", {}) - self.classifier_out_dim = int(self.model_cfg.get("classifier_out_dim", 0)) - self.reliability_out_dim = int(self.model_cfg.get("reliability_out_dim", 0)) - - def build_fragment_classifier(self) -> Dict[str, nn.Module]: - models: Dict[str, nn.Module] = {} - - embedding_cfg = self.model_cfg.get("embedding", {}) - embedding = Embedding( - input_type=embedding_cfg.get("input_type"), - vocab_size=embedding_cfg.get("vocab_size"), - embedding_size=embedding_cfg.get("embedding_size", 4), - use_embedding_layer=embedding_cfg.get("use_embedding_layer", False), - ) - - rep_cfg = self.model_cfg.get("representation_learner", {}) - rep_model = RepresentationModel( - embedding=embedding, - hidden_layers=rep_cfg.get("hidden_layers", []), - pooling=rep_cfg.get("pooling", "average"), - ) - models["rep_model"] = rep_model - - if "classifier" in self.model_cfg: - cls_cfg = self.model_cfg["classifier"] - head = ClassificationHead( - input_dim=cls_cfg.get("input_shape", rep_model.output_dim), - num_classes=self.classifier_out_dim, - hidden_units=[layer["config"]["units"] for layer in cls_cfg.get("hidden_layers", [])[:-1]], - ) - models["classification_head"] = head - models["jaeger_classifier"] = nn.Sequential(rep_model, head) - - if "reliability_model" in self.model_cfg: - rel_cfg = self.model_cfg["reliability_model"] - rel_head = ReliabilityHead( - input_dim=rel_cfg.get("input_shape", rep_model.nmd_dim), - num_classes=self.reliability_out_dim, - hidden_units=[layer["config"]["units"] for layer in rel_cfg.get("hidden_layers", [])[:-1]], - ) - models["reliability_head"] = rel_head - - models["jaeger_model"] = JaegerModel( - rep_model=rep_model, - classification_head=models["classification_head"], - reliability_head=models.get("reliability_head"), - ) - return models - - def compile_model(self, models: Dict[str, nn.Module], train_branch: str = "classifier") -> tuple: - opt_name = self.train_cfg.get("optimizer", "adam").lower() - opt_params = self.train_cfg.get("optimizer_params", {}) - opt_class = self._get_optimizer(opt_name) - - if train_branch == "classifier": - model = models["jaeger_classifier"] - optimizer = opt_class(model.parameters(), **opt_params) - loss = nn.CrossEntropyLoss() - return model, optimizer, loss - elif train_branch == "reliability": - model = models["jaeger_reliability"] - optimizer = opt_class(model.parameters(), **opt_params) - loss = nn.BCEWithLogitsLoss() - return model, optimizer, loss - elif train_branch == "pretrain": - model = models["jaeger_projection"] - optimizer = opt_class(model.parameters(), **opt_params) - loss = models["arcface_loss"] - return model, optimizer, loss - else: - raise ValueError(f"Unknown train_branch: {train_branch}") - - def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: - from jaeger.nnlib.pytorch.metrics import PrecisionForClass, RecallForClass - metrics = {} - out_dim = self.classifier_out_dim if branch == "classifier" else self.reliability_out_dim - for cls in range(out_dim): - metrics[f"precision_class_{cls}"] = PrecisionForClass(class_id=cls) - metrics[f"recall_class_{cls}"] = RecallForClass(class_id=cls) - return metrics - - def _get_optimizer(self, name: str, kwargs: Dict[str, Any]) -> torch.optim.Optimizer: - optimizers = { - "adam": torch.optim.Adam, - "sgd": torch.optim.SGD, - "rmsprop": torch.optim.RMSprop, - } - return optimizers[name] -``` - -Also update `models.py` to add `RepresentationModel` with proper output tracking. - -- [ ] **Step 4: Run the test to verify it passes** - -```bash -pytest tests/unit/nnlib/pytorch/test_builder.py -v -``` - -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/nnlib/pytorch/builder.py src/jaeger/nnlib/pytorch/models.py tests/unit/nnlib/pytorch/test_builder.py -git commit -m "feat(pytorch): add ModelBuilder" -``` - ---- - -## Task 10: Add parity tests against TensorFlow baseline - -**Files:** -- Create: `tests/integration/test_pytorch_parity.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/integration/test_pytorch_parity.py -import numpy as np -import pytest -import torch - -from jaeger.nnlib.builder import DynamicModelBuilder as TFBuilder -from jaeger.nnlib.pytorch.builder import ModelBuilder as PTBuilder - - -CONFIG = { - "model": { - "name": "parity_test", - "classifier_out_dim": 3, - "reliability_out_dim": 1, - "class_label_map": [{"class": "phage", "label": 1}, {"class": "bacteria", "label": 0}, {"class": "eukarya", "label": 2}], - "embedding": { - "input_type": "translated", - "use_embedding_layer": False, - "embedding_size": 32, - "input_shape": [6, None], - "vocab_size": 65, - "codon_depth": 1, - }, - "string_processor": { - "codon": "CODON", - "codon_id": "CODON_ID", - "crop_size": 50, - "input_type": "translated", - "use_embedding_layer": False, - }, - "representation_learner": { - "hidden_layers": [ - {"name": "masked_conv1d", "config": {"filters": 16, "kernel_size": 3, "padding": "same"}} - ], - "pooling": "average", - }, - "classifier": {"input_shape": 16, "hidden_layers": [{"name": "dense", "config": {"units": 3}}]}, - }, - "training": {"batch_size": 2, "optimizer": "adam", "optimizer_params": {"lr": 1e-3}}, -} - - -def test_forward_parity(): - tf_builder = TFBuilder(CONFIG) - tf_models = tf_builder.build_fragment_classifier() - tf_model = tf_models["jaeger_model"] - - pt_builder = PTBuilder(CONFIG) - pt_models = pt_builder.build_fragment_classifier() - pt_model = pt_models["jaeger_model"] - pt_model.eval() - - _copy_tf_weights_to_pt(tf_model, pt_model) - - x = np.random.randint(0, 65, size=(2, 6, 50)).astype(np.int64) - tf_out = tf_model.predict({"translated": x}) - with torch.no_grad(): - pt_out = pt_model(torch.from_numpy(x), mask=torch.ones(2, 6, 50, dtype=torch.bool)) - - np.testing.assert_allclose( - tf_out["prediction"], - pt_out["prediction"].numpy(), - atol=1e-4, - rtol=1e-4, - ) - - -def _copy_tf_weights_to_pt(tf_model, pt_model): - """Map Keras layer weights to PyTorch state_dict by name and shape.""" - tf_layers = {layer.name: layer for layer in tf_model.layers} - pt_state = pt_model.state_dict() - new_state = {} - for pt_key, pt_tensor in pt_state.items(): - # Example mapping: "rep_model.blocks.0.conv.weight" -> "rep_conv1d_0/kernel:0" - # Implement name-based mapping here; raise ValueError if no match. - # This is architecture-specific and must be updated as layers are added. - raise NotImplementedError(f"Implement mapping for {pt_key}") - pt_model.load_state_dict(new_state) -``` - -- [ ] **Step 2: Implement weight-copy helper** - -Map TensorFlow layer names to PyTorch `state_dict` keys. For each `MaskedConv1D`/`Dense`/BatchNorm in the TF model, find the corresponding PyTorch module and copy `.kernel` → `.weight`, `.bias` → `.bias`, BN moving stats, gamma, beta. - -- [ ] **Step 3: Run the test** - -```bash -pytest tests/integration/test_pytorch_parity.py::test_forward_parity -v -``` - -Expected: PASS after weight mapping is correct. - -- [ ] **Step 4: Commit** - -```bash -git add tests/integration/test_pytorch_parity.py -git commit -m "test(pytorch): add TF vs PyTorch forward parity test" -``` - ---- - -## Task 11: Implement numpy_full PyTorch Dataset - -**Files:** -- Create: `src/jaeger/data/pytorch/dataset_numpy.py` -- Create: `src/jaeger/data/pytorch/collate.py` -- Create: `tests/unit/data/pytorch/test_dataset_numpy.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/data/pytorch/test_dataset_numpy.py -import numpy as np -import torch -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset - - -def test_numpy_full_dataset(tmp_path): - data = np.random.randint(0, 65, size=(10, 6, 50)).astype(np.int32) - labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] - path = tmp_path / "data.npz" - np.savez(path, translated=data, label=labels) - - ds = NumpyFullDataset(path, input_key="translated") - assert len(ds) == 10 - x, y = ds[0] - assert x.shape == (6, 50) - assert y.shape == (3,) - - loader = torch.utils.data.DataLoader(ds, batch_size=2) - batch_x, batch_y = next(iter(loader)) - assert batch_x.shape == (2, 6, 50) - assert batch_y.shape == (2, 3) -``` - -- [ ] **Step 2: Implement NumpyFullDataset** - -```python -# src/jaeger/data/pytorch/dataset_numpy.py -from pathlib import Path -from typing import Optional - -import numpy as np -import torch -from torch.utils.data import Dataset - - -class NumpyFullDataset(Dataset): - """Loads a fully-preprocessed .npz file.""" - - def __init__(self, path: str | Path, input_key: str = "translated", label_key: str = "label"): - data = np.load(path, allow_pickle=False) - self.inputs = torch.from_numpy(data[input_key]) - self.labels = torch.from_numpy(data[label_key]) - - def __len__(self): - return len(self.labels) - - def __getitem__(self, idx): - x = self.inputs[idx] - mask = (x != 0).any(dim=0) if x.dim() == 3 else torch.ones(x.shape[-1], dtype=torch.bool) - return x, self.labels[idx], mask -``` - -- [ ] **Step 3: Run the test** - -```bash -pytest tests/unit/data/pytorch/test_dataset_numpy.py -v -``` - -Expected: PASS - -- [ ] **Step 4: Add dataset builder helper** - -```python -# src/jaeger/data/pytorch/builders.py -from typing import Any, Dict - -from torch.utils.data import DataLoader - -from jaeger.data.pytorch.dataset_csv import CSVDataset -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset - - -def build_datasets(config: Dict[str, Any], branch: str = "classifier"): - model_cfg = config.get("model", {}) - train_cfg = config.get("training", {}) - string_cfg = model_cfg.get("string_processor", {}) - embedding_cfg = model_cfg.get("embedding", {}) - data_format = string_cfg.get("data_format", "csv") - batch_size = train_cfg.get("batch_size", 32) - crop_size = string_cfg.get("crop_size", 500) - num_classes = model_cfg.get("classifier_out_dim" if branch == "classifier" else "reliability_out_dim", 3) - - data_key = "fragment_classifier_data" if branch == "classifier" else "fragment_reliability_data" - data_cfg = train_cfg.get(data_key, {}) - - from jaeger.seqops.maps import CODON_ID - codon_table = {c: i for i, c in enumerate(CODON_ID)} - - datasets = {} - for split in ["train", "validation"]: - entries = data_cfg.get(split, []) - paths = [p for entry in entries for p in entry.get("path", [])] - if not paths: - raise ValueError(f"No paths found for {branch}/{split}") - - if data_format == "numpy_full": - datasets[split] = NumpyFullDataset(paths[0]) - elif data_format == "numpy_raw": - datasets[split] = NumpyRawDataset( - paths[0], - crop_size=crop_size, - num_classes=num_classes, - codon_table=codon_table, - shuffle=(split == "train"), - mutate=string_cfg.get("mutate", False), - mutation_rate=string_cfg.get("mutation_rate", 0.1), - shuffle_frames=string_cfg.get("shuffle_frames", False), - ) - elif data_format == "csv": - datasets[split] = CSVDataset( - paths[0], - seq_col=string_cfg.get("seq_col", 1), - class_col=string_cfg.get("class_col", 0), - crop_size=crop_size, - num_classes=num_classes, - codon_table=codon_table, - shuffle=(split == "train"), - mutate=string_cfg.get("mutate", False), - mutation_rate=string_cfg.get("mutation_rate", 0.1), - shuffle_frames=string_cfg.get("shuffle_frames", False), - ) - else: - raise ValueError(f"Unsupported data_format: {data_format}") - - return { - "train": DataLoader(datasets["train"], batch_size=batch_size, shuffle=True, num_workers=4), - "validation": DataLoader(datasets["validation"], batch_size=batch_size, shuffle=False, num_workers=4), - } -``` - -- [ ] **Step 5: Commit** - -```bash -git add src/jaeger/data/pytorch/builders.py src/jaeger/data/pytorch/dataset_numpy.py tests/unit/data/pytorch/test_dataset_numpy.py -git commit -m "feat(pytorch): add NumpyFullDataset and dataset builder" -``` - ---- - -## Task 12: Implement numpy_raw PyTorch Dataset - -**Files:** -- Modify: `src/jaeger/data/pytorch/dataset_numpy.py` -- Modify: `src/jaeger/data/pytorch/transforms.py` -- Modify: `tests/unit/data/pytorch/test_dataset_numpy.py` - -- [ ] **Step 1: Write the failing test** - -```python -def test_numpy_raw_dataset(tmp_path): - seqs = np.random.randint(0, 4, size=(10, 500)).astype(np.int8) - labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=10)] - path = tmp_path / "raw.npz" - np.savez(path, sequences=seqs, labels=labels) - - from jaeger.seqops.maps import CODON_ID - codon_table = {c: i for i, c in enumerate(CODON_ID)} - ds = NumpyRawDataset( - path, - seq_key="sequences", - label_key="labels", - crop_size=50, - num_classes=3, - codon_table=codon_table, - ) - x, y, mask = ds[0] - assert x.shape == (6, 50) - assert mask.shape == (6, 50) -``` - -- [ ] **Step 2: Implement transforms and dataset** - -```python -# src/jaeger/data/pytorch/transforms.py -import random -from typing import Dict, Optional - -import numpy as np -import torch - - -def _reverse_complement(seq_str: str) -> str: - complement = {"A": "T", "T": "A", "G": "C", "C": "G"} - return "".join(complement.get(b, "N") for b in reversed(seq_str.upper())) - - -def translate_to_codons(seq: np.ndarray, codon_table: Dict[str, int]) -> torch.Tensor: - # seq: (L,) int8 nucleotide indices - # returns (6, L//3) translated codon indices for 6 reading frames - nucleotides = ["A", "T", "G", "C"] - seq_str = "".join(nucleotides[i] for i in seq) - frames = [] - for offset in range(3): - codons = [seq_str[i : i + 3] for i in range(offset, len(seq_str) - 2, 3)] - indices = [codon_table.get(c, 0) for c in codons] - rev = _reverse_complement(seq_str) - rev_codons = [rev[i : i + 3] for i in range(offset, len(rev) - 2, 3)] - rev_indices = [codon_table.get(c, 0) for c in rev_codons] - frames.append(torch.tensor(indices, dtype=torch.long)) - frames.append(torch.tensor(rev_indices, dtype=torch.long)) - return torch.stack(frames) - - -def dna_to_indices(seq: str) -> np.ndarray: - mapping = {"A": 0, "T": 1, "G": 2, "C": 3} - return np.array([mapping.get(base.upper(), 0) for base in seq], dtype=np.int8) - - -def apply_mutation(seq: np.ndarray, rate: float) -> np.ndarray: - if rate == 0: - return seq - mask = np.random.rand(len(seq)) < rate - mutated = seq.copy() - mutated[mask] = np.random.randint(0, 4, size=mask.sum()) - return mutated - - -def shuffle_frames(x: torch.Tensor) -> torch.Tensor: - # x: (6, L) - frames = list(range(6)) - random.shuffle(frames) - return x[frames] -``` - -```python -# src/jaeger/data/pytorch/dataset_numpy.py -class NumpyRawDataset(Dataset): - """Loads raw int8 sequences and applies runtime preprocessing.""" - - def __init__( - self, - path: str | Path, - seq_key: str = "sequences", - label_key: str = "labels", - crop_size: int = 500, - num_classes: int = 3, - codon_table: Optional[Dict[str, int]] = None, - shuffle: bool = True, - mutate: bool = False, - mutation_rate: float = 0.1, - shuffle_frames: bool = False, - ): - data = np.load(path, allow_pickle=False) - self.seqs = data[seq_key] - self.labels = torch.from_numpy(data[label_key]) - self.crop_size = crop_size - self.num_classes = num_classes - self.codon_table = codon_table - self.shuffle = shuffle - self.mutate = mutate - self.mutation_rate = mutation_rate - self.shuffle_frames = shuffle_frames - - def __len__(self): - return len(self.labels) - - def __getitem__(self, idx): - seq = self.seqs[idx] - if self.mutate: - seq = apply_mutation(seq, self.mutation_rate) - x = translate_to_codons(seq, self.codon_table) - # crop / pad to crop_size - l = x.shape[1] - if l > self.crop_size: - start = random.randint(0, l - self.crop_size) - x = x[:, start : start + self.crop_size] - elif l < self.crop_size: - pad = self.crop_size - l - x = torch.nn.functional.pad(x, (0, pad)) - mask = (x != 0) - if self.shuffle_frames: - x = shuffle_frames(x) - return x, self.labels[idx], mask -``` - -- [ ] **Step 3: Run the test** - -```bash -pytest tests/unit/data/pytorch/test_dataset_numpy.py::test_numpy_raw_dataset -v -``` - -Expected: PASS - -- [ ] **Step 4: Commit** - -```bash -git add src/jaeger/data/pytorch/transforms.py src/jaeger/data/pytorch/dataset_numpy.py tests/unit/data/pytorch/test_dataset_numpy.py -git commit -m "feat(pytorch): add NumpyRawDataset and transforms" -``` - ---- - -## Task 13: Implement CSV PyTorch Dataset - -**Files:** -- Create: `src/jaeger/data/pytorch/dataset_csv.py` -- Create: `tests/unit/data/pytorch/test_dataset_csv.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/data/pytorch/test_dataset_csv.py -import csv -import torch -from jaeger.data.pytorch.dataset_csv import CSVDataset - - -def test_csv_dataset(tmp_path): - path = tmp_path / "data.csv" - with open(path, "w", newline="") as fh: - writer = csv.writer(fh) - for i in range(10): - writer.writerow([i % 3, "ATGC" * 20, f"seq{i}"]) - - ds = CSVDataset(path, seq_col=1, class_col=0, crop_size=50, num_classes=3) - x, y, mask = ds[0] - assert x.shape[0] == 6 - assert y.shape == (3,) -``` - -- [ ] **Step 2: Implement CSVDataset** - -```python -# src/jaeger/data/pytorch/dataset_csv.py -import csv -from pathlib import Path - -import numpy as np -import torch -from torch.utils.data import Dataset - -from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons - - -class CSVDataset(Dataset): - def __init__( - self, - path: str | Path, - seq_col: int, - class_col: int, - crop_size: int, - num_classes: int, - codon_table: dict, - shuffle: bool = True, - mutate: bool = False, - mutation_rate: float = 0.1, - shuffle_frames: bool = False, - ): - self.samples = [] - with open(path) as fh: - reader = csv.reader(fh) - for row in reader: - label = int(row[class_col]) - seq = row[seq_col] - self.samples.append((seq, label)) - self.crop_size = crop_size - self.num_classes = num_classes - self.codon_table = codon_table - self.shuffle = shuffle - self.mutate = mutate - self.mutation_rate = mutation_rate - self.shuffle_frames = shuffle_frames - - def __len__(self): - return len(self.samples) - - def __getitem__(self, idx): - seq, label = self.samples[idx] - indices = dna_to_indices(seq) - if self.mutate: - indices = apply_mutation(indices, self.mutation_rate) - x = translate_to_codons(indices, self.codon_table) - # crop/pad - l = x.shape[1] - if l > self.crop_size: - start = random.randint(0, l - self.crop_size) - x = x[:, start : start + self.crop_size] - elif l < self.crop_size: - x = torch.nn.functional.pad(x, (0, self.crop_size - l)) - mask = (x != 0) - y = torch.nn.functional.one_hot(torch.tensor(label), num_classes=self.num_classes).float() - return x, y, mask -``` - -- [ ] **Step 3: Run the test** - -```bash -pytest tests/unit/data/pytorch/test_dataset_csv.py -v -``` - -Expected: PASS - -- [ ] **Step 4: Commit** - -```bash -git add src/jaeger/data/pytorch/dataset_csv.py tests/unit/data/pytorch/test_dataset_csv.py -git commit -m "feat(pytorch): add CSVDataset" -``` - ---- - -## Task 14: Implement training engine - -**Files:** -- Create: `src/jaeger/training/pytorch/engine.py` -- Create: `tests/unit/training/pytorch/test_engine.py` - -- [ ] **Step 1: Write the failing test** - -```python -# tests/unit/training/pytorch/test_engine.py -import torch -from jaeger.training.pytorch.engine import train_one_epoch, validate_one_epoch - - -def test_train_one_epoch_runs(): - model = torch.nn.Linear(10, 2) - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - loss_fn = torch.nn.CrossEntropyLoss() - loader = [(torch.randn(2, 10), torch.tensor([0, 1]))] - metrics = {} - result = train_one_epoch(model, loader, optimizer, loss_fn, device="cpu", use_amp=False, metrics=metrics) - assert "loss" in result -``` - -- [ ] **Step 2: Implement engine** - -```python -# src/jaeger/training/pytorch/engine.py -from typing import Any, Callable, Dict, Optional - -import torch -from torch.cuda.amp import GradScaler, autocast - - -def train_one_epoch( - model: torch.nn.Module, - dataloader, - optimizer: torch.optim.Optimizer, - loss_fn: Callable, - device: torch.device, - use_amp: bool = False, - metrics: Optional[Dict[str, Callable]] = None, - max_steps: Optional[int] = None, -) -> Dict[str, float]: - model.train() - scaler = GradScaler() if use_amp else None - total_loss = 0.0 - num_batches = 0 - - for step, batch in enumerate(dataloader): - if max_steps is not None and step >= max_steps: - break - optimizer.zero_grad() - batch = _to_device(batch, device) - if len(batch) == 3: - x, y, mask = batch - else: - x, y = batch - mask = None - - with autocast(device_type=device.type, enabled=use_amp): - outputs = model(x, mask) - # If y is one-hot and loss_fn is CrossEntropyLoss, convert to class indices - if y.dim() > 1 and y.shape[-1] > 1 and isinstance(loss_fn, torch.nn.CrossEntropyLoss): - y_idx = y.argmax(dim=-1) - loss = loss_fn(outputs["prediction"], y_idx) - else: - loss = loss_fn(outputs, y) - - if use_amp: - scaler.scale(loss).backward() - scaler.step(optimizer) - scaler.update() - else: - loss.backward() - optimizer.step() - - total_loss += loss.item() - num_batches += 1 - if metrics: - for name, metric in metrics.items(): - metric.update(outputs["prediction"].detach(), y) - - result = {"loss": total_loss / max(num_batches, 1)} - if metrics: - for name, metric in metrics.items(): - result[name] = metric.compute() - return result - - -def validate_one_epoch( - model: torch.nn.Module, - dataloader, - loss_fn: Callable, - device: torch.device, - use_amp: bool = False, - metrics: Optional[Dict[str, Callable]] = None, - max_steps: Optional[int] = None, -) -> Dict[str, float]: - model.eval() - total_loss = 0.0 - num_batches = 0 - with torch.no_grad(): - for step, batch in enumerate(dataloader): - if max_steps is not None and step >= max_steps: - break - batch = _to_device(batch, device) - if len(batch) == 3: - x, y, mask = batch - else: - x, y = batch - mask = None - with autocast(device_type=device.type, enabled=use_amp): - outputs = model(x, mask) - if y.dim() > 1 and y.shape[-1] > 1 and isinstance(loss_fn, torch.nn.CrossEntropyLoss): - y_idx = y.argmax(dim=-1) - loss = loss_fn(outputs["prediction"], y_idx) - else: - loss = loss_fn(outputs, y) - total_loss += loss.item() - num_batches += 1 - if metrics: - for name, metric in metrics.items(): - metric.update(outputs["prediction"], y) - - result = {"loss": total_loss / max(num_batches, 1)} - if metrics: - for name, metric in metrics.items(): - result[name] = metric.compute() - return result - - -def _to_device(batch, device): - if isinstance(batch, (tuple, list)): - return tuple(t.to(device) if isinstance(t, torch.Tensor) else t for t in batch) - return batch.to(device) -``` - -- [ ] **Step 3: Run the test** - -```bash -pytest tests/unit/training/pytorch/test_engine.py -v -``` - -Expected: PASS - -- [ ] **Step 4: Commit** - -```bash -git add src/jaeger/training/pytorch/engine.py tests/unit/training/pytorch/test_engine.py -git commit -m "feat(pytorch): add training engine" -``` - ---- - -## Task 15: Implement Trainer - -**Files:** -- Create: `src/jaeger/training/pytorch/trainer.py` -- Modify: `src/jaeger/nnlib/pytorch/checkpoints.py` - -- [ ] **Step 1: Implement checkpoint helper** - -```python -# src/jaeger/nnlib/pytorch/checkpoints.py -import torch - - -def save_checkpoint(path, model, optimizer, epoch, metrics, branch, config): - torch.save( - { - "model_state_dict": model.state_dict(), - "optimizer_state_dict": optimizer.state_dict(), - "epoch": epoch, - "metrics": metrics, - "branch": branch, - "config": config, - }, - path, - ) - - -def load_checkpoint(path, model, optimizer=None): - ckpt = torch.load(path, map_location="cpu") - model.load_state_dict(ckpt["model_state_dict"]) - if optimizer is not None and "optimizer_state_dict" in ckpt: - optimizer.load_state_dict(ckpt["optimizer_state_dict"]) - return ckpt -``` - -- [ ] **Step 2: Implement Trainer** - -```python -# src/jaeger/training/pytorch/trainer.py -import json -from pathlib import Path -from typing import Any, Dict, Optional - -import torch -import yaml -from torch.utils.data import DataLoader - -from jaeger.nnlib.pytorch.builder import ModelBuilder -from jaeger.nnlib.pytorch.checkpoints import load_checkpoint, save_checkpoint -from jaeger.training.pytorch.callbacks import CallbackList -from jaeger.training.pytorch.engine import train_one_epoch, validate_one_epoch -from jaeger.utils.logging import get_logger - - -logger = get_logger(log_file=None, log_path=None, level=3) - - -class Trainer: - """Orchestrates classifier/reliability/pretrain training.""" - - def __init__( - self, - config: Dict[str, Any], - device: torch.device, - use_amp: bool = False, - checkpoint_dir: Optional[Path] = None, - ): - self.config = config - self.device = device - self.use_amp = use_amp - self.checkpoint_dir = checkpoint_dir - self.builder = ModelBuilder(config) - self.models = self.builder.build_fragment_classifier() - self.callbacks = CallbackList() - - def fit_classifier( - self, - train_loader: DataLoader, - val_loader: DataLoader, - epochs: int, - initial_epoch: int = 0, - ): - model, optimizer, loss_fn = self.builder.compile_model(self.models, train_branch="classifier") - model = model.to(self.device) - metrics = self.builder.get_metrics(branch="classifier") - - for epoch in range(initial_epoch, epochs): - train_metrics = train_one_epoch( - model, train_loader, optimizer, loss_fn, self.device, self.use_amp, metrics - ) - val_metrics = validate_one_epoch( - model, val_loader, loss_fn, self.device, self.use_amp, metrics - ) - self.callbacks.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) - - def fit_reliability(self, train_loader, val_loader, epochs, initial_epoch=0): - model, optimizer, loss_fn = self.builder.compile_model(self.models, train_branch="reliability") - model = model.to(self.device) - metrics = self.builder.get_metrics(branch="reliability") - for epoch in range(initial_epoch, epochs): - train_metrics = train_one_epoch(model, train_loader, optimizer, loss_fn, self.device, self.use_amp, metrics) - val_metrics = validate_one_epoch(model, val_loader, loss_fn, self.device, self.use_amp, metrics) - self.callbacks.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) - - def fit_projection(self, train_loader, val_loader, epochs, initial_epoch=0): - model, optimizer, loss_fn = self.builder.compile_model(self.models, train_branch="pretrain") - model = model.to(self.device) - for epoch in range(initial_epoch, epochs): - train_metrics = train_one_epoch(model, train_loader, optimizer, loss_fn, self.device, self.use_amp) - val_metrics = validate_one_epoch(model, val_loader, loss_fn, self.device, self.use_amp) - self.callbacks.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) - - def save_model(self, suffix: str = "fragment", metadata: Optional[str] = None): - from jaeger.nnlib.pytorch.checkpoints import save_checkpoint - import yaml - - model_name = self.config["model"]["name"] - experiment = self.config["model"].get("experiment", "1") - base_dir = Path(self.config["training"].get("base_dir", ".")) - save_dir = base_dir / f"{model_name}_exp{experiment}" - save_dir.mkdir(parents=True, exist_ok=True) - - torch.save(self.models["jaeger_model"].state_dict(), save_dir / "model.pt") - with open(save_dir / "model_config.yaml", "w") as fh: - yaml.safe_dump(self.config, fh) - with open(save_dir / "class_label_map.yaml", "w") as fh: - yaml.safe_dump({"classes": self.config["model"].get("class_label_map", [])}, fh) - - if metadata: - Path(metadata).write_text(json.dumps({"model_path": str(save_dir)}, indent=2)) - - logger.info(f"Saved PyTorch model to {save_dir}") - return save_dir -``` - -- [ ] **Step 3: Smoke test** - -```bash -python -c "from jaeger.training.pytorch.trainer import Trainer; print('OK')" -``` - -Expected: `OK` - -- [ ] **Step 4: Commit** - -```bash -git add src/jaeger/training/pytorch/trainer.py src/jaeger/nnlib/pytorch/checkpoints.py -git commit -m "feat(pytorch): add Trainer and checkpoint helpers" -``` - ---- - -## Task 16: Implement callbacks - -**Files:** -- Create: `src/jaeger/training/pytorch/callbacks.py` -- Create: `tests/unit/training/pytorch/test_callbacks.py` - -- [ ] **Step 1: Implement callbacks** - -```python -# src/jaeger/training/pytorch/callbacks.py -from pathlib import Path - -import torch - - -class CallbackList: - def __init__(self, callbacks=None): - self.callbacks = callbacks or [] - - def on_epoch_end(self, epoch, train_metrics, val_metrics, model, optimizer): - for cb in self.callbacks: - cb.on_epoch_end(epoch, train_metrics, val_metrics, model, optimizer) - - -class ModelCheckpoint: - def __init__(self, checkpoint_dir: Path, monitor: str = "val_loss", mode: str = "min"): - self.checkpoint_dir = Path(checkpoint_dir) - self.checkpoint_dir.mkdir(parents=True, exist_ok=True) - self.monitor = monitor - self.mode = mode - self.best = float("inf") if mode == "min" else float("-inf") - - def on_epoch_end(self, epoch, train_metrics, val_metrics, model, optimizer): - value = val_metrics.get(self.monitor) - if value is None: - return - improved = (self.mode == "min" and value < self.best) or (self.mode == "max" and value > self.best) - if improved: - self.best = value - path = self.checkpoint_dir / f"epoch:{epoch}-loss:{value:.4f}.pt" - torch.save( - { - "model_state_dict": model.state_dict(), - "optimizer_state_dict": optimizer.state_dict(), - "epoch": epoch, - "metrics": val_metrics, - }, - path, - ) -``` - -- [ ] **Step 2: Smoke test** - -```bash -pytest tests/unit/training/pytorch/test_callbacks.py -v -``` - -Expected: PASS - -- [ ] **Step 3: Commit** - -```bash -git add src/jaeger/training/pytorch/callbacks.py tests/unit/training/pytorch/test_callbacks.py -git commit -m "feat(pytorch): add training callbacks" -``` - ---- - -## Task 17: Implement distributed training helper - -**Files:** -- Create: `src/jaeger/training/pytorch/distributed.py` - -- [ ] **Step 1: Implement DDP setup** - -```python -# src/jaeger/training/pytorch/distributed.py -import os - -import torch -import torch.distributed as dist -from torch.nn.parallel import DistributedDataParallel as DDP - - -def setup_distributed(): - rank = int(os.environ.get("RANK", 0)) - local_rank = int(os.environ.get("LOCAL_RANK", 0)) - world_size = int(os.environ.get("WORLD_SIZE", 1)) - dist.init_process_group("nccl") - torch.cuda.set_device(local_rank) - return rank, local_rank, world_size - - -def cleanup_distributed(): - dist.destroy_process_group() - - -def wrap_model_ddp(model, device_ids): - return DDP(model, device_ids=device_ids) -``` - -- [ ] **Step 2: Smoke test** - -```bash -python -c "from jaeger.training.pytorch.distributed import setup_distributed; print('OK')" -``` - -Expected: `OK` (does nothing if not launched with torchrun) - -- [ ] **Step 3: Commit** - -```bash -git add src/jaeger/training/pytorch/distributed.py -git commit -m "feat(pytorch): add distributed training helpers" -``` - ---- - -## Task 18: Update `jaeger train` command - -**Files:** -- Modify: `src/jaeger/commands/train.py` - -- [ ] **Step 1: Replace TF training orchestration with PyTorch** - -Keep the Click options. Replace `train_fragment_core` body with: - -```python -# src/jaeger/commands/train.py -def train_fragment_core(**kwargs): - import torch - from jaeger.training.pytorch.trainer import Trainer - from jaeger.data.pytorch.builders import build_datasets - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - config = load_model_config(Path(kwargs["config"])) - config["mix_precision"] = kwargs.get("mixed_precision", False) - config["from_last_checkpoint"] = kwargs.get("from_last_checkpoint", False) - config["force"] = kwargs.get("force", False) - - trainer = Trainer( - config=config, - device=device, - use_amp=kwargs.get("mixed_precision", False), - ) - - classifier_loaders = build_datasets(config, branch="classifier") - reliability_loaders = build_datasets(config, branch="reliability") - - train_cfg = config.get("training", {}) - - if kwargs.get("self_supervised_pretraining", False): - trainer.fit_projection( - classifier_loaders["train"], - classifier_loaders["validation"], - epochs=train_cfg.get("projection_epochs", 10), - ) - - if not kwargs.get("only_reliability_head", False): - trainer.fit_classifier( - classifier_loaders["train"], - classifier_loaders["validation"], - epochs=train_cfg.get("classifier_epochs", 100), - ) - - if not kwargs.get("only_classification_head", False): - trainer.fit_reliability( - reliability_loaders["train"], - reliability_loaders["validation"], - epochs=train_cfg.get("reliability_epochs", 100), - ) - - if kwargs.get("save_model", False) or kwargs.get("only_save", False): - trainer.save_model(suffix="fragment", metadata=kwargs.get("meta")) -``` - -- [ ] **Step 2: Smoke test** - -```bash -jaeger train --help -``` - -Expected: help text shows PyTorch-specific options. - -- [ ] **Step 3: Commit** - -```bash -git add src/jaeger/commands/train.py -git commit -m "feat(pytorch): wire jaeger train to PyTorch trainer" -``` - ---- - -## Task 19: Implement PyTorch inference backend - -**Files:** -- Create: `src/jaeger/inference/pytorch/model.py` -- Create: `src/jaeger/inference/pytorch/engine.py` - -- [ ] **Step 1: Implement model loader** - -```python -# src/jaeger/inference/pytorch/model.py -from pathlib import Path - -import torch -import yaml - -from jaeger.nnlib.pytorch.builder import ModelBuilder - - -class PyTorchJaegerModel: - def __init__(self, model_dir: Path, device: torch.device): - self.model_dir = Path(model_dir) - self.device = device - with open(self.model_dir / "model_config.yaml") as fh: - config = yaml.safe_load(fh) - builder = ModelBuilder(config) - self.models = builder.build_fragment_classifier() - self.jaeger_model = self.models["jaeger_model"].to(device) - self.jaeger_model.load_state_dict(torch.load(self.model_dir / "model.pt", map_location=device)) - self.jaeger_model.eval() - - from jaeger.seqops.maps import CODON_ID - self.codon_table = {c: i for i, c in enumerate(CODON_ID)} - - def predict(self, x, mask=None): - with torch.no_grad(): - return self.jaeger_model(x.to(self.device), mask.to(self.device) if mask is not None else None) -``` - -- [ ] **Step 2: Implement inference engine** - -```python -# src/jaeger/inference/pytorch/engine.py -import numpy as np -import torch -from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons -from jaeger.inference.pytorch.model import PyTorchJaegerModel - - -def run_inference(model_dir, fasta_path, batch_size=96, device="cuda", fsize=2000, stride=2000, **kwargs): - import pyfastx - from jaeger.inference.pytorch.model import PyTorchJaegerModel - - device = torch.device(device if torch.cuda.is_available() else "cpu") - model = PyTorchJaegerModel(model_dir, device) - - results = [] - for name, seq in pyfastx.Fasta(fasta_path, build_index=False): - windows = [] - for i in range(0, max(1, len(seq) - fsize + 1), stride): - window = seq[i : i + fsize] - if len(window) < 100: - continue - indices = dna_to_indices(window) - windows.append((name, len(seq), i, indices)) - - batch_preds = [] - for start in range(0, len(windows), batch_size): - batch = windows[start : start + batch_size] - lengths = [len(w[3]) for w in batch] - max_len = max(lengths) - padded = torch.zeros(len(batch), 6, max_len, dtype=torch.long) - mask = torch.zeros(len(batch), 6, max_len, dtype=torch.bool) - for j, (_, _, _, indices) in enumerate(batch): - translated = translate_to_codons(indices, model.codon_table) - l = translated.shape[1] - padded[j, :, :l] = translated - mask[j, :, :l] = True - with torch.no_grad(): - out = model.predict(padded, mask) - batch_preds.append(out["prediction"].cpu().numpy()) - - results.append({ - "name": name, - "length": len(seq), - "predictions": np.concatenate(batch_preds) if batch_preds else np.array([]), - }) - - return results -``` - -- [ ] **Step 3: Smoke test** - -```bash -python -c "from jaeger.inference.pytorch.model import PyTorchJaegerModel; print('OK')" -``` - -Expected: `OK` - -- [ ] **Step 4: Commit** - -```bash -git add src/jaeger/inference/pytorch/model.py src/jaeger/inference/pytorch/engine.py -git commit -m "feat(pytorch): add PyTorch inference backend" -``` - ---- - -## Task 20: Update `jaeger predict` command - -**Files:** -- Modify: `src/jaeger/commands/predict.py` - -- [ ] **Step 1: Add PyTorch model detection and routing** - -```python -# src/jaeger/commands/predict.py -def run_core(**kwargs): - from jaeger.inference.pytorch.engine import run_inference - return run_inference(**kwargs) -``` - -- [ ] **Step 2: Smoke test** - -```bash -jaeger predict --help -``` - -Expected: help text unchanged. - -- [ ] **Step 3: Commit** - -```bash -git add src/jaeger/commands/predict.py -git commit -m "feat(pytorch): route jaeger predict to PyTorch backend when model.pt present" -``` - ---- - -## Task 21: Update SLURM scripts and container - -**Files:** -- Modify: `slurm/*.slurm` -- Modify: `singularity/jaeger_singularity.def` - -- [ ] **Step 1: Update a training SLURM script** - -Edit `slurm/baseline_500bp.slurm`. Replace the training invocation line: - -```bash -# Before: -jaeger train -c train_config/nn_config_500bp_baseline.yaml - -# After: -torchrun --nproc_per_node=2 --nnodes=1 jaeger train -c train_config/nn_config_500bp_baseline.yaml --mixed_precision -``` - -- [ ] **Step 2: Update Singularity definition** - -Edit `singularity/jaeger_singularity.def`. Replace TensorFlow install lines with: - -```singularity -pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 -``` - -- [ ] **Step 3: Commit** - -```bash -git add slurm/ singularity/ -git commit -m "feat(pytorch): update SLURM scripts and Singularity container for PyTorch" -``` - ---- - -## Task 22: Delete TensorFlow training and inference code - -**Files:** -- Delete: `src/jaeger/nnlib/v1/` -- Delete: `src/jaeger/nnlib/v2/` -- Delete: `src/jaeger/nnlib/builder.py` -- Delete: `src/jaeger/nnlib/inference.py` -- Delete: `src/jaeger/nnlib/conversion.py` -- Delete: `src/jaeger/commands/predict_legacy.py` - -- [ ] **Step 1: Remove directories and files** - -```bash -git rm -rf src/jaeger/nnlib/v1/ src/jaeger/nnlib/v2/ -git rm src/jaeger/nnlib/builder.py src/jaeger/nnlib/inference.py src/jaeger/nnlib/conversion.py -git rm src/jaeger/commands/predict_legacy.py -``` - -- [ ] **Step 2: Verify imports still work** - -```bash -python -c "import jaeger; print('OK')" -``` - -Expected: `OK` - -- [ ] **Step 3: Commit** - -```bash -git commit -m "refactor(pytorch): remove TensorFlow training and inference code" -``` - ---- - -## Task 23: Final QA, benchmarks, and release prep - -**Files:** -- Modify: `docs/_source/train.md` -- Modify: `docs/_source/optimizations.md` -- Modify: `recipes/jaeger-bio/meta.yaml` -- Modify: `install.sh` - -- [ ] **Step 1: Run full test suite** - -```bash -pytest tests/unit tests/integration -v -``` - -Expected: All tests pass. - -- [ ] **Step 2: Run smoke training on Zeus** - -```bash -sbatch slurm/baseline_500bp.slurm -``` - -Expected: Job completes, checkpoints saved, logs show faster batches/sec than TF baseline. - -- [ ] **Step 3: Run inference smoke test** - -```bash -jaeger predict -i tests/fixtures/small.fasta -o /tmp/predict_out --model_path path/to/pytorch/model -``` - -Expected: Predictions written without errors. - -- [ ] **Step 4: Update documentation** - -- `docs/_source/train.md`: replace TensorFlow training instructions with PyTorch examples. -- `docs/_source/optimizations.md`: update inference backend table; remove TF-specific optimizations. - -- [ ] **Step 5: Update Bioconda recipe and install script** - -```yaml -# recipes/jaeger-bio/meta.yaml -requirements: - run: - - pytorch - - torchvision -``` - -```bash -# install.sh -# Replace TensorFlow install lines with PyTorch install lines. -``` - -- [ ] **Step 6: Bump version to 1.27.1** - -```bash -.github/scripts/bump-version.sh 1 27 1 -``` - -- [ ] **Step 7: Final commit and push** - -```bash -git add -A -git commit -m "chore(release): bump version to 1.27.1" -git push origin pytorch_migration -``` - ---- diff --git a/docs/superpowers/specs/2026-06-14-pytorch-migration-design.md b/docs/superpowers/specs/2026-06-14-pytorch-migration-design.md deleted file mode 100644 index 48164e7..0000000 --- a/docs/superpowers/specs/2026-06-14-pytorch-migration-design.md +++ /dev/null @@ -1,299 +0,0 @@ -# PyTorch Migration Design - -## Context - -Jaeger's training stack currently uses TensorFlow/Keras. Empirical testing shows TensorFlow is too slow for architectures that use attention mechanisms and energy-based training regimes. This design migrates the training pipeline — and the inference path that consumes trained models — to PyTorch in a single release. - -## Goals - -- Replace TensorFlow training with a plain PyTorch training stack. -- Migrate `jaeger predict` to support PyTorch-trained models natively. -- Maintain the existing YAML config format so current `train_config/` files remain valid. -- Validate parity with the TensorFlow baseline before merging. -- Run on the Zeus cluster using the existing Singularity + SLURM workflow. - -## Non-goals - -- Loading existing TensorFlow/Keras checkpoints into PyTorch. This is a clean break; models are trained from scratch. -- Supporting TensorFlow and PyTorch training backends simultaneously. The `pytorch_migration` branch becomes PyTorch-only for training. -- Loading old TensorFlow SavedModels in the PyTorch release. The release is a clean break for inference as well; users who need legacy models can stay on the previous Jaeger version. - -## Clarifying decisions - -| Question | Decision | -|---|---| -| Load existing TF checkpoints? | No — train from scratch. | -| PyTorch framework | Plain PyTorch (not Lightning or HF). | -| Data loading | Native `torch.utils.data.Dataset` + `DataLoader`. | -| Inference | Migrate to PyTorch in the same release. | -| Next version | `1.27.1` after merge. | - -## Section 1 — High-level architecture & module layout - -The migration creates a new, self-contained PyTorch stack under `src/jaeger/` while leaving the existing TensorFlow code untouched on the `pytorch_migration` branch. The TF code is deleted only at the very end of the branch, once PyTorch training and inference are validated. - -``` -src/jaeger/ -├── nnlib/ -│ └── pytorch/ # NEW: PyTorch model library -│ ├── __init__.py -│ ├── layers.py # PyTorch equivalents of custom Keras layers -│ ├── models.py # RepModel, ClassifierHead, ReliabilityHead, JaegerModel -│ ├── losses.py # ArcFace, hierarchical loss, classification losses -│ ├── metrics.py # Per-class precision/recall/specificity -│ ├── builder.py # Build nn.Modules from YAML config -│ └── checkpoints.py # Save/load .pt checkpoints + metadata -├── data/ -│ └── pytorch/ # NEW: PyTorch datasets & loaders -│ ├── __init__.py -│ ├── dataset_csv.py # CSV Dataset with padding & augmentations -│ ├── dataset_numpy.py # numpy_full / numpy_raw / numpy_raw_variable -│ ├── dataset_tfrecord.py # TFRecord Dataset (deferred) -│ ├── transforms.py # codon translation, mutation, frame shuffle -│ └── collate.py # padding collate functions -├── training/ -│ └── pytorch/ # NEW: training loops -│ ├── __init__.py -│ ├── trainer.py # Classification / reliability / pretrain loops -│ ├── engine.py # epoch loop, validation, mixed precision -│ ├── distributed.py # DDP setup/teardown -│ └── callbacks.py # ModelCheckpoint, EarlyStopping, LR scheduling -├── inference/ -│ └── pytorch/ # NEW: PyTorch inference backend -│ ├── __init__.py -│ ├── model.py # Load PyTorch checkpoint + config -│ └── engine.py # Batch inference, windowing, aggregation -├── commands/ -│ ├── train.py # MODIFIED: dispatch to PyTorch trainer -│ └── predict.py # MODIFIED: add PyTorch backend option -``` - -Key principles: -- The YAML config format stays the same so existing `train_config/` files remain valid. -- `DynamicModelBuilder` is replaced by `jaeger.nnlib.pytorch.builder.ModelBuilder`, but it reads the same config. -- Training entry point `jaeger train` detects the backend from config or CLI flag; on the `pytorch_migration` branch it defaults to PyTorch. -- Existing TF training and inference files (`nnlib/v1/`, `nnlib/v2/`, current `commands/train.py`, current `commands/predict.py`, etc.) remain during development for reference and are removed before merge. The PyTorch release does not load legacy TF models. - -## Section 2 — Model architecture - -The current Keras builder assembles four pieces: - -1. **Representation learner** — embedding + stack of residual/attention blocks → `(embedding, nmd, [gate])` -2. **Classification head** — dense layers → class logits -3. **Reliability head** — dense layers → confidence score -4. **Projection head** (optional) — for ArcFace self-supervised pretraining - -In PyTorch, each becomes an `nn.Module`: - -- `Embedding` — maps codon indices or one-hot nucleotides to `(B, 6, L, D)` or `(B, 6, L, 3, D)` depending on config. -- `ResidualBlock`, `AxialAttention`, `CrossFrameAttention`, `TransformerEncoder` — replicate the Keras custom layers in `nnlib/pytorch/layers.py`. -- `RepresentationModel` — composes embedding + blocks + pooling → returns `(embedding, nmd, gate)` tuple. -- `ClassificationHead` / `ReliabilityHead` / `ProjectionHead` — small MLPs. -- `JaegerModel` — combines all heads and returns the same output dict as the Keras model: `{"prediction": ..., "embedding": ..., "nmd": ..., "gate": ..., "reliability": ...}`. - -Key design choices: -- **Masking**: Keras uses `Masking` layers and `supports_masking`. In PyTorch we pass a `mask` tensor explicitly and use it in `MaskedBatchNorm`, `MaskedConv1D`, pooling, and attention. -- **Variable length**: PyTorch handles padded batches naturally; we use a collate function to pad sequences and pass the mask. -- **Mixed precision**: Use `torch.cuda.amp.autocast` + `GradScaler` (or `torch.amp` if using newer PyTorch). -- **Weight init**: Keep the same initializers where possible (orthogonal for embeddings/dense, glorot for conv) to make parity tests fair. - -The builder reads the same YAML config and constructs the module graph, preserving hyperparameters (filters, kernel sizes, attention heads, dropout, regularization). - -## Section 3 — Data pipeline - -Current `tf.data` supports five formats: `csv`, `tfrecord`, `numpy_raw`, `numpy_raw_variable`, `numpy_full`. The PyTorch pipeline replaces this with `torch.utils.data.Dataset` + `DataLoader`. - -Priority order: - -1. **Phase B1 — `numpy_full`** (highest priority) - - Fastest loading, no runtime augmentations. - - `Dataset` loads the `.npz` once and returns preprocessed tensors. - - `DataLoader` handles shuffling and batching; no padding needed because all sequences are already cropped. - -2. **Phase B2 — `numpy_raw`** - - Returns int8 sequences + labels. - - Apply codon translation, n-gram extraction, frame shuffle, mutation in a `transform` callable. - - Collate function pads to the batch max length and produces the mask. - -3. **Phase B3 — `csv`** - - Parse CSV lines, apply the same transforms as `numpy_raw`. - - Slower, kept for backward compatibility and small experiments. - -4. **Phase B4 — `tfrecord` / `numpy_raw_variable`** (deferred if time-constrained) - - TFRecord is lower priority now that NumPy formats exist. - - Variable-length can be handled with a custom collate similar to `numpy_raw`. - -Augmentations (from `process_string_train`) move to `src/jaeger/data/pytorch/transforms.py`: -- codon translation -- frame shuffling -- random mutation -- masking -- n-gram encoding - -Each `Dataset` returns a tuple `(input_tensor, label_tensor, mask_tensor)`. The collate function stacks inputs and labels and combines masks. - -For multi-GPU, use `DistributedSampler` with `DataLoader`. - -## Section 4 — Training loop - -The current training flow in `commands/train.py` is: - -1. Build models inside a distribution strategy scope. -2. Compile classifier, reliability, and optionally pretrain branches. -3. Call `.fit()` on each branch sequentially. -4. Save weights + SavedModel graphs. - -In PyTorch, this becomes an explicit training engine: - -- `Trainer` class in `src/jaeger/training/pytorch/trainer.py` owns: - - model - - optimizer(s) - - loss function(s) - - metric trackers - - device / DDP wrapper - - AMP gradient scaler - -- Three training modes, matching the Keras flow: - 1. **Self-supervised pretraining** — train `rep_model` + `projection_head` with ArcFace loss. - 2. **Classifier training** — train `rep_model` + `classification_head`. - 3. **Reliability training** — freeze `rep_model`, train `reliability_head`. - -- `Engine` in `src/jaeger/training/pytorch/engine.py` implements one epoch: - - training step with `autocast`, loss backward, gradient clipping, optimizer step - - validation step - - metric aggregation - - logging - -- `callbacks.py` implements: - - `ModelCheckpoint` — save best/last `.pt` checkpoints - - `EarlyStopping` - - `ReduceLROnPlateau` - - `TensorBoardLogger` / CSV logger - -- Multi-GPU on Zeus: - - `torchrun --nproc_per_node=2 ... jaeger train ...` - - `DistributedDataParallel` wrapper - - `DistributedSampler` for each DataLoader - -- Resume logic: - - Checkpoint file contains model state_dict, optimizer state_dict, epoch, and config. - - `--from_last_checkpoint` loads the latest checkpoint and continues. - -- CLI flags preserved: - - `--mixed_precision`, `--xla` (replaced with AMP), `--from_last_checkpoint`, `--force`, `--only_classification_head`, `--only_reliability_head`, `--only_heads`, `--save_model`, `--only_save`, `--self_supervised_pretraining`. - -## Section 5 — Inference migration - -Current `jaeger predict` supports multiple backends (TF SavedModel, ONNX, TFLite, TensorRT). For the PyTorch release, we add a first-class PyTorch backend. - -- Save format: - - `model.pt` — model state_dict - - `model_config.yaml` — architecture config (same as training config) - - `class_label_map.yaml` — class mapping - - Optional: `model.pt` can be a TorchScript module if tracing works for the variable-length inputs. - -- `src/jaeger/inference/pytorch/model.py`: - - Load config, instantiate `JaegerModel`, load state_dict. - - Move to device (CUDA if available). - - Set to eval mode. - -- `src/jaeger/inference/pytorch/engine.py`: - - Replicate the windowing + batching logic from the current `predict.py`. - - Run forward pass, collect `prediction`, `embedding`, `reliability`. - - Support `--precision fp16` via `autocast`. - - Support `--batch` and `--workers`. - -- CLI integration: - - `jaeger predict` expects PyTorch models (look for `model.pt` + `model_config.yaml`). - - Legacy TensorFlow SavedModels are not loaded by this release. - -- Quantization/optimization: - - Not in the first PyTorch release unless explicitly required. - - Can be added later via `torch.compile` or ONNX export from PyTorch. - -## Section 6 — Checkpointing, AMP, distributed training, and dependencies - -- **Checkpoints**: PyTorch checkpoints are dictionaries containing: - - `model_state_dict` - - `optimizer_state_dict` - - `epoch` - - `config` - - `metrics` - - `branch` (classifier / reliability / projection) - - File naming: `epoch:{epoch}-loss:{val_loss:.4f}.pt` - -- **Mixed precision**: Use `torch.amp.autocast(device_type="cuda")` with `GradScaler`. The Keras `mixed_float16` policy maps directly. - -- **Distributed training**: - - Launch via `torchrun` / `torch.distributed.launch`. - - `DistributedDataParallel` wraps the model. - - `DistributedSampler` for each DataLoader. - - Only the rank-0 process saves checkpoints and logs. - -- **Dependencies**: - - Replace `tensorflow[and-cuda]` with `torch` and `torchvision`. - - Keep `numpy`, `pyyaml`, `click`, `rich`, `scipy`, `pandas`, `polars`, `pyfastx`. - - Add `tensorboard` if using TensorBoard logging. - - Remove TF-only dependencies from the PyTorch branch's training path. - -- **Zeus cluster**: - - Update SLURM scripts to use `torchrun --nproc_per_node=2`. - - Ensure Singularity container has PyTorch + CUDA installed. - - Bind `Jaeger/` source directory at runtime as per cluster rules. - -## Section 7 — Testing & validation strategy - -To make sure the PyTorch models match or exceed TF quality: - -1. **Unit tests** for each new PyTorch module: - - `nnlib/pytorch/layers.py` — forward pass shape tests, mask behavior. - - `nnlib/pytorch/builder.py` — build models from existing configs. - - `data/pytorch/dataset_*.py` — dataset length, batch shapes, augmentation correctness. - - `training/pytorch/engine.py` — one training step, checkpoint save/load. - -2. **Parity tests**: - - Build the same architecture in TF and PyTorch with fixed random seeds. - - Compare forward-pass outputs on identical synthetic input to within tolerance. - - Compare a few training steps on a tiny dataset to verify gradients and weight updates are similar. - -3. **Smoke tests**: - - `jaeger train -c train_config/nn_config_500bp_baseline.yaml` runs for a few epochs. - - `jaeger predict` with a PyTorch-trained model on a small FASTA. - -4. **Benchmarks**: - - Measure batches/sec on Zeus `gpu` partition for attention-based configs. - - Compare against TF baseline on the same hardware. - -5. **CI**: - - Add a GitHub Actions job that installs the PyTorch branch and runs unit + smoke tests. - - Keep the TF tests running until the branch is merged. - -## Section 8 — Branch, release, and migration plan - -- **Branch**: `pytorch_migration` created from `main`. -- **Version policy**: Bump to `1.27.1` when the branch is merged. During development, keep the current version. -- **Phase milestones**: - 1. Model architecture + parity tests - 2. `numpy_full` data loader + classifier training loop - 3. Reliability + pretraining loops - 4. CSV / `numpy_raw` loaders - 5. PyTorch inference backend - 6. SLURM script updates + Zeus benchmarks - 7. Delete TF training/inference code, update docs, final QA -- **Merge criteria**: - - All unit and smoke tests pass. - - Parity tests show acceptable agreement with TF on the same config. - - Zeus benchmark shows faster training for attention/energy-based configs. - - `jaeger predict` works end-to-end with PyTorch-trained models. - - Bioconda recipe and install script updated to PyTorch dependencies. -- **Release**: After merge, tag and release as normal using `.github/scripts/bump-version.sh`. - -## Open questions / risks - -| Risk | Mitigation | -|---|---| -| Custom Keras layers (masked conv, batch norm, attention) are complex to port exactly. | Write parity tests layer-by-layer; fix mismatches before end-to-end training. | -| Mixed precision behavior differs between TF and PyTorch. | Test AMP separately; make it opt-in at first. | -| Zeus Singularity container may need a new PyTorch image. | Build and test the container early in Phase 6. | -| Variable-length padding/collate may be slower than `tf.data`. | Benchmark `numpy_full` first; optimize collate if needed. | -| Removing TF may break other commands (e.g., `quantize`, `convert_graph`, `taxonomy`). | Audit all `commands/` before deleting TF; keep TF-only commands if they have no PyTorch equivalent yet. | diff --git a/entities.json b/entities.json deleted file mode 100644 index 1a247de..0000000 --- a/entities.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "people": [ - "Yasas Wijesekara" - ], - "projects": [ - "jaeger-bio", - "PyPI", - "TensorRT", - "Conda", - "Jaeger" - ], - "topics": [] -} \ No newline at end of file diff --git a/mempalace.yaml b/mempalace.yaml deleted file mode 100644 index b92b69f..0000000 --- a/mempalace.yaml +++ /dev/null @@ -1,60 +0,0 @@ -wing: jaeger -rooms: -- name: slurm - description: Files from slurm/ - keywords: - - slurm - - slurm -- name: testing - description: Files from tests/ - keywords: - - testing - - tests -- name: scripts - description: Files from scripts/ - keywords: - - scripts - - scripts -- name: src - description: Files from src/ - keywords: - - src - - src -- name: singularity - description: Files from singularity/ - keywords: - - singularity - - singularity -- name: train_config - description: Files from train_config/ - keywords: - - train_config - - train_config -- name: documentation - description: Files from docs/ - keywords: - - documentation - - docs -- name: memory - description: Files from memory/ - keywords: - - memory - - memory -- name: recipes - description: Files from recipes/ - keywords: - - recipes - - recipes -- name: test_log - description: Files from test_log/ - keywords: - - test_log - - test_log -- name: test_cli - description: Files from test_cli/ - keywords: - - test_cli - - test_cli -- name: general - description: Files that don't fit other rooms - keywords: [] From cd67bbd7838d670e04a555d8adfbb0d1fcf43b78 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 18:56:56 +0200 Subject: [PATCH 39/64] feat(pytorch): complete PyTorch migration, Numba data conversion, and memory guard - Switch training and inference paths to PyTorch. - Add Numba JIT kernels for nucleotide/translated encoding in optimize-data. - Add memory guard to refuse one-hot conversions that exceed available RAM. - Add nucleotide/translated/DVF training configs. - Update CLI, builders, layers, models, callbacks, engine, and trainer. --- pyproject.toml | 4 + src/jaeger/cli.py | 84 +- src/jaeger/commands/train.py | 108 +- src/jaeger/commands/utils.py | 51 +- src/jaeger/data/pytorch/builders.py | 10 +- src/jaeger/data/pytorch/dataset_numpy.py | 11 +- src/jaeger/dataops/convert.py | 1075 ++++++++++-------- src/jaeger/nnlib/pytorch/builder.py | 83 +- src/jaeger/nnlib/pytorch/layers.py | 121 +- src/jaeger/nnlib/pytorch/models.py | 86 +- src/jaeger/training/pytorch/callbacks.py | 19 +- src/jaeger/training/pytorch/engine.py | 12 + src/jaeger/training/pytorch/trainer.py | 102 +- tests/unit/commands/test_train_pytorch.py | 28 + tests/unit/data/pytorch/test_builders.py | 41 + tests/unit/test_dataops_convert.py | 251 +++- tests/unit/training/pytorch/test_engine.py | 32 + tests/unit/training/pytorch/test_trainer.py | 26 + train_config/nn_config_500bp_dvf.yaml | 198 ++++ train_config/nn_config_500bp_nucleotide.yaml | 214 ++++ train_config/nn_config_500bp_translated.yaml | 209 ++++ 21 files changed, 2106 insertions(+), 659 deletions(-) create mode 100644 train_config/nn_config_500bp_dvf.yaml create mode 100644 train_config/nn_config_500bp_nucleotide.yaml create mode 100644 train_config/nn_config_500bp_translated.yaml diff --git a/pyproject.toml b/pyproject.toml index e150399..f4ce25b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,6 +76,7 @@ legacy = [ test = [ "pytest >=8.0", "pytest-mock >=3.14", + "numba>=0.59", ] onnx = [ "onnxruntime-gpu >=1.26", @@ -87,3 +88,6 @@ taxonomy = [ "taxopy", "faiss-cpu", ] +numba = [ + "numba>=0.59", +] diff --git a/src/jaeger/cli.py b/src/jaeger/cli.py index 3ddd8a6..be37699 100644 --- a/src/jaeger/cli.py +++ b/src/jaeger/cli.py @@ -302,6 +302,13 @@ def tune(**kwargs): required=False, help="Path to write metadata from the container", ) +@click.option( + "--progress-bar", + is_flag=True, + required=False, + default=False, + help="Show a per-batch Rich progress bar with loss and speed", +) @click.option( "-v", "--verbose", @@ -823,21 +830,33 @@ def convert(**kwargs): pass +_VALID_CODON_MAPS = [ + "CODON_ID", + "AA_ID", + "PC5_ID", + "MURPHY10_ID", + "COD_ID", + "PC2_ID", +] + + @utils.command( context_settings=dict(ignore_unknown_options=True, show_default=True), - help=""" - Convert training data to optimized formats for faster loading. - Preprocesses CSV data once so training can skip live preprocessing. + help=f""" + Convert training data to optimized NPZ format for faster loading. + Output is always variable-length and padded to the longest sample. usage ----- - jaeger utils optimize-data -i train.csv -o train.npz --format numpy_full + jaeger utils optimize-data -i train.csv -o train.npz --format translated Supported formats: - numpy_raw - int8 sequences + runtime preprocessing at train time - numpy_full - fully preprocessed, fastest loading - numpy_raw_variable - variable-length int8 sequences + nucleotide - 2-strand nucleotide representation + translated - 6-frame codon representation + both - store both nucleotide and translated arrays + + Available codon maps: {', '.join(_VALID_CODON_MAPS)} """, ) @click.option( @@ -857,18 +876,25 @@ def convert(**kwargs): @click.option( "--format", type=click.Choice( - ["numpy_raw", "numpy_full", "numpy_raw_variable"], + ["nucleotide", "translated", "both"], case_sensitive=False, ), required=True, - help="Output format", + help="Output representation", ) @click.option( "--crop-size", type=int, default=500, show_default=True, - help="Sequence crop size", + help="Maximum sequence length to process", +) +@click.option( + "--stride", + type=int, + default=0, + show_default=True, + help="Sliding-window step (0 = one crop per sequence)", ) @click.option( "--num-classes", @@ -884,20 +910,42 @@ def convert(**kwargs): help="Number of parallel workers (default: all CPUs)", ) @click.option( - "--use-embedding-layer/--no-embedding-layer", - default=True, + "--one-hot", + is_flag=True, + default=False, + show_default=True, + help="Output one-hot float tensors instead of integer indices", +) +@click.option( + "--pad-int", + type=int, + default=0, + show_default=True, + help="Integer value used for padding and stored as 'pad_int' in the NPZ", +) +@click.option( + "--codon-map", + type=click.Choice(_VALID_CODON_MAPS, case_sensitive=False), + default="CODON_ID", show_default=True, - help="Use embedding layer (int indices) vs one-hot (numpy formats)", + help="Codon-to-int map from jaeger.seqops.maps (translated/both only)", ) @click.option( "--max-length", type=int, default=5000, show_default=True, - help="Maximum sequence length (numpy_raw_variable only)", + help="Deprecated; kept for compatibility but ignored", +) +@click.option( + "--compress", + type=click.Choice(["fast", "default", "none"], case_sensitive=False), + default="fast", + show_default=True, + help="NPZ compression level: fast (zlib 1), default (zlib 6), or none.", ) def optimize_data(**kwargs): - """Convert CSV training data to optimized formats""" + """Convert CSV training data to optimized NPZ format""" from jaeger.commands.utils import optimize_data_core optimize_data_core( @@ -905,10 +953,14 @@ def optimize_data(**kwargs): output_path=kwargs.get("output"), format=kwargs.get("format"), crop_size=kwargs.get("crop_size"), + stride=kwargs.get("stride"), num_classes=kwargs.get("num_classes"), num_workers=kwargs.get("num_workers"), - use_embedding_layer=kwargs.get("use_embedding_layer"), + one_hot=kwargs.get("one_hot"), + pad_int=kwargs.get("pad_int"), + codon_map=kwargs.get("codon_map"), max_length=kwargs.get("max_length"), + compress=kwargs.get("compress"), ) pass diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 0f892de..4c2152a 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -66,11 +66,31 @@ def _initialize_lazy_layers( model_cfg = config.get("model", {}) embedding_cfg = model_cfg.get("embedding", {}) sp_cfg = model_cfg.get("string_processor", {}) - frames = int(embedding_cfg.get("frames", 6)) + input_type = embedding_cfg.get("input_type", "translated") + input_shape = embedding_cfg.get("input_shape") crop_size = int(sp_cfg.get("crop_size", 500)) device = next(models["jaeger_classifier"].parameters()).device - dummy = torch.zeros((1, frames, crop_size), dtype=torch.long, device=device) - mask = torch.ones((1, frames, crop_size), dtype=torch.bool, device=device) + + # Use a short dummy length; lazy layers only need to infer channel dims. + length = 64 if input_type == "nucleotide" else max(1, crop_size // 3 - 1) + + if input_shape is not None and len(input_shape) == 3: + # One-hot input, e.g. [2, null, 4] or [6, null, vocab_size]. + frames = int(input_shape[0]) + channels = int(input_shape[2]) + dummy = torch.zeros((1, frames, length, channels), dtype=torch.float32, device=device) + else: + # Integer-index input, e.g. [2, null] or [6, null]. + frames = int(input_shape[0]) if input_shape else ( + 2 if input_type == "nucleotide" else 6 + ) + dummy = torch.zeros((1, frames, length), dtype=torch.long, device=device) + + mask = torch.ones( + dummy.shape[:-1] if dummy.dim() == 4 else dummy.shape, + dtype=torch.bool, + device=device, + ) with torch.no_grad(): models["jaeger_classifier"](dummy, mask) @@ -118,47 +138,44 @@ def _callback_path_from_config( return None +_SUPPORTED_CALLBACKS = {"EarlyStopping", "ModelCheckpoint"} + + def _build_callbacks(train_cfg: Dict[str, Any]) -> List[Any]: """Build PyTorch callbacks from the training config.""" callbacks: List[Any] = [] callbacks_cfg = train_cfg.get("callbacks", {}).get("classifier", []) - # EarlyStopping - early_params = _callback_path_from_config(callbacks_cfg, "EarlyStopping") - patience = train_cfg.get("patience") - if early_params is not None and patience is None: - patience = early_params.get("patience") - if patience is not None: - callbacks.append( - EarlyStopping( - monitor=(early_params or {}).get("monitor", "val_loss") - or train_cfg.get("early_stopping_monitor", "val_loss"), - patience=int(patience), - mode=(early_params or {}).get("mode", "min") - or train_cfg.get("early_stopping_mode", "min"), - restore_best_weights=(early_params or {}).get( - "restore_best_weights", True - ), - ) - ) - - # ModelCheckpoint - checkpoint_params = _callback_path_from_config(callbacks_cfg, "ModelCheckpoint") - checkpoint_path = train_cfg.get("checkpoint_path") - if checkpoint_params is not None and checkpoint_path is None: - checkpoint_path = checkpoint_params.get("filepath") - if checkpoint_path is not None: - callbacks.append( - ModelCheckpoint( - filepath=checkpoint_path, - monitor=(checkpoint_params or {}).get("monitor", "val_loss") - or train_cfg.get("checkpoint_monitor", "val_loss"), - mode=(checkpoint_params or {}).get("mode", "min") - or train_cfg.get("checkpoint_mode", "min"), - save_best_only=(checkpoint_params or {}).get("save_best_only", True), - verbose=int((checkpoint_params or {}).get("verbose", 1) or 0), - ) - ) + for entry in callbacks_cfg: + name = entry.get("name") + params = entry.get("params") or {} + if name == "EarlyStopping": + patience = params.get("patience") or train_cfg.get("patience") + if patience is not None: + callbacks.append( + EarlyStopping( + monitor=params.get("monitor", "val_loss"), + patience=int(patience), + mode=params.get("mode", "min"), + restore_best_weights=params.get("restore_best_weights", True), + ) + ) + elif name == "ModelCheckpoint": + checkpoint_path = params.get("filepath") or train_cfg.get("checkpoint_path") + if checkpoint_path is not None: + callbacks.append( + ModelCheckpoint( + filepath=checkpoint_path, + monitor=params.get("monitor", "val_loss"), + mode=params.get("mode", "min"), + save_best_only=params.get("save_best_only", True), + verbose=int(params.get("verbose", 1) or 0), + ) + ) + elif name in _SUPPORTED_CALLBACKS: + continue + else: + logger.warning("Ignoring unsupported callback: %s", name) return callbacks @@ -332,8 +349,17 @@ def train_fragment_core(**kwargs): opt_params = train_cfg.get("optimizer_params", {}) or {} train_cfg["optimizer_params"] = _normalize_optimizer_params(opt_params) + class_weights = train_cfg.get("classifier_class_weights") + if class_weights is not None: + weight_tensor = torch.zeros(builder.classifier_out_dim) + for idx, w in class_weights.items(): + weight_tensor[int(idx)] = float(w) + class_weights = weight_tensor.to(device) + model, optimizer, loss_fn = builder.compile_model( - models, train_branch="classifier" + models, + train_branch="classifier", + class_weights=class_weights, ) checkpoint_path = None @@ -366,6 +392,7 @@ def train_fragment_core(**kwargs): checkpoint_dir=str(classifier_dir), history_path=str(classifier_dir / "history.json"), branch="classifier", + progress_bar=kwargs.get("progress_bar", False), ) trainer.fit() @@ -422,6 +449,7 @@ def train_fragment_core(**kwargs): checkpoint_dir=str(reliability_dir), history_path=str(reliability_dir / "history.json"), branch="reliability", + progress_bar=kwargs.get("progress_bar", False), ) rel_trainer.fit() diff --git a/src/jaeger/commands/utils.py b/src/jaeger/commands/utils.py index 51e2a74..89dd01c 100644 --- a/src/jaeger/commands/utils.py +++ b/src/jaeger/commands/utils.py @@ -554,17 +554,21 @@ def optimize_data_core( output_path: str, format: str, crop_size: int = 500, + stride: int = 0, num_classes: int = 3, num_workers: int | None = None, - use_embedding_layer: bool = True, + one_hot: bool = False, + pad_int: int = 0, + codon_map: str = "CODON_ID", max_length: int = 5000, + compress: str = "fast", ): - """Convert Jaeger CSV training data to an optimized format. + """Convert Jaeger CSV training data to an optimized NPZ format. - Supports three output formats: - - numpy_raw: int8 sequences + runtime preprocessing at train time - - numpy_full: Fully preprocessed, fastest loading - - numpy_raw_variable: Variable-length int8 sequences + Produces variable-length padded outputs for the requested representation: + - nucleotide: 2-strand nucleotide indices/one-hot + - translated: 6-frame codon indices/one-hot + - both: store both nucleotide and translated arrays Parameters ---------- @@ -573,40 +577,45 @@ def optimize_data_core( output_path : str Path to output file. format : str - One of: numpy_raw, numpy_full, numpy_raw_variable. + One of: nucleotide, translated, both. crop_size : int - Sequence crop size (default: 500). + Maximum sequence length to process (default: 500). + stride : int + Sliding-window step in nucleotides. 0 means one crop per sequence. num_classes : int Number of classes (default: 3). num_workers : int | None - Number of parallel workers for CPU-bound formats. + Number of parallel workers for CPU-bound conversion. Defaults to all CPUs. - use_embedding_layer : bool - Unused (kept for API compatibility). + one_hot : bool + Output float one-hot tensors instead of integer indices. + pad_int : int + Integer value used for padding and stored as ``pad_int`` in the NPZ. + codon_map : str + Name of a length-64 codon map in ``jaeger.seqops.maps``. max_length : int - For numpy_raw_variable: maximum sequence length. + Deprecated; kept for compatibility but ignored. + compress : str + NPZ compression level: fast (zlib 1), default (zlib 6), or none. """ format = format.lower() - valid_formats = [ - "numpy_raw", - "numpy_full", - "numpy_raw_variable", - ] + valid_formats = ["nucleotide", "translated", "both"] if format not in valid_formats: raise ValueError( f"Invalid format: {format}. Choose from: {', '.join(valid_formats)}" ) - print(f"Converting {input_path} -> {output_path}") - print(f"Format: {format}, Crop size: {crop_size}, Num classes: {num_classes}") - convert_dataset( input_path=input_path, output_path=output_path, format=format, crop_size=crop_size, + stride=stride, num_classes=num_classes, num_workers=num_workers, - use_embedding_layer=use_embedding_layer, + one_hot=one_hot, + pad_int=pad_int, + codon_map=codon_map, max_length=max_length, + compress=compress, ) diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/data/pytorch/builders.py index 5129836..000891e 100644 --- a/src/jaeger/data/pytorch/builders.py +++ b/src/jaeger/data/pytorch/builders.py @@ -60,8 +60,16 @@ def build_datasets( raise ValueError(f"No paths found for {branch}/{split}") data_format = string_cfg.get("data_format", "numpy_full") + input_key = string_cfg.get( + "input_key", + "nucleotide" + if model_cfg.get("embedding", {}).get("input_type") == "nucleotide" + else "translated", + ) if data_format == "numpy_full": - split_datasets: list[Dataset] = [NumpyFullDataset(path) for path in paths] + split_datasets: list[Dataset] = [ + NumpyFullDataset(path, input_key=input_key) for path in paths + ] datasets[split] = ( split_datasets[0] if len(split_datasets) == 1 diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/data/pytorch/dataset_numpy.py index 5ba0c57..af8af1a 100644 --- a/src/jaeger/data/pytorch/dataset_numpy.py +++ b/src/jaeger/data/pytorch/dataset_numpy.py @@ -34,18 +34,21 @@ def __init__( data = np.load(path, allow_pickle=False) self.inputs = torch.from_numpy(data[input_key]) self.labels = torch.from_numpy(data[label_key]) + self.pad_int = int(data.get("pad_int", 0)) def __len__(self) -> int: return len(self.labels) def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: x = self.inputs[idx] + if not torch.is_floating_point(x): + x = x.long() if x.dim() == 2: - # (frames, length) - mask = x != 0 + # (frames, length) integer indices + mask = x != self.pad_int elif x.dim() == 3: - # (frames, length, channels) -> mask over frames and length - mask = (x != 0).any(dim=-1) + # (frames, length, channels) one-hot tensor + mask = (x != self.pad_int).any(dim=-1) else: raise ValueError( f"NumpyFullDataset expects 2-D or 3-D inputs, got rank {x.dim()}" diff --git a/src/jaeger/dataops/convert.py b/src/jaeger/dataops/convert.py index ec1a8f6..87b77eb 100644 --- a/src/jaeger/dataops/convert.py +++ b/src/jaeger/dataops/convert.py @@ -1,19 +1,44 @@ """Dataset format converters. -Converts CSV training data to optimized formats: -- ``numpy_raw`` — int8 DNA sequences (fast loading, runtime preprocessing). -- ``numpy_full`` — fully preprocessed tensors (fastest loading, no augmentations). -- ``numpy_raw_variable`` — variable-length int8 sequences. +Converts CSV training data to optimized, variable-length padded NPZ formats: +- ``nucleotide`` - 2-strand nucleotide indices or one-hot tensors. +- ``translated`` - 6-frame codon indices or one-hot tensors. +- ``both`` - store both nucleotide and translated arrays. -The ``numpy_full`` converter uses Numba JIT when available for ~5× speedup. +Integer outputs reserve ``pad_int`` (default 0) for padding. One-hot outputs use +all-zero vectors for padding and for ambiguous bases (N). """ from __future__ import annotations +import multiprocessing import time from functools import partial -from multiprocessing import Pool, cpu_count +from multiprocessing import cpu_count +from typing import Dict, List, Tuple + import numpy as np +import psutil + + +def _save_npz(output_path: str, data: Dict[str, np.ndarray], compress: str): + """Save arrays to NPZ with the requested compression level.""" + if compress == "default": + np.savez_compressed(output_path, **data) + elif compress == "none": + np.savez(output_path, **data) + else: # fast + import io + import zipfile + + with zipfile.ZipFile( + output_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=1 + ) as zf: + for key, arr in data.items(): + buf = io.BytesIO() + np.save(buf, arr) + zf.writestr(f"{key}.npy", buf.getvalue()) + # --------------------------------------------------------------------------- # Numba availability @@ -35,107 +60,141 @@ def wrapper(func): # --------------------------------------------------------------------------- -# numpy_raw (int8 sequences) +# Lookup tables for Numba kernels # --------------------------------------------------------------------------- -_BASE_MAP = {"A": 0, "T": 1, "G": 2, "C": 3, "N": 4} -_BASE_MAP_LOWER = {"a": 0, "t": 1, "g": 2, "c": 3, "n": 4} +def _build_ascii_to_base_idx() -> np.ndarray: + """Return a (256,) int8 array mapping ASCII bytes to A/T/G/C/N indices.""" + table = np.full(256, 4, dtype=np.int8) + table[ord("A")] = 0 + table[ord("T")] = 1 + table[ord("G")] = 2 + table[ord("C")] = 3 + table[ord("a")] = 0 + table[ord("t")] = 1 + table[ord("g")] = 2 + table[ord("c")] = 3 + return table + + +def _build_complement_lut() -> np.ndarray: + """Return a (5,) int8 array mapping A<->T, G<->C, N->N.""" + return np.array([1, 0, 3, 2, 4], dtype=np.int8) + + +def _build_codon_lut(codon_map: np.ndarray) -> np.ndarray: + """Return a flat (125,) int32 lookup table for 3-base combinations. + + Index = base0 * 25 + base1 * 5 + base2 where each base is 0..4. + Unknown codons map to -1. + """ + from jaeger.seqops.maps import CODONS + codon_to_id = {c: i for i, c in enumerate(CODONS)} + lut = np.full(125, -1, dtype=np.int32) + bases = ["A", "T", "G", "C", "N"] + for i in range(5): + for j in range(5): + for k in range(5): + codon = bases[i] + bases[j] + bases[k] + cid = codon_to_id.get(codon, -1) + if 0 <= cid < len(codon_map): + lut[i * 25 + j * 5 + k] = int(codon_map[cid]) + return lut -def _encode_sequence_int8(seq: str, crop_size: int) -> np.ndarray: - """Encode DNA sequence to int8 array.""" - arr = np.zeros(crop_size, dtype=np.int8) - length = min(len(seq), crop_size) - for i in range(length): - c = seq[i] - if c in _BASE_MAP: - arr[i] = _BASE_MAP[c] - elif c in _BASE_MAP_LOWER: - arr[i] = _BASE_MAP_LOWER[c] - else: - arr[i] = 4 - return arr - - -def _process_chunk_numpy_raw( - lines: list[str], crop_size: int, num_classes: int -) -> tuple: - """Process a chunk of CSV lines to int8 sequences.""" - n = len(lines) - sequences = np.zeros((n, crop_size), dtype=np.int8) - labels = np.zeros((n, num_classes), dtype=np.float32) - - for i, line in enumerate(lines): - line = line.strip() - if not line: - continue - comma_idx = line.find(",") - if comma_idx == -1: - continue - label = int(line[:comma_idx]) - seq = line[comma_idx + 1 :] - sequences[i] = _encode_sequence_int8(seq, crop_size) - labels[i, label] = 1.0 - return sequences, labels +_ASCII_TO_BASE_IDX = _build_ascii_to_base_idx() +_COMPLEMENT_LUT = _build_complement_lut() -def _convert_to_numpy_raw( - csv_path: str, - output_path: str, - crop_size: int, - num_classes: int, - num_workers: int | None, -): - """Convert CSV to NumPy raw (int8 sequences).""" - total = sum(1 for _ in open(csv_path)) - print(f"Total samples: {total}") +# Use forkserver on Linux to avoid fork-deadlocks when the parent process has +# background threads (e.g. from PyTorch/NumPy/OpenMP). Fall back to the default +# context on platforms where forkserver is unavailable. +try: + _MP_CONTEXT = multiprocessing.get_context("forkserver") +except ValueError: + _MP_CONTEXT = multiprocessing.get_context() - if num_workers is None: - num_workers = cpu_count() - print(f"Using {num_workers} workers") - with open(csv_path) as f: - lines = f.readlines() +# --------------------------------------------------------------------------- +# Memory guard helpers +# --------------------------------------------------------------------------- +def _estimate_onehot_memory( + total_rows: int, + crop_size: int, + format: str, + one_hot: bool, + codon_map_arr: np.ndarray | None = None, +) -> int: + """Return the estimated raw float32 bytes needed for one-hot output. - chunk_size = max(1, len(lines) // num_workers) - chunks = [lines[i : i + chunk_size] for i in range(0, len(lines), chunk_size)] - print(f"Split into {len(chunks)} chunks") + Returns 0 when ``one_hot`` is False or ``format`` does not require it. + """ + if not one_hot: + return 0 + + total_rows = max(0, total_rows) + crop_size = max(0, crop_size) + estimated = 0 + + if format in ("nucleotide", "both"): + # (2 strands, crop_size, 4 bases) float32 + estimated += total_rows * 2 * crop_size * 4 * np.dtype(np.float32).itemsize + + if format in ("translated", "both") and codon_map_arr is not None: + vocab_size = int(codon_map_arr.max()) + 1 + max_codon_len = max(0, crop_size // 3 - 1) + # (6 frames, max_codon_len, vocab_size) float32 + estimated += ( + total_rows * 6 * max_codon_len * vocab_size * np.dtype(np.float32).itemsize + ) - start = time.time() - process_fn = partial( - _process_chunk_numpy_raw, crop_size=crop_size, num_classes=num_classes - ) + return int(estimated) - with Pool(num_workers) as pool: - results = pool.map(process_fn, chunks) - all_sequences = np.concatenate([r[0] for r in results], axis=0) - all_labels = np.concatenate([r[1] for r in results], axis=0) +def _check_onehot_memory( + estimated_bytes: int, + available_bytes: int, + fraction: float = 0.75, +) -> None: + """Raise MemoryError if the estimated one-hot buffer is too large. + + The ``fraction`` sets a safety margin below 100 % of available RAM to + leave room for the OS, worker overhead, and intermediate copies. + """ + if estimated_bytes <= 0: + return + + if estimated_bytes > available_bytes * fraction: + raise MemoryError( + f"Estimated one-hot output requires {estimated_bytes / (1024 ** 3):.1f} GiB, " + f"but only {available_bytes / (1024 ** 3):.1f} GiB RAM is available " + f"(safety fraction: {fraction:.0%}). " + "Avoid one-hot encoding for large files; use integer encoding instead " + "and let the dataloader/model convert to one-hot at training time." + ) - elapsed = time.time() - start - print( - f"Processed {total} samples in {elapsed:.2f}s ({total / elapsed:.0f} samples/sec)" - ) - np.savez_compressed(output_path, sequences=all_sequences, labels=all_labels) - seq_mb = all_sequences.nbytes / (1024 * 1024) - label_mb = all_labels.nbytes / (1024 * 1024) - print(f"Saved: sequences={seq_mb:.1f}MB, labels={label_mb:.1f}MB") +def _count_lines(path: str) -> int: + """Count newline-terminated rows without loading the whole file.""" + count = 0 + with open(path, "rb") as f: + for _ in f: + count += 1 + return count # --------------------------------------------------------------------------- -# numpy_full (fully preprocessed) +# Constants # --------------------------------------------------------------------------- -_CODON_TO_ID = None -_COMP = {"A": "T", "T": "A", "G": "C", "C": "G", "N": "N"} +_BASE_MAP = {"A": 0, "T": 1, "G": 2, "C": 3, "N": 4} +_BASE_MAP_LOWER = {"a": 0, "t": 1, "g": 2, "c": 3, "n": 4} +_COMPLEMENT = {"A": "T", "T": "A", "G": "C", "C": "G", "N": "N"} +_NUCLEOTIDE_ONEHOT_ORDER = "ATGC" -# Precomputed lookup tables for numba -_CODON_LUT = None -_ASCII_TO_IDX = None -_COMP_LUT = None +_CODON_TO_ID: Dict[str, int] | None = None -def _get_codon_to_id(): +def _get_codon_to_id() -> Dict[str, int]: """Lazy-load codon mapping to avoid E402 module-level import.""" global _CODON_TO_ID if _CODON_TO_ID is None: @@ -145,500 +204,540 @@ def _get_codon_to_id(): return _CODON_TO_ID -def _build_numba_lookups(): - """Build flat lookup tables for numba-optimized processing.""" - global _CODON_LUT, _ASCII_TO_IDX, _COMP_LUT - if _CODON_LUT is not None: - return _CODON_LUT, _ASCII_TO_IDX, _COMP_LUT +def _get_codon_map(codon_map_name: str) -> np.ndarray: + """Return the requested length-64 codon map from jaeger.seqops.maps.""" + import jaeger.seqops.maps as maps - from jaeger.seqops.maps import CODONS + mapping = getattr(maps, codon_map_name, None) + if mapping is None: + raise ValueError(f"Unknown codon map: {codon_map_name}") + mapping = np.asarray(mapping) + if mapping.shape != (64,): + raise ValueError( + f"Codon map {codon_map_name} must have shape (64,), got {mapping.shape}" + ) + return mapping - base_to_idx = {"A": 0, "T": 1, "G": 2, "C": 3, "N": 4} - # Complement lookup: A<->T, G<->C, N<->N - _COMP_LUT = np.array([1, 0, 3, 2, 4], dtype=np.int8) +def _base_index(char: str) -> int: + """Return 0..4 for ACGTN (case-insensitive).""" + if char in _BASE_MAP: + return _BASE_MAP[char] + return _BASE_MAP_LOWER.get(char, 4) - # Flat codon lookup: 125 entries (5x5x5) - codon_to_id = {c: i for i, c in enumerate(CODONS)} - _CODON_LUT = np.full(125, -1, dtype=np.int32) - bases = ["A", "T", "G", "C", "N"] - for i in range(5): - for j in range(5): - for k in range(5): - codon = bases[i] + bases[j] + bases[k] - _CODON_LUT[i * 25 + j * 5 + k] = codon_to_id.get(codon, -1) - # ASCII to base index lookup - _ASCII_TO_IDX = np.full(256, 4, dtype=np.int8) - for c, i in base_to_idx.items(): - _ASCII_TO_IDX[ord(c)] = i - - return _CODON_LUT, _ASCII_TO_IDX, _COMP_LUT +def _reverse_complement(seq: str) -> str: + """Return the reverse-complement DNA string.""" + return "".join(_COMPLEMENT.get(b, "N") for b in reversed(seq.upper())) +# --------------------------------------------------------------------------- +# Nucleotide encoding +# --------------------------------------------------------------------------- @njit(cache=False) -def _process_sequence_numba( - seq_bytes, crop_size, seq_len, codon_lut, comp_lut, ascii_lut -): - """Process DNA sequence to 6-frame codon embeddings (numba-optimized).""" - length = min(len(seq_bytes), crop_size) +def _encode_nucleotide_integer_numba( + seq_bytes: np.ndarray, crop_size: int, pad_int: int +) -> Tuple[np.ndarray, int]: + """Numba kernel for 2-strand integer nucleotide encoding.""" + length = min(seq_bytes.shape[0], crop_size) + arr = np.full((2, crop_size), pad_int, dtype=np.int32) + ascii_lut = _ASCII_TO_BASE_IDX + comp_lut = _COMPLEMENT_LUT - # Convert ASCII bytes to base indices - bases = np.empty(length, dtype=np.int8) for i in range(length): - bases[i] = ascii_lut[seq_bytes[i]] - - # Compute frame lengths - offset = [-2, -1, 0][length % 3] - ngrams_len = length - 3 + 1 + idx = ascii_lut[seq_bytes[i]] + arr[0, i] = idx + 1 - end_f1 = ngrams_len - 3 + offset - end_f2 = ngrams_len - 2 + offset - end_f3 = ngrams_len - 1 + offset - - nf1 = max(0, (end_f1 + 2) // 3) - nf2 = max(0, (end_f2 + 1) // 3) - nf3 = max(0, end_f3 // 3) - - # Reverse complement - rev_bases = np.empty(length, dtype=np.int8) for i in range(length): - rev_bases[i] = comp_lut[bases[length - 1 - i]] + idx = comp_lut[ascii_lut[seq_bytes[length - 1 - i]]] + arr[1, i] = idx + 1 - nr1 = max(0, (end_f1 + 2) // 3) - nr2 = max(0, (end_f2 + 1) // 3) - nr3 = max(0, end_f3 // 3) + return arr, length - min_len = min(nf1, nf2, nf3, nr1, nr2, nr3) - if min_len <= 0: - return np.empty((6, 0), dtype=np.int32) - min_len = min(min_len, seq_len) - frames = np.empty((6, min_len), dtype=np.int32) +def _encode_nucleotide_integer( + seq: str, crop_size: int, pad_int: int +) -> Tuple[np.ndarray, int]: + """Encode a DNA sequence as integer nucleotide indices (2 strands).""" + if HAS_NUMBA: + seq_bytes = np.frombuffer(seq[:crop_size].encode("ascii"), dtype=np.uint8) + return _encode_nucleotide_integer_numba(seq_bytes, crop_size, pad_int) - # Forward frames - for i in range(min_len): - idx0 = bases[i * 3] - idx1 = bases[i * 3 + 1] - idx2 = bases[i * 3 + 2] - frames[0, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] - - idx0 = bases[i * 3 + 1] - idx1 = bases[i * 3 + 2] - idx2 = bases[i * 3 + 3] - frames[1, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] - - idx0 = bases[i * 3 + 2] - idx1 = bases[i * 3 + 3] - idx2 = bases[i * 3 + 4] - frames[2, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] - - # Reverse frames - for i in range(min_len): - idx0 = rev_bases[i * 3] - idx1 = rev_bases[i * 3 + 1] - idx2 = rev_bases[i * 3 + 2] - frames[3, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + seq = seq[:crop_size] + length = len(seq) + arr = np.full((2, crop_size), pad_int, dtype=np.int32) - idx0 = rev_bases[i * 3 + 1] - idx1 = rev_bases[i * 3 + 2] - idx2 = rev_bases[i * 3 + 3] - frames[4, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + for i, ch in enumerate(seq): + arr[0, i] = _base_index(ch) + 1 - idx0 = rev_bases[i * 3 + 2] - idx1 = rev_bases[i * 3 + 3] - idx2 = rev_bases[i * 3 + 4] - frames[5, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + rev = _reverse_complement(seq) + for i, ch in enumerate(rev): + arr[1, i] = _base_index(ch) + 1 - # +1 for embedding layer - for i in range(6): - for j in range(min_len): - frames[i, j] = frames[i, j] + 1 + return arr, length - return frames +@njit(cache=False) +def _encode_nucleotide_onehot_numba( + seq_bytes: np.ndarray, crop_size: int +) -> Tuple[np.ndarray, int]: + """Numba kernel for 2-strand one-hot nucleotide encoding.""" + length = min(seq_bytes.shape[0], crop_size) + arr = np.zeros((2, crop_size, 4), dtype=np.float32) + ascii_lut = _ASCII_TO_BASE_IDX + comp_lut = _COMPLEMENT_LUT -def _process_sequence_full( - seq: str, crop_size: int, ngram_width: int = 3 -) -> np.ndarray: - """Process a single DNA sequence to 6-frame codon embeddings (pure Python).""" - codon_to_id = _get_codon_to_id() - seq = seq[:crop_size].upper() - length = len(seq) - offset = [-2, -1, 0][length % ngram_width] + for i in range(length): + idx = ascii_lut[seq_bytes[i]] + if idx < 4: + arr[0, i, idx] = 1.0 - ngrams_len = length - ngram_width + 1 - ngrams = [seq[i : i + ngram_width] for i in range(ngrams_len)] + for i in range(length): + idx = comp_lut[ascii_lut[seq_bytes[length - 1 - i]]] + if idx < 4: + arr[1, i, idx] = 1.0 - end_f1 = ngrams_len - 3 + offset - end_f2 = ngrams_len - 2 + offset - end_f3 = ngrams_len - 1 + offset + return arr, length - f1 = [codon_to_id.get(ngrams[i], -1) for i in range(0, end_f1, ngram_width)] - f2 = [codon_to_id.get(ngrams[i], -1) for i in range(1, end_f2, ngram_width)] - f3 = [codon_to_id.get(ngrams[i], -1) for i in range(2, end_f3, ngram_width)] - rev_comp = "".join(_COMP.get(b, "N") for b in reversed(seq)) - rev_ngrams_len = len(rev_comp) - ngram_width + 1 - rev_ngrams = [rev_comp[i : i + ngram_width] for i in range(rev_ngrams_len)] +def _encode_nucleotide_onehot(seq: str, crop_size: int) -> Tuple[np.ndarray, int]: + """Encode a DNA sequence as a one-hot nucleotide tensor (2 strands).""" + if HAS_NUMBA: + seq_bytes = np.frombuffer(seq[:crop_size].encode("ascii"), dtype=np.uint8) + return _encode_nucleotide_onehot_numba(seq_bytes, crop_size) - end_r1 = rev_ngrams_len - 3 + offset - end_r2 = rev_ngrams_len - 2 + offset - end_r3 = rev_ngrams_len - 1 + offset + seq = seq[:crop_size] + length = len(seq) + arr = np.zeros((2, crop_size, 4), dtype=np.float32) - r1 = [codon_to_id.get(rev_ngrams[i], -1) for i in range(0, end_r1, ngram_width)] - r2 = [codon_to_id.get(rev_ngrams[i], -1) for i in range(1, end_r2, ngram_width)] - r3 = [codon_to_id.get(rev_ngrams[i], -1) for i in range(2, end_r3, ngram_width)] + for i, ch in enumerate(seq): + idx = _base_index(ch) + if idx < 4: + arr[0, i, idx] = 1.0 - min_len = min(len(f1), len(f2), len(f3), len(r1), len(r2), len(r3)) - frames = np.array( - [ - f1[:min_len], - f2[:min_len], - f3[:min_len], - r1[:min_len], - r2[:min_len], - r3[:min_len], - ], - dtype=np.int32, - ) - frames = frames + 1 - return frames + rev = _reverse_complement(seq) + for i, ch in enumerate(rev): + idx = _base_index(ch) + if idx < 4: + arr[1, i, idx] = 1.0 + return arr, length -@njit(cache=False) -def _process_batch_numba( - sequences, lengths, crop_size, seq_len, codon_lut, comp_lut, ascii_lut -): - """Process a batch of DNA sequences to 6-frame codon embeddings.""" - n = len(lengths) - out = np.zeros((n, 6, seq_len), dtype=np.int32) - for s in range(n): - length = min(lengths[s], crop_size) +# --------------------------------------------------------------------------- +# Translated (codon) encoding +# --------------------------------------------------------------------------- +@njit(cache=False) +def _translate_sequence_integer_numba( + seq_bytes: np.ndarray, crop_size: int, codon_lut: np.ndarray, pad_int: int +) -> Tuple[np.ndarray, int]: + """Numba kernel for 6-frame codon indices using *codon_lut*.""" + length = min(seq_bytes.shape[0], crop_size) + ngram_width = 3 + ngrams_len = max(0, length - ngram_width + 1) + offset = [-2, -1, 0][length % ngram_width] + ascii_lut = _ASCII_TO_BASE_IDX + comp_lut = _COMPLEMENT_LUT - # Convert ASCII to base indices - bases = np.empty(length, dtype=np.int8) - for i in range(length): - bases[i] = ascii_lut[sequences[s, i]] + end_f1 = ngrams_len - 3 + offset + end_f2 = ngrams_len - 2 + offset + end_f3 = ngrams_len - 1 + offset - # Compute frame lengths - offset = [-2, -1, 0][length % 3] - ngrams_len = length - 3 + 1 + nf1 = max(0, (end_f1 + 2) // ngram_width) + nf2 = max(0, (end_f2 + 1) // ngram_width) + nf3 = max(0, end_f3 // ngram_width) - end_f1 = ngrams_len - 3 + offset - end_f2 = ngrams_len - 2 + offset - end_f3 = ngrams_len - 1 + offset + # Reverse complement bytes + rev_len = length + rev_ngrams_len = max(0, rev_len - ngram_width + 1) + end_r1 = rev_ngrams_len - 3 + offset + end_r2 = rev_ngrams_len - 2 + offset + end_r3 = rev_ngrams_len - 1 + offset - nf1 = max(0, (end_f1 + 2) // 3) - nf2 = max(0, (end_f2 + 1) // 3) - nf3 = max(0, end_f3 // 3) + nr1 = max(0, (end_r1 + 2) // ngram_width) + nr2 = max(0, (end_r2 + 1) // ngram_width) + nr3 = max(0, end_r3 // ngram_width) - # Reverse complement - rev_bases = np.empty(length, dtype=np.int8) - for i in range(length): - rev_bases[i] = comp_lut[bases[length - 1 - i]] + min_len = min(nf1, nf2, nf3, nr1, nr2, nr3) + max_codon_len = max(0, crop_size // 3 - 1) + frames = np.full((6, max_codon_len), pad_int, dtype=np.int32) - nr1 = max(0, (end_f1 + 2) // 3) - nr2 = max(0, (end_f2 + 1) // 3) - nr3 = max(0, end_f3 // 3) + if min_len <= 0: + return frames, 0 - min_len = min(nf1, nf2, nf3, nr1, nr2, nr3) - if min_len <= 0: - continue - min_len = min(min_len, seq_len) - - for i in range(min_len): - idx0 = bases[i * 3] - idx1 = bases[i * 3 + 1] - idx2 = bases[i * 3 + 2] - out[s, 0, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + 1 - - idx0 = bases[i * 3 + 1] - idx1 = bases[i * 3 + 2] - idx2 = bases[i * 3 + 3] - out[s, 1, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + 1 - - idx0 = bases[i * 3 + 2] - idx1 = bases[i * 3 + 3] - idx2 = bases[i * 3 + 4] - out[s, 2, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + 1 - - for i in range(min_len): - idx0 = rev_bases[i * 3] - idx1 = rev_bases[i * 3 + 1] - idx2 = rev_bases[i * 3 + 2] - out[s, 3, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + 1 - - idx0 = rev_bases[i * 3 + 1] - idx1 = rev_bases[i * 3 + 2] - idx2 = rev_bases[i * 3 + 3] - out[s, 4, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + 1 - - idx0 = rev_bases[i * 3 + 2] - idx1 = rev_bases[i * 3 + 3] - idx2 = rev_bases[i * 3 + 4] - out[s, 5, i] = codon_lut[idx0 * 25 + idx1 * 5 + idx2] + 1 + for i in range(min_len): + # Forward frame 1 + b0 = ascii_lut[seq_bytes[i * 3]] + b1 = ascii_lut[seq_bytes[i * 3 + 1]] + b2 = ascii_lut[seq_bytes[i * 3 + 2]] + cid = codon_lut[b0 * 25 + b1 * 5 + b2] + frames[0, i] = pad_int if cid < 0 else cid + 1 + + # Forward frame 2 + b0 = ascii_lut[seq_bytes[i * 3 + 1]] + b1 = ascii_lut[seq_bytes[i * 3 + 2]] + b2 = ascii_lut[seq_bytes[i * 3 + 3]] + cid = codon_lut[b0 * 25 + b1 * 5 + b2] + frames[1, i] = pad_int if cid < 0 else cid + 1 + + # Forward frame 3 + b0 = ascii_lut[seq_bytes[i * 3 + 2]] + b1 = ascii_lut[seq_bytes[i * 3 + 3]] + b2 = ascii_lut[seq_bytes[i * 3 + 4]] + cid = codon_lut[b0 * 25 + b1 * 5 + b2] + frames[2, i] = pad_int if cid < 0 else cid + 1 + + # Reverse frames use reversed sequence + ri = rev_len - 1 - i * 3 + b0 = comp_lut[ascii_lut[seq_bytes[ri]]] + b1 = comp_lut[ascii_lut[seq_bytes[ri - 1]]] + b2 = comp_lut[ascii_lut[seq_bytes[ri - 2]]] + cid = codon_lut[b0 * 25 + b1 * 5 + b2] + frames[3, i] = pad_int if cid < 0 else cid + 1 + + b0 = comp_lut[ascii_lut[seq_bytes[ri - 1]]] + b1 = comp_lut[ascii_lut[seq_bytes[ri - 2]]] + b2 = comp_lut[ascii_lut[seq_bytes[ri - 3]]] + cid = codon_lut[b0 * 25 + b1 * 5 + b2] + frames[4, i] = pad_int if cid < 0 else cid + 1 + + b0 = comp_lut[ascii_lut[seq_bytes[ri - 2]]] + b1 = comp_lut[ascii_lut[seq_bytes[ri - 3]]] + b2 = comp_lut[ascii_lut[seq_bytes[ri - 4]]] + cid = codon_lut[b0 * 25 + b1 * 5 + b2] + frames[5, i] = pad_int if cid < 0 else cid + 1 + + return frames, min_len + + +def _translate_sequence_integer( + seq: str, crop_size: int, codon_map: np.ndarray, pad_int: int, codon_lut: np.ndarray | None = None +) -> Tuple[np.ndarray, int]: + """Encode a DNA sequence as 6-frame codon indices using *codon_map*.""" + if HAS_NUMBA: + seq_bytes = np.frombuffer( + seq[:crop_size].upper().encode("ascii"), dtype=np.uint8 + ) + if codon_lut is None: + codon_lut = _build_codon_lut(codon_map) + return _translate_sequence_integer_numba( + seq_bytes, crop_size, codon_lut, pad_int + ) - return out + seq = seq[:crop_size].upper() + length = len(seq) + codon_to_id = _get_codon_to_id() + ngram_width = 3 + ngrams_len = max(0, length - ngram_width + 1) + ngrams = [seq[i : i + ngram_width] for i in range(ngrams_len)] + offset = [-2, -1, 0][length % ngram_width] -def _process_chunk_numpy_full( - lines: list[str], crop_size: int, num_classes: int -) -> tuple: - """Process a chunk of CSV lines to fully preprocessed tensors.""" - n = len(lines) - seq_len = crop_size // 3 - 1 - sequences = np.zeros((n, 6, seq_len), dtype=np.int32) - labels = np.zeros((n, num_classes), dtype=np.float32) + end_f1 = ngrams_len - 3 + offset + end_f2 = ngrams_len - 2 + offset + end_f3 = ngrams_len - 1 + offset - for i, line in enumerate(lines): - line = line.strip() - if not line: - continue - comma_idx = line.find(",") - if comma_idx == -1: - continue - label = int(line[:comma_idx]) - seq = line[comma_idx + 1 :] - frames = _process_sequence_full(seq, crop_size) - actual_len = frames.shape[1] - sequences[i, :, :actual_len] = frames - labels[i, label] = 1.0 + f1 = [ngrams[i] for i in range(0, end_f1, ngram_width)] + f2 = [ngrams[i] for i in range(1, end_f2, ngram_width)] + f3 = [ngrams[i] for i in range(2, end_f3, ngram_width)] - return sequences, labels + rev_comp = _reverse_complement(seq) + rev_ngrams_len = max(0, len(rev_comp) - ngram_width + 1) + rev_ngrams = [rev_comp[i : i + ngram_width] for i in range(rev_ngrams_len)] + end_r1 = rev_ngrams_len - 3 + offset + end_r2 = rev_ngrams_len - 2 + offset + end_r3 = rev_ngrams_len - 1 + offset -def _convert_to_numpy_full( - csv_path: str, - output_path: str, - crop_size: int, - num_classes: int, - num_workers: int | None, -): - """Convert CSV to fully preprocessed NumPy format.""" - total = sum(1 for _ in open(csv_path)) - print(f"Total samples: {total}") + r1 = [rev_ngrams[i] for i in range(0, end_r1, ngram_width)] + r2 = [rev_ngrams[i] for i in range(1, end_r2, ngram_width)] + r3 = [rev_ngrams[i] for i in range(2, end_r3, ngram_width)] - if num_workers is None: - num_workers = cpu_count() - print(f"Using {num_workers} workers") + min_len = min(len(f1), len(f2), len(f3), len(r1), len(r2), len(r3)) + max_codon_len = max(0, crop_size // 3 - 1) + frames = np.full((6, max_codon_len), pad_int, dtype=np.int32) - # Read and parse CSV - with open(csv_path) as f: - lines = f.readlines() + for frame_idx, codons in enumerate([f1, f2, f3, r1, r2, r3]): + for i, codon in enumerate(codons[:min_len]): + cid = codon_to_id.get(codon, -1) + if 0 <= cid < len(codon_map): + frames[frame_idx, i] = int(codon_map[cid]) + 1 + else: + frames[frame_idx, i] = pad_int - start = time.time() + return frames, min_len - if HAS_NUMBA and total >= 1000: - # Fast path: batch numba processing (no multiprocessing overhead) - print("Using numba-optimized batch processing") - codon_lut, ascii_lut, comp_lut = _build_numba_lookups() - seq_len = crop_size // 3 - 1 - - # Parse labels and encode sequences as uint8 array - labels = np.zeros((total, num_classes), dtype=np.float32) - sequences = np.full((total, crop_size), ord("N"), dtype=np.uint8) - lengths = np.zeros(total, dtype=np.int32) - - for i, line in enumerate(lines): - line = line.strip() - if not line: - continue - comma_idx = line.find(",") - if comma_idx == -1: - continue - label = int(line[:comma_idx]) - seq = line[comma_idx + 1 :] - seq_bytes = seq.encode() - length = min(len(seq_bytes), crop_size) - sequences[i, :length] = np.frombuffer(seq_bytes[:length], dtype=np.uint8) - lengths[i] = length - labels[i, label] = 1.0 - - # Process entire batch in numba - all_sequences = _process_batch_numba( - sequences, lengths, crop_size, seq_len, codon_lut, comp_lut, ascii_lut - ) - all_labels = labels - else: - # Fallback: multiprocessing with pure Python - if HAS_NUMBA: - print("Numba available but dataset too small for batch mode") - else: - print("Numba not available, using pure Python") - - chunk_size = max(1, len(lines) // num_workers) - chunks = [lines[i : i + chunk_size] for i in range(0, len(lines), chunk_size)] - process_fn = partial( - _process_chunk_numpy_full, crop_size=crop_size, num_classes=num_classes - ) - with Pool(num_workers) as pool: - results = pool.map(process_fn, chunks) - all_sequences = np.concatenate([r[0] for r in results], axis=0) - all_labels = np.concatenate([r[1] for r in results], axis=0) - elapsed = time.time() - start - print( - f"Processed {total} samples in {elapsed:.2f}s ({total / elapsed:.0f} samples/sec)" - ) +def _to_one_hot(arr: np.ndarray, vocab_size: int, pad_int: int) -> np.ndarray: + """Convert an integer array to a one-hot float tensor. - np.savez_compressed(output_path, translated=all_sequences, label=all_labels) - seq_mb = all_sequences.nbytes / (1024 * 1024) - label_mb = all_labels.nbytes / (1024 * 1024) - print(f"Saved: translated={seq_mb:.1f}MB, label={label_mb:.1f}MB") + Padding positions (``arr == pad_int``) remain all-zero. + """ + one_hot = np.zeros(arr.shape + (vocab_size,), dtype=np.float32) + mask = arr != pad_int + if mask.any(): + one_hot[mask] = np.eye(vocab_size)[arr[mask] - 1] + return one_hot # --------------------------------------------------------------------------- -# numpy_raw_variable (variable-length int8) +# Chunk processing # --------------------------------------------------------------------------- +def _parse_csv_line(line: str) -> Tuple[int, str]: + """Parse a 'label,sequence[,id]' CSV line.""" + parts = line.strip().split(",", 2) + if len(parts) < 2: + raise ValueError(f"Invalid CSV line (no comma): {line!r}") + return int(parts[0]), parts[1] -def _encode_variable_length(line: str, max_length: int) -> tuple: - """Encode a single CSV line to variable-length int8 sequence.""" - line = line.strip() - if not line: - return None, None, None - comma_idx = line.find(",") - if comma_idx == -1: - return None, None, None - label = int(line[:comma_idx]) - seq = line[comma_idx + 1 :] - length = min(len(seq), max_length) - - arr = np.full(max_length, 4, dtype=np.int8) - for i in range(length): - c = seq[i] - if c in _BASE_MAP: - arr[i] = _BASE_MAP[c] - elif c in _BASE_MAP_LOWER: - arr[i] = _BASE_MAP_LOWER[c] - else: - arr[i] = 4 - return arr, length, label - - -def _process_chunk_numpy_raw_variable( - lines: list[str], max_length: int, num_classes: int -) -> tuple: - """Process a chunk of CSV lines with variable lengths.""" - n = len(lines) - sequences = np.full((n, max_length), 4, dtype=np.int8) - lengths = np.zeros(n, dtype=np.int32) - labels = np.zeros((n, num_classes), dtype=np.float32) - - for i, line in enumerate(lines): - seq_arr, length, label = _encode_variable_length(line, max_length) - if seq_arr is None: - continue - sequences[i] = seq_arr - lengths[i] = length - labels[i, label] = 1.0 - - return sequences, lengths, labels - +def _sliding_windows(seq: str, crop_size: int, stride: int): + """Yield windows from *seq* using a sliding window of *crop_size*. -def _convert_to_numpy_raw_variable( - csv_path: str, - output_path: str, - max_length: int, - num_classes: int, - num_workers: int | None, -): - """Convert CSV to NumPy with variable-length sequences.""" - total = sum(1 for _ in open(csv_path)) - print(f"Total samples: {total}") - - if num_workers is None: - num_workers = cpu_count() - print(f"Using {num_workers} workers") + If *stride* is 0 (or negative), a single truncated window is returned. + Sequences shorter than *crop_size* yield one padded window. + """ + if stride <= 0: + yield seq[:crop_size] + return - with open(csv_path) as f: - lines = f.readlines() + length = len(seq) + if length <= crop_size: + yield seq + return - chunk_size = max(1, len(lines) // num_workers) - chunks = [lines[i : i + chunk_size] for i in range(0, len(lines), chunk_size)] - print(f"Split into {len(chunks)} chunks") + for start in range(0, length, stride): + window = seq[start : start + crop_size] + if not window: + break + yield window - start = time.time() - process_fn = partial( - _process_chunk_numpy_raw_variable, - max_length=max_length, - num_classes=num_classes, - ) - with Pool(num_workers) as pool: - results = pool.map(process_fn, chunks) +def _process_chunk( + lines: List[str], + crop_size: int, + stride: int, + num_classes: int, + codon_map: np.ndarray, + format: str, + one_hot: bool, + pad_int: int, +) -> Dict[str, List[np.ndarray]]: + """Process a chunk of CSV lines and return per-sample arrays.""" + out: Dict[str, List[np.ndarray]] = { + "nucleotide": [], + "translated": [], + "labels": [], + "lengths": [], + } + + codon_lut = None + if format in ("translated", "both"): + codon_lut = _build_codon_lut(codon_map) + + for line in lines: + if not line.strip(): + continue + label, seq = _parse_csv_line(line) + label_vec = np.zeros(num_classes, dtype=np.float32) + label_vec[label] = 1.0 + + for window in _sliding_windows(seq, crop_size, stride): + if format in ("nucleotide", "both"): + if one_hot: + nuc, nuc_len = _encode_nucleotide_onehot(window, crop_size) + else: + nuc, nuc_len = _encode_nucleotide_integer( + window, crop_size, pad_int + ) + out["nucleotide"].append(nuc) + + if format in ("translated", "both"): + trans, trans_len = _translate_sequence_integer( + window, crop_size, codon_map, pad_int, codon_lut=codon_lut + ) + if one_hot: + trans = _to_one_hot(trans, int(codon_map.max()) + 1, pad_int) + out["translated"].append(trans) + + out["labels"].append(label_vec) + # Store nucleotide length when available; otherwise codon length. + out["lengths"].append( + nuc_len if format in ("nucleotide", "both") else trans_len + ) - all_sequences = np.concatenate([r[0] for r in results], axis=0) - all_lengths = np.concatenate([r[1] for r in results], axis=0) - all_labels = np.concatenate([r[2] for r in results], axis=0) + return out - elapsed = time.time() - start - print( - f"Processed {total} samples in {elapsed:.2f}s ({total / elapsed:.0f} samples/sec)" - ) - np.savez_compressed( - output_path, - sequences=all_sequences, - lengths=all_lengths, - labels=all_labels, - ) - seq_mb = all_sequences.nbytes / (1024 * 1024) - length_mb = all_lengths.nbytes / (1024 * 1024) - label_mb = all_labels.nbytes / (1024 * 1024) - print( - f"Saved: sequences={seq_mb:.1f}MB, lengths={length_mb:.1f}MB, labels={label_mb:.1f}MB" - ) +def _trim_to_max(arrays: List[np.ndarray], lengths: List[int]) -> np.ndarray: + """Stack a list of padded arrays and trim to the actual maximum length.""" + stacked = np.stack(arrays, axis=0) + max_len = max(lengths) if lengths else 0 + if max_len == 0: + return stacked + # Time/length is the second-to-last axis for 1-D feature arrays and the + # second axis for nucleotide/translated arrays. Trim all axes after batch + # and frame/strand dims to max_len. + if stacked.ndim == 3: + return stacked[:, :, :max_len] + if stacked.ndim == 4: + return stacked[:, :, :max_len, :] + return stacked # --------------------------------------------------------------------------- # Public dispatcher # --------------------------------------------------------------------------- - - def convert_dataset( input_path: str, output_path: str, format: str, crop_size: int = 500, + stride: int = 0, num_classes: int = 3, - use_embedding_layer: bool = True, - max_length: int = 1000, num_workers: int | None = None, + one_hot: bool = False, + pad_int: int = 0, + codon_map: str = "CODON_ID", + max_length: int = 5000, + compress: str = "fast", + _available_memory_bytes: int | None = None, ): """Convert a CSV dataset to an optimized training format. Args: input_path: Path to input CSV file (label,sequence format). output_path: Path to output file. - format: Target format — ``numpy_raw``, ``numpy_full``, - or ``numpy_raw_variable``. - crop_size: Sequence crop size (for fixed-length formats). + format: Target representation — ``nucleotide``, ``translated``, or ``both``. + crop_size: Maximum sequence length to process per window. + stride: Sliding-window step in nucleotides. 0 = one window per sequence. num_classes: Number of output classes. - use_embedding_layer: Unused (kept for API compatibility). - max_length: Maximum sequence length (variable-length only). num_workers: Number of parallel workers (``None`` = auto). + one_hot: Output float one-hot tensors instead of integer indices. + pad_int: Integer value used for padding and stored as ``pad_int``. + codon_map: Name of a length-64 codon map in ``jaeger.seqops.maps``. + max_length: Deprecated; kept for compatibility. + compress: NPZ compression level — ``fast`` (zlib 1), ``default`` (zlib 6), + or ``none``. + _available_memory_bytes: Internal override for memory guard tests; do not + use from production code. + + Raises: + MemoryError: If ``one_hot=True`` and the estimated raw float32 tensor + exceeds the available RAM safety margin. """ - valid_formats = ["numpy_raw", "numpy_full", "numpy_raw_variable"] + del max_length # no longer used + + format = format.lower() + total_lines = _count_lines(input_path) + valid_formats = ["nucleotide", "translated", "both"] if format not in valid_formats: raise ValueError( f"Invalid format: {format}. Choose from: {', '.join(valid_formats)}" ) - print(f"Converting {input_path} -> {output_path}") - print(f"Format: {format}, Crop size: {crop_size}, Num classes: {num_classes}") - - if format == "numpy_raw": - _convert_to_numpy_raw( - input_path, output_path, crop_size, num_classes, num_workers + compress = compress.lower() + valid_compress = ["fast", "default", "none"] + if compress not in valid_compress: + raise ValueError( + f"Invalid compress: {compress}. Choose from: {', '.join(valid_compress)}" ) - elif format == "numpy_full": - _convert_to_numpy_full( - input_path, output_path, crop_size, num_classes, num_workers + + codon_map_arr = _get_codon_map(codon_map) if format in ("translated", "both") else None + + # Guard against one-hot outputs that cannot fit in RAM. + if one_hot: + estimated = _estimate_onehot_memory( + total_rows=total_lines, + crop_size=crop_size, + format=format, + one_hot=one_hot, + codon_map_arr=codon_map_arr, ) - elif format == "numpy_raw_variable": - _convert_to_numpy_raw_variable( - input_path, output_path, max_length, num_classes, num_workers + available = ( + _available_memory_bytes + if _available_memory_bytes is not None + else psutil.virtual_memory().available ) + _check_onehot_memory(estimated, available) + + # Validate pad_int does not collide with actual token indices. + if not one_hot and format in ("nucleotide", "both"): + actual_nuc = {1, 2, 3, 4, 5} # A,T,G,C,N after +1 offset + if pad_int in actual_nuc: + raise ValueError( + f"pad_int ({pad_int}) collides with a nucleotide token index " + f"({sorted(actual_nuc)})" + ) + if not one_hot and format in ("translated", "both") and codon_map_arr is not None: + actual_trans = set(int(v) + 1 for v in codon_map_arr) + if pad_int in actual_trans: + raise ValueError( + f"pad_int ({pad_int}) collides with a translated token index " + f"for codon map {codon_map}" + ) + + print(f"Converting {input_path} -> {output_path}") + print( + f"Format: {format}, one_hot: {one_hot}, pad_int: {pad_int}, " + f"codon_map: {codon_map}, crop_size: {crop_size}, stride: {stride}, " + f"num_classes: {num_classes}, compress: {compress}" + ) + + with open(input_path) as f: + lines = f.readlines() + + if num_workers is None: + num_workers = cpu_count() + print(f"Total CSV rows: {total_lines}, using {num_workers} workers") + + chunk_size = max(1, len(lines) // num_workers) + chunks = [lines[i : i + chunk_size] for i in range(0, len(lines), chunk_size)] + + start = time.time() + process_fn = partial( + _process_chunk, + crop_size=crop_size, + stride=stride, + num_classes=num_classes, + codon_map=codon_map_arr if codon_map_arr is not None else np.arange(64), + format=format, + one_hot=one_hot, + pad_int=pad_int, + ) + + if num_workers == 1: + results = [process_fn(chunk) for chunk in chunks] + elif HAS_NUMBA: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=num_workers) as executor: + results = list(executor.map(process_fn, chunks)) + else: + with _MP_CONTEXT.Pool(num_workers) as pool: + results = pool.map(process_fn, chunks) + + # Merge chunks + nuc_arrays = [a for r in results for a in r["nucleotide"]] + trans_arrays = [a for r in results for a in r["translated"]] + labels = np.stack([a for r in results for a in r["labels"]], axis=0) + lengths = np.array([a for r in results for a in r["lengths"]], dtype=np.int32) + + to_save: Dict[str, np.ndarray] = {"label": labels, "lengths": lengths, "pad_int": np.array(pad_int)} + + if format in ("nucleotide", "both"): + to_save["nucleotide"] = _trim_to_max(nuc_arrays, lengths.tolist()) + if format in ("translated", "both") and codon_map_arr is not None: + to_save["translated"] = _trim_to_max(trans_arrays, [t.shape[1] for t in trans_arrays]) + to_save["vocab_size"] = np.array(int(codon_map_arr.max()) + 1) + + _save_npz(output_path, to_save, compress) + + total_windows = len(labels) + elapsed = time.time() - start + print( + f"Processed {total_lines} CSV rows into {total_windows} windows in " + f"{elapsed:.2f}s ({total_windows / elapsed:.0f} windows/sec)" + ) + for key, arr in to_save.items(): + print(f" {key}: shape={arr.shape}, dtype={arr.dtype}") diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py index 69e7330..cef84b1 100644 --- a/src/jaeger/nnlib/pytorch/builder.py +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional import torch import torch.nn as nn @@ -12,6 +12,51 @@ ) +def _resolve_vocab_size( + embedding_cfg: Dict[str, Any], string_cfg: Dict[str, Any] +) -> int: + """Infer vocab_size from config, falling back to the codon table length.""" + vocab_size = embedding_cfg.get("vocab_size") + if vocab_size is not None: + return int(vocab_size) + + codon_id = string_cfg.get("codon_id") + table_len = None + if isinstance(codon_id, (list, tuple)): + table_len = len(codon_id) + elif isinstance(codon_id, str): + import jaeger.seqops.maps as maps + + table = getattr(maps, codon_id, None) + if isinstance(table, (list, tuple)): + table_len = len(table) + if table_len is not None: + # Index 0 is reserved for padding/unknown, so the embedding size is + # one larger than the raw codon table length. + return table_len + 1 + # Standard 6-frame codon vocabulary size (including padding/unknown index 0) + return 65 + + +def _extract_hidden_units(hidden_layers: List[Dict[str, Any]]) -> List[int]: + """Return intermediate dense unit sizes from a Keras-style layer list.""" + units: List[int] = [] + # The final layer produces the output classes, so exclude it. + for layer in hidden_layers[:-1]: + cfg = layer.get("config", {}) + if "units" in cfg: + units.append(int(cfg["units"])) + return units + + +def _extract_dropout_rate(hidden_layers: List[Dict[str, Any]]) -> float: + """Return the first explicit dropout rate found in the layer list, if any.""" + for layer in hidden_layers: + if layer.get("name", "").lower() == "dropout": + return float(layer.get("config", {}).get("rate", 0.0)) + return 0.0 + + class _ClassifierPipeline(nn.Module): """Runs the representation model and classification head, propagating masks.""" @@ -48,9 +93,10 @@ def build_fragment_classifier(self) -> Dict[str, nn.Module]: models: Dict[str, nn.Module] = {} embedding_cfg = self.model_cfg.get("embedding", {}) + string_cfg = self.model_cfg.get("string_processor", {}) embedding = Embedding( input_type=embedding_cfg.get("input_type"), - vocab_size=embedding_cfg.get("vocab_size"), + vocab_size=_resolve_vocab_size(embedding_cfg, string_cfg), embedding_size=embedding_cfg.get("embedding_size", 4), use_embedding_layer=embedding_cfg.get("use_embedding_layer", False), ) @@ -64,28 +110,30 @@ def build_fragment_classifier(self) -> Dict[str, nn.Module]: models["rep_model"] = rep_model cls_cfg = self.model_cfg["classifier"] + input_dim = cls_cfg.get("input_shape", rep_model.output_dim) + if isinstance(input_dim, (list, tuple)) and len(input_dim) == 1: + input_dim = input_dim[0] head = ClassificationHead( - input_dim=cls_cfg.get("input_shape", rep_model.output_dim), + input_dim=input_dim, num_classes=self.classifier_out_dim, - hidden_units=[ - layer["config"]["units"] - for layer in cls_cfg.get("hidden_layers", [])[:-1] - ], - dropout=cls_cfg.get("dropout", 0.0), + hidden_units=_extract_hidden_units(cls_cfg.get("hidden_layers", [])), + dropout=_extract_dropout_rate(cls_cfg.get("hidden_layers", [])) + or cls_cfg.get("dropout", 0.0), ) models["classification_head"] = head models["jaeger_classifier"] = _ClassifierPipeline(rep_model, head) if "reliability_model" in self.model_cfg: rel_cfg = self.model_cfg["reliability_model"] + rel_input_dim = rel_cfg.get("input_shape", rep_model.nmd_dim) + if isinstance(rel_input_dim, (list, tuple)) and len(rel_input_dim) == 1: + rel_input_dim = rel_input_dim[0] rel_head = ReliabilityHead( - input_dim=rel_cfg.get("input_shape", rep_model.nmd_dim), + input_dim=rel_input_dim, num_classes=self.reliability_out_dim, - hidden_units=[ - layer["config"]["units"] - for layer in rel_cfg.get("hidden_layers", [])[:-1] - ], - dropout=rel_cfg.get("dropout", 0.0), + hidden_units=_extract_hidden_units(rel_cfg.get("hidden_layers", [])), + dropout=_extract_dropout_rate(rel_cfg.get("hidden_layers", [])) + or rel_cfg.get("dropout", 0.0), ) models["reliability_head"] = rel_head @@ -97,7 +145,10 @@ def build_fragment_classifier(self) -> Dict[str, nn.Module]: return models def compile_model( - self, models: Dict[str, nn.Module], train_branch: str = "classifier" + self, + models: Dict[str, nn.Module], + train_branch: str = "classifier", + class_weights: Optional[torch.Tensor] = None, ) -> tuple: opt_name = self.train_cfg.get("optimizer", "adam").lower() opt_params = self.train_cfg.get("optimizer_params", {}) @@ -106,7 +157,7 @@ def compile_model( if train_branch == "classifier": model = models["jaeger_classifier"] optimizer = opt_class(model.parameters(), **opt_params) - loss = nn.CrossEntropyLoss() + loss = nn.CrossEntropyLoss(weight=class_weights) return model, optimizer, loss if train_branch == "reliability": raise NotImplementedError( diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 537dca4..353f432 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -1,5 +1,5 @@ import math -from typing import Callable, Optional, Tuple +from typing import Callable, List, Optional, Tuple import torch import torch.nn as nn @@ -48,7 +48,10 @@ def __init__( use_bias: bool = True, kernel_initializer: str = "glorot_uniform", bias_initializer: str = "zeros", - kernel_regularizer: Optional[float] = None, + kernel_regularizer: Optional[str] = None, + kernel_regularizer_w: float = 0.0, + bias_regularizer: Optional[str] = None, + **kwargs, ): super().__init__() padding_norm = padding.lower() @@ -65,10 +68,12 @@ def __init__( f"bias_initializer {bias_initializer!r} is not supported; " "only 'zeros' is implemented." ) - if kernel_regularizer is not None: - raise NotImplementedError( - f"kernel_regularizer {kernel_regularizer!r} is not supported." - ) + + # Regularizers are accepted for Keras config compatibility but are not + # implemented at the layer level; use optimizer weight_decay instead. + self.kernel_regularizer = kernel_regularizer + self.kernel_regularizer_w = kernel_regularizer_w + self.bias_regularizer = bias_regularizer self.filters = filters self.kernel_size = kernel_size @@ -322,17 +327,38 @@ def forward( class AxialAttention(nn.Module): - """Axial attention over the sequence axis.""" + """Axial attention over the sequence axis. - def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0): + Supports Keras-style configs with ``feed_forward_dim`` (defaults to + ``4 * embed_dim``), ``num_blocks`` (defaults to 1), and ``dropout_rate``. + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + feed_forward_dim: Optional[int] = None, + num_blocks: int = 1, + dropout_rate: Optional[float] = None, + ): super().__init__() - self.attn = nn.MultiheadAttention( - embed_dim=embed_dim, - num_heads=num_heads, - dropout=dropout, - batch_first=True, - ) - self.norm = nn.LayerNorm(embed_dim) + self.embed_dim = embed_dim + dropout = dropout_rate if dropout_rate is not None else dropout + feed_forward_dim = feed_forward_dim or embed_dim * 4 + layers = [] + for _ in range(num_blocks): + layers.append( + nn.TransformerEncoderLayer( + d_model=embed_dim, + nhead=num_heads, + dim_feedforward=feed_forward_dim, + dropout=dropout, + batch_first=True, + norm_first=False, + ) + ) + self.blocks = nn.ModuleList(layers) def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None @@ -343,12 +369,10 @@ def forward( key_padding_mask = None if mask is not None: key_padding_mask = ~mask.reshape(b * f, seq_len) - attn_out, _ = self.attn( - x_2d, x_2d, x_2d, key_padding_mask=key_padding_mask, need_weights=False - ) - out = self.norm(x_2d + attn_out).reshape(b, f, seq_len, d) - out_mask = mask - return out, out_mask + out = x_2d + for block in self.blocks: + out = block(out, src_key_padding_mask=key_padding_mask) + return out.reshape(b, f, seq_len, d), mask class MaskedGlobalAvgPooling(nn.Module): @@ -390,22 +414,55 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ResidualBlock(nn.Module): - def __init__(self, layer: nn.Module): + """Residual wrapper around a layer or a small stack of MaskedConv1D layers. + + For Keras-style configs, ``config`` may contain ``block_size``, + ``filters``, ``kernel_size``, ``use_bias``, and regularizer keys to build + a stack of ``MaskedConv1D`` layers. If ``layer`` is provided directly, + it is wrapped as before. + """ + + def __init__(self, layer: Optional[nn.Module] = None, config: Optional[dict] = None): super().__init__() - self.layer = layer + if layer is not None: + self.layer = layer + elif config is not None: + cfg = dict(config) + block_size = cfg.pop("block_size", 1) + cfg.setdefault("padding", "same") + cfg.setdefault("activation", None) + layers: List[nn.Module] = [] + for _ in range(block_size): + layers.append(MaskedConv1D(**cfg)) + self.layer = nn.ModuleList(layers) + else: + raise ValueError("ResidualBlock requires either ``layer`` or ``config``.") def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): - if ( - hasattr(self.layer, "forward") - and "mask" in self.layer.forward.__code__.co_varnames - ): - out = self.layer(x, mask) + residual = x + out = x + out_mask = mask + if isinstance(self.layer, nn.ModuleList): + for layer in self.layer: + out, out_mask = layer(out, out_mask) + elif "mask" in self.layer.forward.__code__.co_varnames: + out = self.layer(out, mask) if isinstance(out, tuple): out, out_mask = out - return out + x, out_mask - return out + x, mask - out = self.layer(x) - return out + x, mask + else: + out = self.layer(out) + + # If dimensions differ, project the residual with a 1x1 conv. + if out.shape[-1] != residual.shape[-1]: + proj = nn.Conv1d( + residual.shape[-1], out.shape[-1], kernel_size=1, bias=False + ).to(out.device) + residual = proj( + residual.permute(0, 1, 3, 2).reshape(-1, residual.shape[-1], residual.shape[2]) + ) + residual = residual.permute(0, 2, 1).reshape_as(out) + + return out + residual, out_mask class TransformerEncoder(nn.Module): diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index f5224cb..ddc0633 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -32,6 +32,8 @@ def __init__( ): super().__init__() self.input_type = input_type + self.vocab_size = vocab_size + self.embedding_size = embedding_size self.use_embedding_layer = use_embedding_layer self.use_positional_embeddings = use_positional_embeddings @@ -43,7 +45,14 @@ def __init__( self.embed = nn.Linear(vocab_size, embedding_size, bias=False) nn.init.orthogonal_(self.embed.weight) elif input_type == "nucleotide": - self.embed = None + # Nucleotide input can be provided as a float one-hot tensor or as + # integer indices. Integer indices are one-hot encoded and then + # optionally projected to ``embedding_size``. + if use_embedding_layer: + self.embed = None + else: + self.embed = nn.Linear(vocab_size, embedding_size, bias=False) + nn.init.orthogonal_(self.embed.weight) else: raise ValueError(f"Invalid input_type: {input_type}") @@ -63,6 +72,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.embed.weight.dtype ) x = self.embed(x) + elif self.input_type == "nucleotide": + if not torch.is_floating_point(x): + x = F.one_hot(x, num_classes=self.vocab_size).to( + torch.float32 + ) + if self.embed is not None: + x = self.embed(x) if self.use_positional_embeddings: x = self.positional(x) return x @@ -145,29 +161,64 @@ def __init__( self.embedding = embedding self.blocks = nn.ModuleList() self.return_nmd = False + current_dim = self._embedding_output_dim(embedding) for cfg in hidden_layers: name = cfg["name"].lower() config = cfg.get("config", {}) if name == "masked_conv1d": - self.blocks.append(MaskedConv1D(**config)) + block = MaskedConv1D(**config) + current_dim = block.filters + self.blocks.append(block) elif name == "masked_batchnorm": - self.blocks.append(MaskedBatchNorm(**config)) + config = dict(config) + config.setdefault("num_features", current_dim) + block = MaskedBatchNorm(**config) + current_dim = block.num_features + self.blocks.append(block) elif name == "masked_layer_norm": - self.blocks.append(MaskedLayerNorm(**config)) + config = dict(config) + config.setdefault("num_features", current_dim) + block = MaskedLayerNorm(**config) + current_dim = block.num_features + self.blocks.append(block) elif name == "axial_attention": - self.blocks.append(AxialAttention(**config)) + config = dict(config) + config.setdefault("embed_dim", current_dim) + block = AxialAttention(**config) + current_dim = block.embed_dim + self.blocks.append(block) elif name == "cross_frame_attention": - self.blocks.append(CrossFrameAttention(**config)) + config = dict(config) + config.setdefault("embed_dim", current_dim) + block = CrossFrameAttention(**config) + current_dim = block.embed_dim + self.blocks.append(block) elif name == "transformer_encoder": - self.blocks.append(TransformerEncoder(**config)) + config = dict(config) + config.setdefault("embed_dim", current_dim) + block = TransformerEncoder(**config) + current_dim = self._block_output_dim(block) + self.blocks.append(block) elif name == "residual_block": - inner = self.blocks.pop(-1) if self.blocks else None - self.blocks.append(ResidualBlock(inner)) + # Config-driven residual block: ``config`` defines the internal + # stack (e.g. ``block_size`` MaskedConv1D layers). If an older + # style config passes the previous layer directly, wrap it. + if "layer" in config: + block = ResidualBlock(layer=config["layer"]) + else: + cfg_copy = dict(config) + cfg_copy.setdefault("filters", current_dim) + block = ResidualBlock(config=cfg_copy) + current_dim = cfg_copy.get("filters", current_dim) + self.blocks.append(block) elif name == "dense": + config = dict(config) if "units" in config: - config = dict(config) config["out_features"] = config.pop("units") - self.blocks.append(nn.Linear(**config)) + config.setdefault("in_features", current_dim) + block = nn.Linear(**config) + current_dim = block.out_features + self.blocks.append(block) elif name == "activation": self.blocks.append( GeLU() if config.get("activation") == "gelu" else nn.ReLU() @@ -179,6 +230,17 @@ def __init__( self.pooler = self._build_pooler(pooling) self.output_dim, self.nmd_dim = self._infer_dims() + @staticmethod + def _embedding_output_dim(embedding: nn.Module) -> int: + if isinstance(embedding, Embedding): + return int(embedding.embedding_size) + if hasattr(embedding, "embed") and embedding.embed is not None: + if hasattr(embedding.embed, "embedding_dim"): + return int(embedding.embed.embedding_dim) + if hasattr(embedding.embed, "out_features"): + return int(embedding.embed.out_features) + return 4 + def _infer_dims(self): output_dim = None for block in reversed(self.blocks): @@ -203,7 +265,7 @@ def _block_output_dim(block: nn.Module) -> Optional[int]: if isinstance(block, (MaskedBatchNorm, MaskedLayerNorm)): return int(block.num_features) if isinstance(block, (AxialAttention, CrossFrameAttention)): - return int(block.attn.embed_dim) + return int(block.embed_dim) if isinstance(block, TransformerEncoder): return int(block.encoder.layers[0].self_attn.embed_dim) if isinstance(block, nn.Linear): diff --git a/src/jaeger/training/pytorch/callbacks.py b/src/jaeger/training/pytorch/callbacks.py index bbf2bb8..7741cd8 100644 --- a/src/jaeger/training/pytorch/callbacks.py +++ b/src/jaeger/training/pytorch/callbacks.py @@ -90,24 +90,31 @@ def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): if improved: self.best_value = current if self.save_best_only: - self._save(trainer, epoch) + self._save(trainer, epoch, logs) if self.verbose: print( f"Epoch {epoch}: {self.monitor} improved to {current:.4f}; saving model" ) elif not self.save_best_only: - self._save(trainer, epoch) - - def _save(self, trainer, epoch: int): - self.filepath.parent.mkdir(parents=True, exist_ok=True) + self._save(trainer, epoch, logs) + + def _save(self, trainer, epoch: int, logs: Optional[Dict[str, float]] = None): + logs = dict(logs or {}) + # ``epoch`` is passed explicitly so templates without it still work. + logs.setdefault("epoch", epoch) + formatted = str(self.filepath).format(**logs) + path = Path(formatted) + path.parent.mkdir(parents=True, exist_ok=True) torch.save( { "epoch": epoch, "model_state_dict": trainer.model.state_dict(), "optimizer_state_dict": trainer.optimizer.state_dict(), }, - self.filepath, + path, ) + if self.verbose: + print(f"Epoch {epoch}: saved checkpoint to {path}") def on_train_end(self, trainer): pass diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py index 15c059a..2181b74 100644 --- a/src/jaeger/training/pytorch/engine.py +++ b/src/jaeger/training/pytorch/engine.py @@ -15,6 +15,8 @@ def train_one_epoch( forward_key: str = "prediction", label_key: str = "label", branch: str = "classifier", + progress: Optional[Any] = None, + task_id: Optional[Any] = None, ) -> Dict[str, float]: """Train for one epoch and return averaged loss and metrics.""" model.train() @@ -51,6 +53,10 @@ def train_one_epoch( total_loss += loss.item() * batch_size total_samples += batch_size + if progress is not None and task_id is not None: + progress.advance(task_id, 1) + progress.update(task_id, description=f"loss={loss.item():.4f}") + if metrics: for metric in metrics.values(): if hasattr(metric, "update"): @@ -76,6 +82,8 @@ def evaluate( metrics: Optional[Dict[str, Any]] = None, forward_key: str = "prediction", branch: str = "classifier", + progress: Optional[Any] = None, + task_id: Optional[Any] = None, ) -> Dict[str, float]: """Evaluate and return averaged loss and metrics.""" model.eval() @@ -100,6 +108,10 @@ def evaluate( total_loss += loss.item() * batch_size total_samples += batch_size + if progress is not None and task_id is not None: + progress.advance(task_id, 1) + progress.update(task_id, description=f"loss={loss.item():.4f}") + if metrics: for metric in metrics.values(): if hasattr(metric, "update"): diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 27bab82..06ef7d8 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -4,11 +4,31 @@ import torch import torch.nn as nn +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + ProgressColumn, + TextColumn, + TimeElapsedColumn, + TimeRemainingColumn, +) +from rich.text import Text from torch.utils.data import DataLoader from jaeger.training.pytorch.engine import evaluate, train_one_epoch +class _SpeedMsColumn(ProgressColumn): + """Rich column showing average milliseconds per batch.""" + + def render(self, task): + if task.completed == 0: + return Text("? ms/batch") + ms = task.elapsed / task.completed * 1000 + return Text(f"{ms:.1f} ms/batch") + + class Trainer: """High-level training loop for Jaeger PyTorch models.""" @@ -26,6 +46,7 @@ def __init__( checkpoint_dir: Optional[str] = None, history_path: Optional[str] = None, branch: str = "classifier", + progress_bar: bool = False, ): self.model = model self.train_loader = train_loader @@ -39,6 +60,7 @@ def __init__( self.checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir else None self.history_path = Path(history_path) if history_path else None self.branch = branch + self.progress_bar = progress_bar self.history: List[Dict[str, float]] = [] self.should_stop = False @@ -53,23 +75,62 @@ def fit(self) -> List[Dict[str, float]]: if hasattr(callback, "on_epoch_begin"): callback.on_epoch_begin(self, epoch) - train_metrics = train_one_epoch( - self.model, - self.train_loader, - self.loss_fn, - self.optimizer, - self.device, - metrics=self.metrics, - branch=self.branch, - ) - val_metrics = evaluate( - self.model, - self.val_loader, - self.loss_fn, - self.device, - metrics=self.metrics, - branch=self.branch, - ) + if self.progress_bar: + with Progress( + TextColumn("[progress.description]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TimeElapsedColumn(), + TimeRemainingColumn(), + _SpeedMsColumn(), + ) as progress: + train_task = progress.add_task( + f"train epoch {epoch}/{self.epochs}", + total=self._loader_length(self.train_loader), + ) + val_task = progress.add_task( + f"val epoch {epoch}/{self.epochs}", + total=self._loader_length(self.val_loader), + ) + train_metrics = train_one_epoch( + self.model, + self.train_loader, + self.loss_fn, + self.optimizer, + self.device, + metrics=self.metrics, + branch=self.branch, + progress=progress, + task_id=train_task, + ) + val_metrics = evaluate( + self.model, + self.val_loader, + self.loss_fn, + self.device, + metrics=self.metrics, + branch=self.branch, + progress=progress, + task_id=val_task, + ) + else: + train_metrics = train_one_epoch( + self.model, + self.train_loader, + self.loss_fn, + self.optimizer, + self.device, + metrics=self.metrics, + branch=self.branch, + ) + val_metrics = evaluate( + self.model, + self.val_loader, + self.loss_fn, + self.device, + metrics=self.metrics, + branch=self.branch, + ) epoch_log = {"epoch": epoch} for k, v in train_metrics.items(): @@ -96,6 +157,13 @@ def fit(self) -> List[Dict[str, float]]: return self.history + @staticmethod + def _loader_length(loader: DataLoader) -> Optional[int]: + try: + return len(loader) + except TypeError: + return None + def _save_checkpoint(self, epoch: int): if self.checkpoint_dir is None: return diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index da75421..9463f86 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -145,3 +145,31 @@ def test_train_fragment_core_only_classification_head(tmp_path): xla=False, meta=None, ) + + +def test_train_fragment_core_progress_bar(tmp_path): + """``train_fragment_core`` should support the ``--progress-bar`` path.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + progress_bar=True, + ) diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/data/pytorch/test_builders.py index ca8f187..dcbc4f7 100644 --- a/tests/unit/data/pytorch/test_builders.py +++ b/tests/unit/data/pytorch/test_builders.py @@ -110,6 +110,47 @@ def test_build_datasets_multiple_paths(tmp_path): assert batch_mask.shape == (8, 6, 50) +def _make_nucleotide_npz(path, n_samples, length=50): + # One-hot nucleotide tensor: (N, 2 strands, length, 4 bases) + data = np.eye(4, dtype=np.float32)[ + np.random.randint(0, 4, size=(n_samples, 2, length)) + ].astype(np.float32) + labels = np.eye(3, dtype=np.float32)[np.random.randint(0, 3, size=n_samples)] + np.savez(path, nucleotide=data, label=labels) + + +def test_build_datasets_numpy_full_nucleotide(tmp_path): + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_nucleotide_npz(train_path, n_samples=8) + _make_nucleotide_npz(val_path, n_samples=4) + + config = _build_config([train_path], [val_path], batch_size=4) + config["model"]["embedding"] = {"input_type": "nucleotide"} + loaders = build_datasets(config, branch="classifier") + + batch_x, batch_y, batch_mask = next(iter(loaders["train"])) + assert batch_x.shape == (4, 2, 50, 4) + assert batch_y.shape == (4, 3) + assert batch_mask.shape == (4, 2, 50) + + +def test_build_datasets_numpy_full_input_key_override(tmp_path): + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_nucleotide_npz(train_path, n_samples=8) + _make_nucleotide_npz(val_path, n_samples=4) + + config = _build_config([train_path], [val_path], batch_size=4) + # Even though embedding says translated, the explicit input_key wins. + config["model"]["embedding"] = {"input_type": "translated"} + config["model"]["string_processor"]["input_key"] = "nucleotide" + loaders = build_datasets(config, branch="classifier") + + batch_x, _, _ = next(iter(loaders["train"])) + assert batch_x.shape == (4, 2, 50, 4) + + def test_build_datasets_unsupported_format(tmp_path): train_path = tmp_path / "train.npz" val_path = tmp_path / "val.npz" diff --git a/tests/unit/test_dataops_convert.py b/tests/unit/test_dataops_convert.py index da28675..601702a 100644 --- a/tests/unit/test_dataops_convert.py +++ b/tests/unit/test_dataops_convert.py @@ -10,13 +10,70 @@ from jaeger.dataops import convert +class TestOneHotMemoryGuard: + """Memory guard prevents OOM on huge one-hot conversions.""" + + def test_estimate_nucleotide_onehot_memory(self): + # 1M rows, crop 500, one-hot nucleotide + # 1_000_000 * 2 * 500 * 4 * 4 = 16_000_000_000 bytes + assert convert._estimate_onehot_memory( + total_rows=1_000_000, + crop_size=500, + format="nucleotide", + one_hot=True, + ) == 16_000_000_000 + + def test_estimate_no_memory_for_non_onehot(self): + assert ( + convert._estimate_onehot_memory( + total_rows=1_000_000, + crop_size=500, + format="nucleotide", + one_hot=False, + ) + == 0 + ) + + def test_check_onehot_memory_raises_when_too_large(self): + with pytest.raises(MemoryError, match="one-hot"): + convert._check_onehot_memory( + estimated_bytes=100, + available_bytes=100, + fraction=0.5, + ) + + def test_check_onehot_memory_passes_when_small(self): + # Should not raise + convert._check_onehot_memory( + estimated_bytes=1, + available_bytes=100, + fraction=0.5, + ) + + def test_guard_rejects_huge_onehot_conversion(self, tmp_path: Path): + """A tiny CSV with enormous crop_size should be rejected before allocation.""" + csv = tmp_path / "big.csv" + csv.write_text("0,ATGC\n") + with pytest.raises(MemoryError, match="one-hot"): + convert.convert_dataset( + input_path=str(csv), + output_path=str(tmp_path / "out.npz"), + format="nucleotide", + one_hot=True, + crop_size=1_000_000_000, + num_classes=2, + num_workers=1, + _available_memory_bytes=1024, + ) + + class TestConvertDataset: - def test_numpy_full(self, simple_csv_path: str, tmp_path: Path): + def test_translated_integer(self, simple_csv_path: str, tmp_path: Path): out = tmp_path / "out.npz" convert.convert_dataset( input_path=simple_csv_path, output_path=str(out), - format="numpy_full", + format="translated", crop_size=24, num_classes=2, num_workers=1, @@ -25,30 +82,135 @@ def test_numpy_full(self, simple_csv_path: str, tmp_path: Path): data = np.load(out) assert "translated" in data assert "label" in data + assert "pad_int" in data + assert "lengths" in data + assert data["pad_int"].item() == 0 + assert data["translated"].dtype == np.int32 - def test_numpy_raw(self, simple_csv_path: str, tmp_path: Path): + def test_nucleotide_integer(self, simple_csv_path: str, tmp_path: Path): out = tmp_path / "out.npz" convert.convert_dataset( input_path=simple_csv_path, output_path=str(out), - format="numpy_raw", + format="nucleotide", crop_size=24, num_classes=2, num_workers=1, ) assert out.exists() + data = np.load(out) + assert "nucleotide" in data + assert "label" in data + assert "pad_int" in data + assert data["nucleotide"].dtype == np.int32 + + def test_lowercase_consistency(self, tmp_path: Path): + """Lowercase sequences must encode identically to uppercase.""" + csv = tmp_path / "mixed.csv" + csv.write_text("0,ATGCatgcNN\n") + out_upper = tmp_path / "out_upper.npz" + convert.convert_dataset( + input_path=str(csv), + output_path=str(out_upper), + format="nucleotide", + crop_size=24, + num_classes=2, + num_workers=1, + ) + csv_lower = tmp_path / "mixed_lower.csv" + csv_lower.write_text("0,atgcATGCnn\n") + out_lower = tmp_path / "out_lower.npz" + convert.convert_dataset( + input_path=str(csv_lower), + output_path=str(out_lower), + format="nucleotide", + crop_size=24, + num_classes=2, + num_workers=1, + ) + upper_data = np.load(out_upper) + lower_data = np.load(out_lower) + assert np.array_equal(upper_data["nucleotide"], lower_data["nucleotide"]) - def test_numpy_raw_variable(self, simple_csv_path: str, tmp_path: Path): + def test_both_onehot(self, simple_csv_path: str, tmp_path: Path): out = tmp_path / "out.npz" convert.convert_dataset( input_path=simple_csv_path, output_path=str(out), - format="numpy_raw_variable", + format="both", + one_hot=True, crop_size=24, num_classes=2, num_workers=1, ) assert out.exists() + data = np.load(out) + assert "nucleotide" in data + assert "translated" in data + assert "label" in data + assert data["nucleotide"].dtype == np.float32 + assert data["translated"].dtype == np.float32 + + def test_codon_map_aa_id(self, simple_csv_path: str, tmp_path: Path): + out = tmp_path / "out.npz" + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(out), + format="translated", + codon_map="AA_ID", + crop_size=24, + num_classes=2, + num_workers=1, + ) + data = np.load(out) + assert "translated" in data + assert data["vocab_size"].item() == 21 + + def test_pad_int(self, simple_csv_path: str, tmp_path: Path): + out = tmp_path / "out.npz" + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(out), + format="nucleotide", + pad_int=7, + crop_size=30, + num_classes=2, + num_workers=1, + ) + data = np.load(out) + assert data["pad_int"].item() == 7 + nuc = data["nucleotide"] + # Actual token indices are offset by +1, so 0 is never used. + assert (nuc == 0).sum() == 0 + + def test_pad_int_collision_raises(self, simple_csv_path: str, tmp_path: Path): + with pytest.raises(ValueError, match="collides"): + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(tmp_path / "out.npz"), + format="nucleotide", + pad_int=1, + crop_size=24, + num_classes=2, + num_workers=1, + ) + + def test_stride_sliding_window(self, tmp_path: Path): + csv = tmp_path / "long.csv" + # 1000 bp sequence, label 0 + csv.write_text(f"0,{'ATGC' * 250}\n") + out = tmp_path / "out.npz" + convert.convert_dataset( + input_path=str(csv), + output_path=str(out), + format="nucleotide", + crop_size=500, + stride=500, + num_classes=2, + num_workers=1, + ) + data = np.load(out) + assert data["nucleotide"].shape[0] == 2 def test_invalid_format(self, simple_csv_path: str, tmp_path: Path): with pytest.raises(ValueError): @@ -59,3 +221,80 @@ def test_invalid_format(self, simple_csv_path: str, tmp_path: Path): crop_size=24, num_classes=2, ) + + def test_default_workers(self, simple_csv_path: str, tmp_path: Path): + """convert_dataset should complete with the default worker count. + + This is a regression guard against hangs caused by the default + ``fork``-based ``multiprocessing.Pool`` on Linux when the parent + process has background threads. + """ + out = tmp_path / "out.npz" + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(out), + format="translated", + crop_size=24, + num_classes=2, + ) + assert out.exists() + data = np.load(out) + assert "translated" in data + assert "label" in data + + @pytest.mark.skipif( + not convert.HAS_NUMBA, reason="Numba not installed" + ) + def test_numba_outputs_match_python(self, simple_csv_path: str, tmp_path: Path, monkeypatch): + """Numba path must produce the same arrays as the pure-Python path.""" + assert convert.HAS_NUMBA, "Numba must be available for this test" + + out_numba = tmp_path / "out_numba.npz" + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(out_numba), + format="both", + one_hot=True, + crop_size=30, + num_classes=2, + num_workers=1, + ) + data_numba = dict(np.load(out_numba)) + + monkeypatch.setattr(convert, "HAS_NUMBA", False) + out_py = tmp_path / "out_py.npz" + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(out_py), + format="both", + one_hot=True, + crop_size=30, + num_classes=2, + num_workers=1, + ) + data_py = dict(np.load(out_py)) + + assert np.array_equal(data_numba["nucleotide"], data_py["nucleotide"]) + assert np.array_equal(data_numba["translated"], data_py["translated"]) + assert np.array_equal(data_numba["label"], data_py["label"]) + assert np.array_equal(data_numba["lengths"], data_py["lengths"]) + assert data_numba["pad_int"].item() == data_py["pad_int"].item() + assert data_numba["vocab_size"].item() == data_py["vocab_size"].item() + + def test_fallback_when_numba_disabled(self, simple_csv_path: str, tmp_path: Path, monkeypatch): + """Pure-Python fallback must still produce valid output.""" + monkeypatch.setattr(convert, "HAS_NUMBA", False) + out = tmp_path / "out_py.npz" + convert.convert_dataset( + input_path=simple_csv_path, + output_path=str(out), + format="translated", + crop_size=24, + num_classes=2, + num_workers=1, + ) + data = dict(np.load(out)) + assert "translated" in data + assert "label" in data + assert data["translated"].dtype == np.int32 + assert data["translated"].shape[0] > 0 diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py index 7290649..ccaef2f 100644 --- a/tests/unit/training/pytorch/test_engine.py +++ b/tests/unit/training/pytorch/test_engine.py @@ -1,4 +1,5 @@ import torch +from rich.progress import Progress from jaeger.training.pytorch.engine import evaluate, train_one_epoch @@ -47,3 +48,34 @@ def test_evaluate(): metrics = {} history = evaluate(model, loader, loss_fn, torch.device("cpu"), metrics=metrics) assert "loss" in history + + +def test_train_one_epoch_with_progress(): + model = _make_dummy_classifier() + loader = _make_dummy_loader() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + with Progress(transient=True) as progress: + task_id = progress.add_task("train", total=len(loader)) + history = train_one_epoch( + model, + loader, + loss_fn, + optimizer, + torch.device("cpu"), + progress=progress, + task_id=task_id, + ) + assert "loss" in history + + +def test_evaluate_with_progress(): + model = _make_dummy_classifier() + loader = _make_dummy_loader() + loss_fn = torch.nn.CrossEntropyLoss() + with Progress(transient=True) as progress: + task_id = progress.add_task("val", total=len(loader)) + history = evaluate( + model, loader, loss_fn, torch.device("cpu"), progress=progress, task_id=task_id + ) + assert "loss" in history diff --git a/tests/unit/training/pytorch/test_trainer.py b/tests/unit/training/pytorch/test_trainer.py index 87c8a60..01f2c14 100644 --- a/tests/unit/training/pytorch/test_trainer.py +++ b/tests/unit/training/pytorch/test_trainer.py @@ -61,3 +61,29 @@ def test_trainer_fit(): assert "val_loss" in history[0] assert (Path(tmpdir) / "history.json").exists() assert len(list(Path(tmpdir).glob("checkpoint_epoch_*.pt"))) == 2 + + +def test_trainer_fit_with_progress_bar(): + model = _make_dummy_model() + train_loader = _make_dummy_data() + val_loader = _make_dummy_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=1, + device=torch.device("cpu"), + checkpoint_dir=tmpdir, + history_path=Path(tmpdir) / "history.json", + progress_bar=True, + ) + history = trainer.fit() + assert len(history) == 1 + assert "train_loss" in history[0] + assert "val_loss" in history[0] diff --git a/train_config/nn_config_500bp_dvf.yaml b/train_config/nn_config_500bp_dvf.yaml new file mode 100644 index 0000000..a31e598 --- /dev/null +++ b/train_config/nn_config_500bp_dvf.yaml @@ -0,0 +1,198 @@ +model: + name: jaeger_500bp_axial_npz + experiment: 500bp_axial_npz + seed: 42 + classifier_out_dim: 3 + reliability_out_dim: 0 + base_dir: /home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/training/genomad_train_data/jaeger_format + class_label_map: + - class: chromosome + label: 0 + - class: virus + label: 1 + - class: plasmid + label: 2 + activation: gelu + mode: training + embedding: + use_embedding_layer: true + input_type: translated + strands: 2 + frames: 6 + length: null + input_shape: + - 6 + - null + embedding_size: 64 + embedding_regularizer: l2 + embedding_regularizer_w: 1.0e-05 + string_processor: + data_format: numpy_full + seq_onehot: false + codon: CODON + codon_id: CODON_ID + crop_size: 500 + buffer_size: 5000 + shuffle: true + reshuffle_each_iteration: true + mutate: false + mutation_rate: 0.1 + shuffle_frames: false + masking: false + classifier_labels: + - 0 + - 1 + - 2 + classifier_labels_map: + - 0 + - 1 + - 2 + representation_learner: + hidden_layers: + - name: masked_conv1d + config: + filters: 32 + kernel_size: 7 + strides: 1 + dilation_rate: 1 + use_bias: true + activation: null + kernel_regularizer: l2 + kernel_regularizer_w: 1.0e-05 + - name: masked_batchnorm + config: + return_nmd: false + - name: activation + config: + activation: gelu + - name: residual_block + config: + use_1x1conv: false + block_size: 2 + filters: 32 + kernel_size: 3 + use_bias: true + kernel_regularizer: l2 + kernel_regularizer_w: 1.0e-05 + - name: masked_batchnorm + config: + return_nmd: false + - name: activation + config: + activation: gelu + - name: axial_attention + config: + embed_dim: 32 + num_heads: 4 + feed_forward_dim: 128 + dropout_rate: 0.1 + num_blocks: 1 + - name: masked_batchnorm + config: + return_nmd: false + pooling: average + classifier: + input_shape: 32 + hidden_layers: + - name: dropout + config: + rate: 0.3 + - name: dense + config: + units: 3 + activation: null + dtype: float32 + use_bias: true + kernel_regularizer: l2 + kernel_regularizer_w: 0.01 +training: + data_dir: /home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/training/genomad_train_data/jaeger_format + experiment_root: experiments/experiment_{{ model.experiment }}_{{ model.seed }} + classifier_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" + reliability_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/reliability" + classifier_epochs: 5 + reliability_epochs: 0 + projection_epochs: 0 + classifier_train_steps: 1000 + reliability_train_steps: 0 + classifier_validation_steps: 100 + reliability_validation_steps: 0 + batch_size: 64 + optimizer: adam + optimizer_params: + learning_rate: 0.0001 + clipnorm: 5 + loss_classifier: categorical_crossentropy + loss_params_classifier: + from_logits: true + classifier_class_weights: + 0: 0.6912 + 1: 0.7641 + 2: 4.0898 + loss_reliability: binary_crossentropy + loss_params_reliability: + from_logits: true + metrics_classifier: + - name: categorical_accuracy + params: null + metrics_reliability: + - name: binary_accuracy + params: null + callbacks: + clean_old: true + directories: + - "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" + classifier: + - name: EarlyStopping + params: + monitor: val_loss + patience: 3 + mode: min + restore_best_weights: true + - name: ModelCheckpoint + params: + filepath: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/epoch:{epoch:02d}-loss:{val_loss:.2f}.weights.h5" + monitor: val_loss + mode: min + save_weights_only: true + verbose: 1 + save_best_only: false + - name: ReduceLROnPlateau + params: + monitor: val_loss + patience: 2 + factor: 0.5 + min_lr: 1.0e-05 + - name: TerminateOnNaN + - name: CSVLogger + params: + filename: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/training.log" + separator: "," + append: true + model_saving: + path: "{{ model.base_dir }}/{{ training.experiment_root }}/model" + save_weights: true + save_exec_graph: true + fragment_classifier_data: + train: + - class: + - chromosome + - virus + - plasmid + path: + - "{{ training.data_dir }}/train_shuffled_full.npz" + label: + - 0 + - 1 + - 2 + validation: + - class: + - chromosome + - virus + - plasmid + path: + - "{{ training.data_dir }}/val_shuffled_full.npz" + label: + - 0 + - 1 + - 2 diff --git a/train_config/nn_config_500bp_nucleotide.yaml b/train_config/nn_config_500bp_nucleotide.yaml new file mode 100644 index 0000000..bcaf316 --- /dev/null +++ b/train_config/nn_config_500bp_nucleotide.yaml @@ -0,0 +1,214 @@ +# 500 bp Jaeger model configuration for one-hot nucleotide input. +# The preprocessed NPZ files must contain a 'nucleotide' array of shape +# (N, 2, length, 4) with one-hot encoded strands and a 'label' array. +# +# Update model.base_dir and training.data_dir to point to your own dataset. + +model: + name: jaeger_500bp_nucleotide + experiment: 500bp_nucleotide + seed: 42 + classifier_out_dim: 3 + reliability_out_dim: 0 + base_dir: /path/to/base_dir + class_label_map: + - class: chromosome + label: 0 + - class: virus + label: 1 + - class: plasmid + label: 2 + activation: gelu + mode: training + + embedding: + # Integer-index nucleotide input. The Embedding layer one-hot encodes the + # integer indices and projects them to embedding_size. Set + # use_embedding_layer: true and input_shape: [2, null, 4] if you generated + # one-hot nucleotide data with --one-hot instead. + use_embedding_layer: false + input_type: nucleotide + strands: 2 + frames: null + length: null + input_shape: + - 2 + - null + vocab_size: 6 + embedding_size: 4 + embedding_regularizer: l2 + embedding_regularizer_w: 1.0e-05 + + string_processor: + data_format: numpy_full + # Explicit input key; builders.py will also infer 'nucleotide' from + # model.embedding.input_type automatically. + input_key: nucleotide + crop_size: 500 + buffer_size: 5000 + shuffle: false + reshuffle_each_iteration: false + mutate: false + mutation_rate: 0.1 + shuffle_frames: false + masking: false + classifier_labels: + - 0 + - 1 + - 2 + classifier_labels_map: + - 0 + - 1 + - 2 + + representation_learner: + hidden_layers: + - name: masked_conv1d + config: + filters: 32 + kernel_size: 7 + strides: 1 + dilation_rate: 1 + use_bias: true + activation: null + kernel_regularizer: l2 + kernel_regularizer_w: 1.0e-05 + - name: masked_batchnorm + config: + return_nmd: false + - name: activation + config: + activation: gelu + - name: residual_block + config: + use_1x1conv: false + block_size: 2 + filters: 32 + kernel_size: 3 + use_bias: true + kernel_regularizer: l2 + kernel_regularizer_w: 1.0e-05 + - name: masked_batchnorm + config: + return_nmd: false + - name: activation + config: + activation: gelu + - name: axial_attention + config: + embed_dim: 32 + num_heads: 4 + feed_forward_dim: 128 + dropout_rate: 0.1 + num_blocks: 1 + - name: masked_batchnorm + config: + return_nmd: false + pooling: average + + classifier: + input_shape: 32 + hidden_layers: + - name: dropout + config: + rate: 0.3 + - name: dense + config: + units: 3 + activation: null + dtype: float32 + use_bias: true + kernel_regularizer: l2 + kernel_regularizer_w: 0.01 + +training: + data_dir: /path/to/data_dir + experiment_root: experiments/experiment_{{ model.experiment }}_{{ model.seed }} + classifier_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" + reliability_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/reliability" + classifier_epochs: 5 + reliability_epochs: 0 + projection_epochs: 0 + classifier_train_steps: 1000 + reliability_train_steps: 0 + classifier_validation_steps: 100 + reliability_validation_steps: 0 + batch_size: 64 + optimizer: adam + optimizer_params: + learning_rate: 0.0001 + clipnorm: 5 + loss_classifier: categorical_crossentropy + loss_params_classifier: + from_logits: true + classifier_class_weights: + 0: 0.6912 + 1: 0.7641 + 2: 4.0898 + loss_reliability: binary_crossentropy + loss_params_reliability: + from_logits: true + metrics_classifier: + - name: categorical_accuracy + params: null + metrics_reliability: + - name: binary_accuracy + params: null + callbacks: + clean_old: true + directories: + - "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" + classifier: + - name: EarlyStopping + params: + monitor: val_loss + patience: 3 + mode: min + restore_best_weights: true + - name: ModelCheckpoint + params: + filepath: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/epoch:{epoch:02d}-loss:{val_loss:.2f}.weights.h5" + monitor: val_loss + mode: min + save_weights_only: true + verbose: 1 + save_best_only: false + - name: ReduceLROnPlateau + params: + monitor: val_loss + patience: 2 + factor: 0.5 + min_lr: 1.0e-05 + - name: TerminateOnNaN + - name: CSVLogger + params: + filename: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/training.log" + separator: "," + append: true + model_saving: + path: "{{ model.base_dir }}/{{ training.experiment_root }}/model" + save_weights: true + save_exec_graph: true + fragment_classifier_data: + train: + - class: + - chromosome + - virus + - plasmid + path: + - "{{ training.data_dir }}/train_nucleotide.npz" + label: + - 0 + - 1 + - 2 + validation: + - class: + - chromosome + - virus + - plasmid + path: + - "{{ training.data_dir }}/val_nucleotide.npz" + label: + - 0 + - 1 + - 2 diff --git a/train_config/nn_config_500bp_translated.yaml b/train_config/nn_config_500bp_translated.yaml new file mode 100644 index 0000000..23e953a --- /dev/null +++ b/train_config/nn_config_500bp_translated.yaml @@ -0,0 +1,209 @@ +# 500 bp Jaeger model configuration for 6-frame translated (codon) input. +# The preprocessed NPZ files must contain a 'translated' array of shape +# (N, 6, length) with integer codon indices and a 'label' array. +# +# Update model.base_dir and training.data_dir to point to your own dataset. + +model: + name: jaeger_500bp_translated + experiment: 500bp_translated + seed: 42 + classifier_out_dim: 3 + reliability_out_dim: 0 + base_dir: /path/to/base_dir + class_label_map: + - class: chromosome + label: 0 + - class: virus + label: 1 + - class: plasmid + label: 2 + activation: gelu + mode: training + + embedding: + use_embedding_layer: true + input_type: translated + strands: 2 + frames: 6 + length: null + input_shape: + - 6 + - null + embedding_size: 64 + embedding_regularizer: l2 + embedding_regularizer_w: 1.0e-05 + + string_processor: + data_format: numpy_full + seq_onehot: false + codon: CODON + codon_id: CODON_ID + crop_size: 500 + buffer_size: 5000 + shuffle: true + reshuffle_each_iteration: true + mutate: false + mutation_rate: 0.1 + shuffle_frames: false + masking: false + classifier_labels: + - 0 + - 1 + - 2 + classifier_labels_map: + - 0 + - 1 + - 2 + + representation_learner: + hidden_layers: + - name: masked_conv1d + config: + filters: 32 + kernel_size: 7 + strides: 1 + dilation_rate: 1 + use_bias: true + activation: null + kernel_regularizer: l2 + kernel_regularizer_w: 1.0e-05 + - name: masked_batchnorm + config: + return_nmd: false + - name: activation + config: + activation: gelu + - name: residual_block + config: + use_1x1conv: false + block_size: 2 + filters: 32 + kernel_size: 3 + use_bias: true + kernel_regularizer: l2 + kernel_regularizer_w: 1.0e-05 + - name: masked_batchnorm + config: + return_nmd: false + - name: activation + config: + activation: gelu + - name: axial_attention + config: + embed_dim: 32 + num_heads: 4 + feed_forward_dim: 128 + dropout_rate: 0.1 + num_blocks: 1 + - name: masked_batchnorm + config: + return_nmd: false + pooling: average + + classifier: + input_shape: 32 + hidden_layers: + - name: dropout + config: + rate: 0.3 + - name: dense + config: + units: 3 + activation: null + dtype: float32 + use_bias: true + kernel_regularizer: l2 + kernel_regularizer_w: 0.01 + +training: + data_dir: /path/to/data_dir + experiment_root: experiments/experiment_{{ model.experiment }}_{{ model.seed }} + classifier_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" + reliability_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/reliability" + classifier_epochs: 5 + reliability_epochs: 0 + projection_epochs: 0 + classifier_train_steps: 1000 + reliability_train_steps: 0 + classifier_validation_steps: 100 + reliability_validation_steps: 0 + batch_size: 64 + optimizer: adam + optimizer_params: + learning_rate: 0.0001 + clipnorm: 5 + loss_classifier: categorical_crossentropy + loss_params_classifier: + from_logits: true + classifier_class_weights: + 0: 0.6912 + 1: 0.7641 + 2: 4.0898 + loss_reliability: binary_crossentropy + loss_params_reliability: + from_logits: true + metrics_classifier: + - name: categorical_accuracy + params: null + metrics_reliability: + - name: binary_accuracy + params: null + callbacks: + clean_old: true + directories: + - "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" + classifier: + - name: EarlyStopping + params: + monitor: val_loss + patience: 3 + mode: min + restore_best_weights: true + - name: ModelCheckpoint + params: + filepath: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/epoch:{epoch:02d}-loss:{val_loss:.2f}.weights.h5" + monitor: val_loss + mode: min + save_weights_only: true + verbose: 1 + save_best_only: false + - name: ReduceLROnPlateau + params: + monitor: val_loss + patience: 2 + factor: 0.5 + min_lr: 1.0e-05 + - name: TerminateOnNaN + - name: CSVLogger + params: + filename: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/training.log" + separator: "," + append: true + model_saving: + path: "{{ model.base_dir }}/{{ training.experiment_root }}/model" + save_weights: true + save_exec_graph: true + fragment_classifier_data: + train: + - class: + - chromosome + - virus + - plasmid + path: + - "{{ training.data_dir }}/train_shuffled_full.npz" + label: + - 0 + - 1 + - 2 + validation: + - class: + - chromosome + - virus + - plasmid + path: + - "{{ training.data_dir }}/val_shuffled_full.npz" + label: + - 0 + - 1 + - 2 From 322ca0a25f2305c51369758147aec66ecebad471 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:03:40 +0200 Subject: [PATCH 40/64] feat(nnlib): native PyTorch layer registry, SiameseModel, and DVF baseline config - Add NativeSequential and a registry for native torch.nn layers (conv1d, linear, dropout, pooling, activations, shape helpers). - Add SiameseModel that runs two identical native-layer branches on Jaeger's 2-strand nucleotide input and averages the outputs. - Extend Embedding with onehot_dim so padding is always an all-zero vector and N can be excluded from the one-hot channels. - Update ModelBuilder to detect model_type: siamese / branch_layers. - Add train_config/nn_config_500bp_dvf.yaml DVF-style baseline. - Add unit tests for Embedding onehot_dim, SiameseModel shapes, and builder siamese integration. --- src/jaeger/nnlib/pytorch/builder.py | 21 +++- src/jaeger/nnlib/pytorch/layers.py | 116 ++++++++++++++++++++- src/jaeger/nnlib/pytorch/models.py | 82 +++++++++++++-- tests/unit/nnlib/pytorch/test_builder.py | 52 ++++++++++ tests/unit/nnlib/pytorch/test_models.py | 71 +++++++++++++ train_config/nn_config_500bp_dvf.yaml | 124 ++++++++++------------- 6 files changed, 385 insertions(+), 81 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py index cef84b1..27f0a07 100644 --- a/src/jaeger/nnlib/pytorch/builder.py +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -9,6 +9,7 @@ JaegerModel, ReliabilityHead, RepresentationModel, + SiameseModel, ) @@ -99,14 +100,26 @@ def build_fragment_classifier(self) -> Dict[str, nn.Module]: vocab_size=_resolve_vocab_size(embedding_cfg, string_cfg), embedding_size=embedding_cfg.get("embedding_size", 4), use_embedding_layer=embedding_cfg.get("use_embedding_layer", False), + onehot_dim=embedding_cfg.get("onehot_dim"), ) rep_cfg = self.model_cfg.get("representation_learner", {}) - rep_model = RepresentationModel( - embedding=embedding, - hidden_layers=rep_cfg.get("hidden_layers", []), - pooling=rep_cfg.get("pooling", "average"), + is_siamese = ( + self.model_cfg.get("model_type") == "siamese" + or "branch_layers" in rep_cfg ) + if is_siamese: + rep_model = SiameseModel( + embedding=embedding, + branch_layers=rep_cfg.get("branch_layers", []), + pooling=rep_cfg.get("pooling"), + ) + else: + rep_model = RepresentationModel( + embedding=embedding, + hidden_layers=rep_cfg.get("hidden_layers", []), + pooling=rep_cfg.get("pooling", "average"), + ) models["rep_model"] = rep_model cls_cfg = self.model_cfg["classifier"] diff --git a/src/jaeger/nnlib/pytorch/layers.py b/src/jaeger/nnlib/pytorch/layers.py index 353f432..edf0f9f 100644 --- a/src/jaeger/nnlib/pytorch/layers.py +++ b/src/jaeger/nnlib/pytorch/layers.py @@ -1,5 +1,5 @@ import math -from typing import Callable, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -505,3 +505,117 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): out, _ = self.attn(x_t, x_t, x_t, key_padding_mask=key_mask, need_weights=False) out = self.norm(x_t + out).reshape(b, length, f, d).permute(0, 2, 1, 3) return out, mask + + +# --------------------------------------------------------------------------- +# Shape helpers for native PyTorch layer pipelines +# --------------------------------------------------------------------------- +class Permute(nn.Module): + """Permute tensor dimensions.""" + + def __init__(self, dims: List[int]): + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.permute(*self.dims) + + +class Flatten(nn.Module): + """Flatten contiguous dimensions, defaulting to all dims after batch.""" + + def __init__(self, start_dim: int = 1, end_dim: int = -1): + super().__init__() + self.start_dim = start_dim + self.end_dim = end_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.flatten(x, start_dim=self.start_dim, end_dim=self.end_dim) + + +class SqueezeLast(nn.Module): + """Squeeze the last dimension if it is size 1.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[-1] == 1: + return x.squeeze(-1) + return x + + +# --------------------------------------------------------------------------- +# Native PyTorch layer registry +# --------------------------------------------------------------------------- +def _build_native_layer(name: str, config: Dict[str, Any]) -> nn.Module: + """Build a native ``torch.nn`` module from a config dict. + + The registry maps Keras-style layer names to PyTorch constructors. + All constructor arguments are passed through ``config``. + """ + config = config or {} + name = name.lower() + + registry: Dict[str, Callable[..., nn.Module]] = { + # Linear / dense + "linear": nn.Linear, + "dense": nn.Linear, + # Convolutions + "conv1d": nn.Conv1d, + "conv2d": nn.Conv2d, + # Normalization + "batchnorm1d": nn.BatchNorm1d, + "batchnorm2d": nn.BatchNorm2d, + "layernorm": nn.LayerNorm, + # Pooling + "adaptive_avg_pool1d": nn.AdaptiveAvgPool1d, + "adaptive_max_pool1d": nn.AdaptiveMaxPool1d, + "avg_pool1d": nn.AvgPool1d, + "max_pool1d": nn.MaxPool1d, + # Regularization + "dropout": nn.Dropout, + "dropout1d": nn.Dropout1d, + "dropout2d": nn.Dropout2d, + # Activations + "relu": nn.ReLU, + "gelu": nn.GELU, + "sigmoid": nn.Sigmoid, + "softmax": nn.Softmax, + "tanh": nn.Tanh, + # Shape helpers + "permute": Permute, + "flatten": Flatten, + "squeeze_last": SqueezeLast, + } + + if name not in registry: + raise ValueError( + f"Unknown native layer: {name!r}. " + f"Supported layers: {sorted(registry.keys())}" + ) + + return registry[name](**config) + + +class NativeSequential(nn.Module): + """Sequential container for native PyTorch layers. + + Accepts an optional mask argument for compatibility with Jaeger's masked + layer pipeline; the mask is returned unchanged. Shape-helper layers + (``permute``, ``flatten``, ``squeeze_last``) can be interleaved with + native ``torch.nn`` modules. + """ + + def __init__(self, layers: List[Dict[str, Any]]): + super().__init__() + modules: List[nn.Module] = [] + for layer_cfg in layers: + name = layer_cfg.get("name") + config = layer_cfg.get("config", {}) + modules.append(_build_native_layer(name, config)) + self.layers = nn.ModuleList(modules) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + for layer in self.layers: + x = layer(x) + return x, mask diff --git a/src/jaeger/nnlib/pytorch/models.py b/src/jaeger/nnlib/pytorch/models.py index ddc0633..9269076 100644 --- a/src/jaeger/nnlib/pytorch/models.py +++ b/src/jaeger/nnlib/pytorch/models.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import torch import torch.nn as nn @@ -29,6 +29,7 @@ def __init__( use_embedding_layer: bool, use_positional_embeddings: bool = False, positional_embedding_length: Optional[int] = None, + onehot_dim: Optional[int] = None, ): super().__init__() self.input_type = input_type @@ -36,6 +37,7 @@ def __init__( self.embedding_size = embedding_size self.use_embedding_layer = use_embedding_layer self.use_positional_embeddings = use_positional_embeddings + self.onehot_dim = onehot_dim if input_type == "translated": if use_embedding_layer: @@ -47,12 +49,28 @@ def __init__( elif input_type == "nucleotide": # Nucleotide input can be provided as a float one-hot tensor or as # integer indices. Integer indices are one-hot encoded and then - # optionally projected to ``embedding_size``. + # optionally projected to ``embedding_size``. Padding (token 0) is + # always mapped to an all-zero vector. + effective_onehot_dim = ( + onehot_dim if onehot_dim is not None else (vocab_size - 1) + ) + self.effective_onehot_dim = effective_onehot_dim if use_embedding_layer: self.embed = None + elif effective_onehot_dim == embedding_size: + self.embed = None else: - self.embed = nn.Linear(vocab_size, embedding_size, bias=False) + self.embed = nn.Linear(effective_onehot_dim, embedding_size, bias=False) nn.init.orthogonal_(self.embed.weight) + + # Mapping matrix from integer tokens to one-hot vectors. Row 0 + # (padding) is all-zero; token i maps to channel i-1 when it fits. + mapping = torch.zeros(vocab_size, effective_onehot_dim) + for token in range(1, vocab_size): + channel = token - 1 + if channel < effective_onehot_dim: + mapping[token, channel] = 1.0 + self.register_buffer("_nucleotide_onehot_mapping", mapping) else: raise ValueError(f"Invalid input_type: {input_type}") @@ -74,9 +92,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.embed(x) elif self.input_type == "nucleotide": if not torch.is_floating_point(x): - x = F.one_hot(x, num_classes=self.vocab_size).to( - torch.float32 - ) + x = F.embedding( + x, self._nucleotide_onehot_mapping.to(x.device) + ).to(torch.float32) if self.embed is not None: x = self.embed(x) if self.use_positional_embeddings: @@ -320,6 +338,58 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): return pooled +class SiameseModel(nn.Module): + """Siamese representation model with two identical native-layer branches. + + The input is expected to be a Jaeger nucleotide tensor of shape + ``(B, 2, L, D)`` (two strands after the embedding layer). Each strand is + processed by the same branch network and the resulting feature vectors are + averaged. + """ + + def __init__( + self, + embedding: nn.Module, + branch_layers: List[Dict[str, Any]], + pooling: Optional[str] = None, + ): + super().__init__() + from jaeger.nnlib.pytorch.layers import NativeSequential + + self.embedding = embedding + self.branch = NativeSequential(branch_layers) + self.output_dim = self._infer_branch_output_dim() + self.nmd_dim = self.output_dim + + def _infer_branch_output_dim(self) -> int: + """Return the output dimension of the last Linear layer, if known.""" + for layer in reversed(self.branch.layers): + if isinstance(layer, nn.Linear): + return int(layer.out_features) + raise ValueError( + "SiameseModel branch_layers must end with a linear layer " + "so the output dimension can be inferred." + ) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + x = self.embedding(x) # (B, 2, L, D) + x0 = x[:, 0] # (B, L, D) + x1 = x[:, 1] # (B, L, D) + + # Optional: zero padded positions before the branch. + if mask is not None: + mask0 = mask[:, 0].unsqueeze(-1).to(x.dtype) + mask1 = mask[:, 1].unsqueeze(-1).to(x.dtype) + x0 = x0 * mask0 + x1 = x1 * mask1 + + f0, _ = self.branch(x0) + f1, _ = self.branch(x1) + return (f0 + f1) * 0.5 + + class JaegerModel(nn.Module): """Combined model exposing all outputs for inference.""" diff --git a/tests/unit/nnlib/pytorch/test_builder.py b/tests/unit/nnlib/pytorch/test_builder.py index dbbf409..a07901d 100644 --- a/tests/unit/nnlib/pytorch/test_builder.py +++ b/tests/unit/nnlib/pytorch/test_builder.py @@ -98,3 +98,55 @@ def test_builder_mask_propagation(): assert out_all.shape == (2, 3) assert not torch.allclose(out_all[0], out_masked[0], atol=1e-6) assert torch.allclose(out_all[1], out_masked[1]) + + +def _dvf_config(): + return { + "model": { + "name": "dvf_test", + "classifier_out_dim": 3, + "reliability_out_dim": 0, + "model_type": "siamese", + "embedding": { + "input_type": "nucleotide", + "use_embedding_layer": False, + "vocab_size": 6, + "onehot_dim": 4, + "embedding_size": 4, + }, + "string_processor": {"input_key": "nucleotide"}, + "representation_learner": { + "branch_layers": [ + {"name": "permute", "config": {"dims": [0, 2, 1]}}, + { + "name": "conv1d", + "config": { + "in_channels": 4, + "out_channels": 8, + "kernel_size": 3, + }, + }, + {"name": "relu"}, + {"name": "adaptive_max_pool1d", "config": {"output_size": 1}}, + {"name": "squeeze_last"}, + {"name": "linear", "config": {"in_features": 8, "out_features": 16}}, + ] + }, + "classifier": { + "hidden_layers": [ + {"name": "linear", "config": {"in_features": 16, "out_features": 3}} + ] + }, + }, + "training": {"batch_size": 2, "optimizer": "adam", "optimizer_params": {}}, + } + + +def test_builder_siamese_dvf(): + builder = ModelBuilder(_dvf_config()) + models = builder.build_fragment_classifier() + x = torch.randint(0, 6, (2, 2, 20)) + mask = torch.ones(2, 2, 20, dtype=torch.bool) + out = models["jaeger_model"](x, mask) + assert out["prediction"].shape == (2, 3) + assert out["embedding"].shape == (2, 16) diff --git a/tests/unit/nnlib/pytorch/test_models.py b/tests/unit/nnlib/pytorch/test_models.py index 8c0f8c6..ed24021 100644 --- a/tests/unit/nnlib/pytorch/test_models.py +++ b/tests/unit/nnlib/pytorch/test_models.py @@ -1,3 +1,4 @@ +import pytest import torch from jaeger.nnlib.pytorch.models import ( ClassificationHead, @@ -6,6 +7,7 @@ ProjectionHead, ReliabilityHead, RepresentationModel, + SiameseModel, ) @@ -116,3 +118,72 @@ def test_jaeger_model_forward(): assert out["embedding"].shape == (2, 16) assert out["nmd"].shape == (2, 16) assert out["reliability"].shape == (2, 1) + + +def test_embedding_nucleotide_onehot_dim_4(): + """Padding and N must map to an all-zero one-hot vector.""" + emb = Embedding( + input_type="nucleotide", + vocab_size=6, + embedding_size=4, + use_embedding_layer=False, + onehot_dim=4, + ) + # tokens: pad=0, A=1, T=2, G=3, C=4, N=5 + x = torch.arange(6).view(6, 1) + out = emb(x) + assert out.shape == (6, 1, 4) + assert torch.allclose(out[0], torch.zeros(4)) # padding + assert torch.allclose(out[5], torch.zeros(4)) # N + assert torch.allclose(out[1], torch.tensor([1.0, 0.0, 0.0, 0.0])) + assert torch.allclose(out[4], torch.tensor([0.0, 0.0, 0.0, 1.0])) + + +def test_embedding_nucleotide_default_onehot_dim_excludes_padding(): + """Without onehot_dim, default to vocab_size-1 and keep padding zero.""" + emb = Embedding( + input_type="nucleotide", + vocab_size=6, + embedding_size=5, + use_embedding_layer=False, + ) + x = torch.arange(6).view(6, 1) + out = emb(x) + assert out.shape == (6, 1, 5) + assert torch.allclose(out[0], torch.zeros(5)) + + +def test_siamese_model_forward_shape(): + """SiameseModel averages two branch outputs.""" + embedding = Embedding( + input_type="nucleotide", + vocab_size=6, + embedding_size=4, + use_embedding_layer=False, + onehot_dim=4, + ) + branch_layers = [ + {"name": "permute", "config": {"dims": [0, 2, 1]}}, + {"name": "conv1d", "config": {"in_channels": 4, "out_channels": 8, "kernel_size": 3}}, + {"name": "relu"}, + {"name": "adaptive_max_pool1d", "config": {"output_size": 1}}, + {"name": "squeeze_last"}, + {"name": "linear", "config": {"in_features": 8, "out_features": 16}}, + ] + model = SiameseModel(embedding=embedding, branch_layers=branch_layers) + x = torch.randint(0, 6, (2, 2, 20)) + out = model(x) + assert out.shape == (2, 16) + + +def test_siamese_model_branch_must_end_with_linear(): + """SiameseModel needs a final linear layer to infer output_dim.""" + embedding = Embedding( + input_type="nucleotide", + vocab_size=6, + embedding_size=4, + use_embedding_layer=False, + onehot_dim=4, + ) + with pytest.raises(ValueError, match="linear layer"): + SiameseModel(embedding=embedding, branch_layers=[{"name": "relu"}]) diff --git a/train_config/nn_config_500bp_dvf.yaml b/train_config/nn_config_500bp_dvf.yaml index a31e598..7c870d6 100644 --- a/train_config/nn_config_500bp_dvf.yaml +++ b/train_config/nn_config_500bp_dvf.yaml @@ -1,10 +1,16 @@ +# 500 bp DeepVirFinder-style siamese Conv1d baseline for Jaeger. +# Trains on integer-encoded nucleotide NPZ files (2 strands) and one-hots +# A,T,G,C to 4 channels inside the model. N and padding map to all-zero. +# +# Update model.base_dir and training.data_dir to point to your own dataset. + model: - name: jaeger_500bp_axial_npz - experiment: 500bp_axial_npz + name: jaeger_500bp_dvf_baseline + experiment: 500bp_dvf_baseline seed: 42 classifier_out_dim: 3 reliability_out_dim: 0 - base_dir: /home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/training/genomad_train_data/jaeger_format + base_dir: /path/to/base_dir class_label_map: - class: chromosome label: 0 @@ -14,27 +20,28 @@ model: label: 2 activation: gelu mode: training + model_type: siamese + embedding: - use_embedding_layer: true - input_type: translated + use_embedding_layer: false + input_type: nucleotide strands: 2 - frames: 6 + frames: null length: null input_shape: - - 6 + - 2 - null - embedding_size: 64 - embedding_regularizer: l2 - embedding_regularizer_w: 1.0e-05 + vocab_size: 6 # integer tokens 0..5 (pad, A, T, G, C, N) + onehot_dim: 4 # A,T,G,C only; N and padding -> all-zero + embedding_size: 4 # no projection since onehot_dim == embedding_size + string_processor: data_format: numpy_full - seq_onehot: false - codon: CODON - codon_id: CODON_ID + input_key: nucleotide crop_size: 500 buffer_size: 5000 - shuffle: true - reshuffle_each_iteration: true + shuffle: false + reshuffle_each_iteration: false mutate: false mutation_rate: 0.1 shuffle_frames: false @@ -47,66 +54,43 @@ model: - 0 - 1 - 2 + representation_learner: - hidden_layers: - - name: masked_conv1d - config: - filters: 32 - kernel_size: 7 - strides: 1 - dilation_rate: 1 - use_bias: true - activation: null - kernel_regularizer: l2 - kernel_regularizer_w: 1.0e-05 - - name: masked_batchnorm - config: - return_nmd: false - - name: activation + branch_layers: + - name: permute config: - activation: gelu - - name: residual_block + dims: [0, 2, 1] # (B, L, 4) -> (B, 4, L) + - name: conv1d config: - use_1x1conv: false - block_size: 2 - filters: 32 - kernel_size: 3 - use_bias: true - kernel_regularizer: l2 - kernel_regularizer_w: 1.0e-05 - - name: masked_batchnorm + in_channels: 4 + out_channels: 32 + kernel_size: 7 + - name: relu + - name: adaptive_max_pool1d config: - return_nmd: false - - name: activation + output_size: 1 + - name: squeeze_last + - name: dropout config: - activation: gelu - - name: axial_attention + p: 0.1 + - name: linear config: - embed_dim: 32 - num_heads: 4 - feed_forward_dim: 128 - dropout_rate: 0.1 - num_blocks: 1 - - name: masked_batchnorm + in_features: 32 + out_features: 64 + - name: relu + - name: dropout config: - return_nmd: false - pooling: average + p: 0.1 + classifier: - input_shape: 32 hidden_layers: - - name: dropout - config: - rate: 0.3 - - name: dense + - name: linear config: - units: 3 - activation: null - dtype: float32 - use_bias: true - kernel_regularizer: l2 - kernel_regularizer_w: 0.01 + in_features: 64 + out_features: 3 + training: - data_dir: /home/yasas-wijesekara/ssd/Projects/Jaeger_revisions/training/genomad_train_data/jaeger_format + data_dir: /path/to/data_dir experiment_root: experiments/experiment_{{ model.experiment }}_{{ model.seed }} classifier_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" reliability_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/reliability" @@ -151,7 +135,7 @@ training: restore_best_weights: true - name: ModelCheckpoint params: - filepath: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/epoch:{epoch:02d}-loss:{val_loss:.2f}.weights.h5" + filepath: '{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/epoch:{epoch:02d}-loss:{val_loss:.2f}.weights.h5' monitor: val_loss mode: min save_weights_only: true @@ -166,11 +150,11 @@ training: - name: TerminateOnNaN - name: CSVLogger params: - filename: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/training.log" - separator: "," + filename: '{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/training.log' + separator: ',' append: true model_saving: - path: "{{ model.base_dir }}/{{ training.experiment_root }}/model" + path: '{{ model.base_dir }}/{{ training.experiment_root }}/model' save_weights: true save_exec_graph: true fragment_classifier_data: @@ -180,7 +164,7 @@ training: - virus - plasmid path: - - "{{ training.data_dir }}/train_shuffled_full.npz" + - "{{ training.data_dir }}/train_nucleotide.npz" label: - 0 - 1 @@ -191,7 +175,7 @@ training: - virus - plasmid path: - - "{{ training.data_dir }}/val_shuffled_full.npz" + - "{{ training.data_dir }}/val_nucleotide.npz" label: - 0 - 1 From 2a7e507c1c9aef6be5f5901af8bcb1795a500abd Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:10:39 +0200 Subject: [PATCH 41/64] refactor(data): move pytorch dataset loaders from data/ to dataops/ - Move src/jaeger/data/pytorch/ to src/jaeger/dataops/pytorch/. - Move corresponding tests from tests/unit/data/pytorch/ to tests/unit/dataops/pytorch/. - Update all imports from jaeger.data.pytorch to jaeger.dataops.pytorch. - Update AGENTS.md so data/ is described as bundled data/model artifacts and dataops/ includes the pytorch dataset loaders/builders. --- src/jaeger/{data => dataops}/pytorch/__init__.py | 0 src/jaeger/{data => dataops}/pytorch/builders.py | 0 src/jaeger/{data => dataops}/pytorch/collate.py | 0 src/jaeger/{data => dataops}/pytorch/dataset_csv.py | 0 src/jaeger/{data => dataops}/pytorch/dataset_numpy.py | 0 src/jaeger/{data => dataops}/pytorch/transforms.py | 0 tests/unit/{data => dataops}/pytorch/test_builders.py | 0 tests/unit/{data => dataops}/pytorch/test_collate.py | 0 tests/unit/{data => dataops}/pytorch/test_dataset_csv.py | 0 tests/unit/{data => dataops}/pytorch/test_dataset_numpy.py | 0 10 files changed, 0 insertions(+), 0 deletions(-) rename src/jaeger/{data => dataops}/pytorch/__init__.py (100%) rename src/jaeger/{data => dataops}/pytorch/builders.py (100%) rename src/jaeger/{data => dataops}/pytorch/collate.py (100%) rename src/jaeger/{data => dataops}/pytorch/dataset_csv.py (100%) rename src/jaeger/{data => dataops}/pytorch/dataset_numpy.py (100%) rename src/jaeger/{data => dataops}/pytorch/transforms.py (100%) rename tests/unit/{data => dataops}/pytorch/test_builders.py (100%) rename tests/unit/{data => dataops}/pytorch/test_collate.py (100%) rename tests/unit/{data => dataops}/pytorch/test_dataset_csv.py (100%) rename tests/unit/{data => dataops}/pytorch/test_dataset_numpy.py (100%) diff --git a/src/jaeger/data/pytorch/__init__.py b/src/jaeger/dataops/pytorch/__init__.py similarity index 100% rename from src/jaeger/data/pytorch/__init__.py rename to src/jaeger/dataops/pytorch/__init__.py diff --git a/src/jaeger/data/pytorch/builders.py b/src/jaeger/dataops/pytorch/builders.py similarity index 100% rename from src/jaeger/data/pytorch/builders.py rename to src/jaeger/dataops/pytorch/builders.py diff --git a/src/jaeger/data/pytorch/collate.py b/src/jaeger/dataops/pytorch/collate.py similarity index 100% rename from src/jaeger/data/pytorch/collate.py rename to src/jaeger/dataops/pytorch/collate.py diff --git a/src/jaeger/data/pytorch/dataset_csv.py b/src/jaeger/dataops/pytorch/dataset_csv.py similarity index 100% rename from src/jaeger/data/pytorch/dataset_csv.py rename to src/jaeger/dataops/pytorch/dataset_csv.py diff --git a/src/jaeger/data/pytorch/dataset_numpy.py b/src/jaeger/dataops/pytorch/dataset_numpy.py similarity index 100% rename from src/jaeger/data/pytorch/dataset_numpy.py rename to src/jaeger/dataops/pytorch/dataset_numpy.py diff --git a/src/jaeger/data/pytorch/transforms.py b/src/jaeger/dataops/pytorch/transforms.py similarity index 100% rename from src/jaeger/data/pytorch/transforms.py rename to src/jaeger/dataops/pytorch/transforms.py diff --git a/tests/unit/data/pytorch/test_builders.py b/tests/unit/dataops/pytorch/test_builders.py similarity index 100% rename from tests/unit/data/pytorch/test_builders.py rename to tests/unit/dataops/pytorch/test_builders.py diff --git a/tests/unit/data/pytorch/test_collate.py b/tests/unit/dataops/pytorch/test_collate.py similarity index 100% rename from tests/unit/data/pytorch/test_collate.py rename to tests/unit/dataops/pytorch/test_collate.py diff --git a/tests/unit/data/pytorch/test_dataset_csv.py b/tests/unit/dataops/pytorch/test_dataset_csv.py similarity index 100% rename from tests/unit/data/pytorch/test_dataset_csv.py rename to tests/unit/dataops/pytorch/test_dataset_csv.py diff --git a/tests/unit/data/pytorch/test_dataset_numpy.py b/tests/unit/dataops/pytorch/test_dataset_numpy.py similarity index 100% rename from tests/unit/data/pytorch/test_dataset_numpy.py rename to tests/unit/dataops/pytorch/test_dataset_numpy.py From 54c136f80945f93f72604bdd84b5b47ed7660605 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:23:58 +0200 Subject: [PATCH 42/64] feat(training): moving-average loss in progress bar + complete dataops refactor docs - Add loss_window_size parameter to train_one_epoch, evaluate, and Trainer. - Track a deque of recent batch losses and display the moving average in the Rich progress bar instead of the noisy single-batch loss. - Add unit test verifying the progress bar shows the moving average. - Commit the AGENTS.md and CHANGELOG.md updates that were missed in the previous dataops move refactor. --- AGENTS.md | 4 +-- CHANGELOG.md | 2 +- src/jaeger/commands/train.py | 2 +- src/jaeger/dataops/pytorch/__init__.py | 8 +++--- src/jaeger/dataops/pytorch/builders.py | 6 ++--- src/jaeger/dataops/pytorch/dataset_csv.py | 6 ++--- src/jaeger/dataops/pytorch/dataset_numpy.py | 2 +- src/jaeger/training/pytorch/engine.py | 19 +++++++++++--- src/jaeger/training/pytorch/trainer.py | 6 +++++ tests/unit/dataops/pytorch/test_builders.py | 2 +- tests/unit/dataops/pytorch/test_collate.py | 2 +- .../unit/dataops/pytorch/test_dataset_csv.py | 2 +- .../dataops/pytorch/test_dataset_numpy.py | 4 +-- tests/unit/training/pytorch/test_engine.py | 25 +++++++++++++++++++ 14 files changed, 66 insertions(+), 24 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f9cf951..6d7e93e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -83,8 +83,8 @@ jaeger health | `src/jaeger/cli.py` | Click CLI definition, PyTorch/legacy model routing, and TF import suppression for legacy workflows | | `src/jaeger/commands/` | Command implementations: `predict`, `predict_legacy`, `train`, `tune`, `health`, `quantize`, `convert_graph`, `taxonomy`, `downloads`, `test`, plus `utils*.py` | | `src/jaeger/nnlib/` | Neural-network library: `v1/` and `v2/` layers, losses, metrics, inference, conversion, builder | -| `src/jaeger/dataops/` | Data operations: `convert.py`, `dataset.py`, `ood.py`, `split.py` | -| `src/jaeger/data/` | Bundled data (`config.json`), TFRecord helpers, dataset loaders | +| `src/jaeger/dataops/` | Data operations: `convert.py`, `dataset.py`, `ood.py`, `split.py`, and the `pytorch/` dataset loaders/builders | +| `src/jaeger/data/` | Bundled data (`config.json`), model artifacts, and test fixtures | | `src/jaeger/seqops/` | Sequence operations: encode, maps, stats, synthetic, transform, validate, io | | `src/jaeger/postprocess/` | Post-processing: `collect.py`, `helpers.py`, `prophages.py` | | `src/jaeger/preprocess/` | Preprocessing maps/converters (`v1/`, `v2/`) | diff --git a/CHANGELOG.md b/CHANGELOG.md index 0650d8a..d6d03e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,7 +13,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. - Migrated the training and inference backend from TensorFlow to PyTorch. - New `jaeger.nnlib.pytorch` modules provide model building, layers, losses, metrics, and conversion utilities. - New `jaeger.training.pytorch` trainer, callbacks, and distributed helpers support multi-GPU and mixed-precision training. - - New `jaeger.data.pytorch` CSV/NumPy dataset builders and collators feed variable-length sequences to PyTorch models. + - New `jaeger.dataops.pytorch` CSV/NumPy dataset builders and collators feed variable-length sequences to PyTorch models. - `jaeger train` now runs the PyTorch training path by default. - `jaeger predict` now uses the PyTorch inference runner; legacy TensorFlow SavedModel support is retained only for deprecated models. diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 4c2152a..a781192 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -22,7 +22,7 @@ import torch import torch.nn as nn -from jaeger.data.pytorch.builders import build_datasets +from jaeger.dataops.pytorch.builders import build_datasets from jaeger.nnlib.pytorch.builder import ModelBuilder from jaeger.training.pytorch.callbacks import EarlyStopping, ModelCheckpoint from jaeger.training.pytorch.distributed import ( diff --git a/src/jaeger/dataops/pytorch/__init__.py b/src/jaeger/dataops/pytorch/__init__.py index 8797f46..51a391a 100644 --- a/src/jaeger/dataops/pytorch/__init__.py +++ b/src/jaeger/dataops/pytorch/__init__.py @@ -2,10 +2,10 @@ from __future__ import annotations -from jaeger.data.pytorch.builders import build_datasets -from jaeger.data.pytorch.collate import pad_collate -from jaeger.data.pytorch.dataset_csv import CSVDataset -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset +from jaeger.dataops.pytorch.builders import build_datasets +from jaeger.dataops.pytorch.collate import pad_collate +from jaeger.dataops.pytorch.dataset_csv import CSVDataset +from jaeger.dataops.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset __all__ = [ "build_datasets", diff --git a/src/jaeger/dataops/pytorch/builders.py b/src/jaeger/dataops/pytorch/builders.py index 000891e..59a8c1a 100644 --- a/src/jaeger/dataops/pytorch/builders.py +++ b/src/jaeger/dataops/pytorch/builders.py @@ -6,9 +6,9 @@ from torch.utils.data import ConcatDataset, DataLoader, Dataset -from jaeger.data.pytorch.collate import pad_collate -from jaeger.data.pytorch.dataset_csv import CSVDataset -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset +from jaeger.dataops.pytorch.collate import pad_collate +from jaeger.dataops.pytorch.dataset_csv import CSVDataset +from jaeger.dataops.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset from jaeger.seqops.maps import CODONS diff --git a/src/jaeger/dataops/pytorch/dataset_csv.py b/src/jaeger/dataops/pytorch/dataset_csv.py index b91365f..4a54243 100644 --- a/src/jaeger/dataops/pytorch/dataset_csv.py +++ b/src/jaeger/dataops/pytorch/dataset_csv.py @@ -10,7 +10,7 @@ import torch from torch.utils.data import Dataset -from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons +from jaeger.dataops.pytorch.transforms import dna_to_indices, translate_to_codons class CSVDataset(Dataset): @@ -58,7 +58,7 @@ def __getitem__(self, idx): label, seq = self.rows[idx] seq_indices = dna_to_indices(seq) if self.mutate: - from jaeger.data.pytorch.transforms import apply_mutation + from jaeger.dataops.pytorch.transforms import apply_mutation seq_indices = apply_mutation(seq_indices, self.mutation_rate) @@ -77,7 +77,7 @@ def __getitem__(self, idx): mask = x != 0 if self.shuffle_frames: - from jaeger.data.pytorch.transforms import shuffle_frames + from jaeger.dataops.pytorch.transforms import shuffle_frames x = shuffle_frames(x) diff --git a/src/jaeger/dataops/pytorch/dataset_numpy.py b/src/jaeger/dataops/pytorch/dataset_numpy.py index af8af1a..4ee12c8 100644 --- a/src/jaeger/dataops/pytorch/dataset_numpy.py +++ b/src/jaeger/dataops/pytorch/dataset_numpy.py @@ -10,7 +10,7 @@ import torch from torch.utils.data import Dataset -from jaeger.data.pytorch.transforms import ( +from jaeger.dataops.pytorch.transforms import ( apply_mutation, shuffle_frames as shuffle_frames_fn, translate_to_codons, diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py index 2181b74..ac6e8e7 100644 --- a/src/jaeger/training/pytorch/engine.py +++ b/src/jaeger/training/pytorch/engine.py @@ -1,3 +1,4 @@ +from collections import deque from typing import Any, Callable, Dict, Optional import torch @@ -17,11 +18,13 @@ def train_one_epoch( branch: str = "classifier", progress: Optional[Any] = None, task_id: Optional[Any] = None, + loss_window_size: int = 50, ) -> Dict[str, float]: """Train for one epoch and return averaged loss and metrics.""" model.train() total_loss = 0.0 total_samples = 0 + recent_losses: deque[float] = deque(maxlen=loss_window_size) if metrics: for m in metrics.values(): if hasattr(m, "reset"): @@ -50,12 +53,15 @@ def train_one_epoch( optimizer.step() batch_size = y.size(0) - total_loss += loss.item() * batch_size + loss_value = loss.item() + total_loss += loss_value * batch_size total_samples += batch_size + recent_losses.append(loss_value) if progress is not None and task_id is not None: + moving_avg = sum(recent_losses) / len(recent_losses) progress.advance(task_id, 1) - progress.update(task_id, description=f"loss={loss.item():.4f}") + progress.update(task_id, description=f"loss={moving_avg:.4f}") if metrics: for metric in metrics.values(): @@ -84,11 +90,13 @@ def evaluate( branch: str = "classifier", progress: Optional[Any] = None, task_id: Optional[Any] = None, + loss_window_size: int = 50, ) -> Dict[str, float]: """Evaluate and return averaged loss and metrics.""" model.eval() total_loss = 0.0 total_samples = 0 + recent_losses: deque[float] = deque(maxlen=loss_window_size) if metrics: for m in metrics.values(): if hasattr(m, "reset"): @@ -105,12 +113,15 @@ def evaluate( loss = loss_fn(preds, y) batch_size = y.size(0) - total_loss += loss.item() * batch_size + loss_value = loss.item() + total_loss += loss_value * batch_size total_samples += batch_size + recent_losses.append(loss_value) if progress is not None and task_id is not None: + moving_avg = sum(recent_losses) / len(recent_losses) progress.advance(task_id, 1) - progress.update(task_id, description=f"loss={loss.item():.4f}") + progress.update(task_id, description=f"loss={moving_avg:.4f}") if metrics: for metric in metrics.values(): diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 06ef7d8..0bb1b84 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -47,6 +47,7 @@ def __init__( history_path: Optional[str] = None, branch: str = "classifier", progress_bar: bool = False, + loss_window_size: int = 50, ): self.model = model self.train_loader = train_loader @@ -61,6 +62,7 @@ def __init__( self.history_path = Path(history_path) if history_path else None self.branch = branch self.progress_bar = progress_bar + self.loss_window_size = loss_window_size self.history: List[Dict[str, float]] = [] self.should_stop = False @@ -102,6 +104,7 @@ def fit(self) -> List[Dict[str, float]]: branch=self.branch, progress=progress, task_id=train_task, + loss_window_size=self.loss_window_size, ) val_metrics = evaluate( self.model, @@ -112,6 +115,7 @@ def fit(self) -> List[Dict[str, float]]: branch=self.branch, progress=progress, task_id=val_task, + loss_window_size=self.loss_window_size, ) else: train_metrics = train_one_epoch( @@ -122,6 +126,7 @@ def fit(self) -> List[Dict[str, float]]: self.device, metrics=self.metrics, branch=self.branch, + loss_window_size=self.loss_window_size, ) val_metrics = evaluate( self.model, @@ -130,6 +135,7 @@ def fit(self) -> List[Dict[str, float]]: self.device, metrics=self.metrics, branch=self.branch, + loss_window_size=self.loss_window_size, ) epoch_log = {"epoch": epoch} diff --git a/tests/unit/dataops/pytorch/test_builders.py b/tests/unit/dataops/pytorch/test_builders.py index dcbc4f7..d2fcd07 100644 --- a/tests/unit/dataops/pytorch/test_builders.py +++ b/tests/unit/dataops/pytorch/test_builders.py @@ -2,7 +2,7 @@ import pytest import torch -from jaeger.data.pytorch.builders import build_datasets +from jaeger.dataops.pytorch.builders import build_datasets from jaeger.nnlib.pytorch.builder import ModelBuilder diff --git a/tests/unit/dataops/pytorch/test_collate.py b/tests/unit/dataops/pytorch/test_collate.py index 9ff2554..0e63a61 100644 --- a/tests/unit/dataops/pytorch/test_collate.py +++ b/tests/unit/dataops/pytorch/test_collate.py @@ -1,6 +1,6 @@ import torch -from jaeger.data.pytorch.collate import pad_collate +from jaeger.dataops.pytorch.collate import pad_collate def test_pad_collate_variable_length_2d(): diff --git a/tests/unit/dataops/pytorch/test_dataset_csv.py b/tests/unit/dataops/pytorch/test_dataset_csv.py index bb7d24b..24c2c8d 100644 --- a/tests/unit/dataops/pytorch/test_dataset_csv.py +++ b/tests/unit/dataops/pytorch/test_dataset_csv.py @@ -1,6 +1,6 @@ import csv -from jaeger.data.pytorch.dataset_csv import CSVDataset +from jaeger.dataops.pytorch.dataset_csv import CSVDataset from jaeger.seqops.maps import CODONS diff --git a/tests/unit/dataops/pytorch/test_dataset_numpy.py b/tests/unit/dataops/pytorch/test_dataset_numpy.py index c7f9b37..ed2147d 100644 --- a/tests/unit/dataops/pytorch/test_dataset_numpy.py +++ b/tests/unit/dataops/pytorch/test_dataset_numpy.py @@ -1,7 +1,7 @@ import numpy as np import torch -from jaeger.data.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset -from jaeger.data.pytorch.transforms import dna_to_indices, translate_to_codons +from jaeger.dataops.pytorch.dataset_numpy import NumpyFullDataset, NumpyRawDataset +from jaeger.dataops.pytorch.transforms import dna_to_indices, translate_to_codons from jaeger.seqops.maps import CODONS diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py index ccaef2f..35b90bc 100644 --- a/tests/unit/training/pytorch/test_engine.py +++ b/tests/unit/training/pytorch/test_engine.py @@ -79,3 +79,28 @@ def test_evaluate_with_progress(): model, loader, loss_fn, torch.device("cpu"), progress=progress, task_id=task_id ) assert "loss" in history + + +def test_train_one_epoch_progress_uses_moving_average(): + """The progress bar description should show the moving average, not a single batch.""" + model = _make_dummy_classifier() + loader = _make_dummy_loader(n=8) + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + with Progress(transient=True) as progress: + task_id = progress.add_task("train", total=len(loader)) + train_one_epoch( + model, + loader, + loss_fn, + optimizer, + torch.device("cpu"), + progress=progress, + task_id=task_id, + loss_window_size=4, + ) + description = progress.tasks[task_id].description + assert description.startswith("loss=") + # With a window of 4 over 8 identical batches the average should be stable. + loss_value = float(description.split("=")[1]) + assert loss_value > 0 From 526c598cbdfee32b59fa97c4a3aa1f4108a3734b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:29:22 +0200 Subject: [PATCH 43/64] chore(config): untrack train_config/nn_config_500bp_dvf.yaml The DVF baseline config contains user-specific paths and is tuned per experiment. Remove it from version control and add it to .gitignore so it can be edited locally without showing up in git status. --- .gitignore | 1 + train_config/nn_config_500bp_dvf.yaml | 182 -------------------------- 2 files changed, 1 insertion(+), 182 deletions(-) delete mode 100644 train_config/nn_config_500bp_dvf.yaml diff --git a/.gitignore b/.gitignore index ecb6cfd..1a27497 100644 --- a/.gitignore +++ b/.gitignore @@ -63,3 +63,4 @@ kimi-export-*.md # Local/test train configs train_config/nn_config_500bp_*_test.yaml train_config/nn_config_500bp_*_zeus_v2.yaml +train_config/nn_config_500bp_dvf.yaml diff --git a/train_config/nn_config_500bp_dvf.yaml b/train_config/nn_config_500bp_dvf.yaml deleted file mode 100644 index 7c870d6..0000000 --- a/train_config/nn_config_500bp_dvf.yaml +++ /dev/null @@ -1,182 +0,0 @@ -# 500 bp DeepVirFinder-style siamese Conv1d baseline for Jaeger. -# Trains on integer-encoded nucleotide NPZ files (2 strands) and one-hots -# A,T,G,C to 4 channels inside the model. N and padding map to all-zero. -# -# Update model.base_dir and training.data_dir to point to your own dataset. - -model: - name: jaeger_500bp_dvf_baseline - experiment: 500bp_dvf_baseline - seed: 42 - classifier_out_dim: 3 - reliability_out_dim: 0 - base_dir: /path/to/base_dir - class_label_map: - - class: chromosome - label: 0 - - class: virus - label: 1 - - class: plasmid - label: 2 - activation: gelu - mode: training - model_type: siamese - - embedding: - use_embedding_layer: false - input_type: nucleotide - strands: 2 - frames: null - length: null - input_shape: - - 2 - - null - vocab_size: 6 # integer tokens 0..5 (pad, A, T, G, C, N) - onehot_dim: 4 # A,T,G,C only; N and padding -> all-zero - embedding_size: 4 # no projection since onehot_dim == embedding_size - - string_processor: - data_format: numpy_full - input_key: nucleotide - crop_size: 500 - buffer_size: 5000 - shuffle: false - reshuffle_each_iteration: false - mutate: false - mutation_rate: 0.1 - shuffle_frames: false - masking: false - classifier_labels: - - 0 - - 1 - - 2 - classifier_labels_map: - - 0 - - 1 - - 2 - - representation_learner: - branch_layers: - - name: permute - config: - dims: [0, 2, 1] # (B, L, 4) -> (B, 4, L) - - name: conv1d - config: - in_channels: 4 - out_channels: 32 - kernel_size: 7 - - name: relu - - name: adaptive_max_pool1d - config: - output_size: 1 - - name: squeeze_last - - name: dropout - config: - p: 0.1 - - name: linear - config: - in_features: 32 - out_features: 64 - - name: relu - - name: dropout - config: - p: 0.1 - - classifier: - hidden_layers: - - name: linear - config: - in_features: 64 - out_features: 3 - -training: - data_dir: /path/to/data_dir - experiment_root: experiments/experiment_{{ model.experiment }}_{{ model.seed }} - classifier_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" - reliability_dir: "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/reliability" - classifier_epochs: 5 - reliability_epochs: 0 - projection_epochs: 0 - classifier_train_steps: 1000 - reliability_train_steps: 0 - classifier_validation_steps: 100 - reliability_validation_steps: 0 - batch_size: 64 - optimizer: adam - optimizer_params: - learning_rate: 0.0001 - clipnorm: 5 - loss_classifier: categorical_crossentropy - loss_params_classifier: - from_logits: true - classifier_class_weights: - 0: 0.6912 - 1: 0.7641 - 2: 4.0898 - loss_reliability: binary_crossentropy - loss_params_reliability: - from_logits: true - metrics_classifier: - - name: categorical_accuracy - params: null - metrics_reliability: - - name: binary_accuracy - params: null - callbacks: - clean_old: true - directories: - - "{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier" - classifier: - - name: EarlyStopping - params: - monitor: val_loss - patience: 3 - mode: min - restore_best_weights: true - - name: ModelCheckpoint - params: - filepath: '{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/epoch:{epoch:02d}-loss:{val_loss:.2f}.weights.h5' - monitor: val_loss - mode: min - save_weights_only: true - verbose: 1 - save_best_only: false - - name: ReduceLROnPlateau - params: - monitor: val_loss - patience: 2 - factor: 0.5 - min_lr: 1.0e-05 - - name: TerminateOnNaN - - name: CSVLogger - params: - filename: '{{ model.base_dir }}/{{ training.experiment_root }}/checkpoints/classifier/training.log' - separator: ',' - append: true - model_saving: - path: '{{ model.base_dir }}/{{ training.experiment_root }}/model' - save_weights: true - save_exec_graph: true - fragment_classifier_data: - train: - - class: - - chromosome - - virus - - plasmid - path: - - "{{ training.data_dir }}/train_nucleotide.npz" - label: - - 0 - - 1 - - 2 - validation: - - class: - - chromosome - - virus - - plasmid - path: - - "{{ training.data_dir }}/val_nucleotide.npz" - label: - - 0 - - 1 - - 2 From fe98769db1b53ee00bb4876bf09573b5288cd1eb Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:37:14 +0200 Subject: [PATCH 44/64] feat(training): show cumulative epoch mean loss in progress bar Replace the per-batch and moving-average loss displays with the running mean over all batches processed so far in the current epoch. This gives a smoother, more stable progress indicator. --- src/jaeger/training/pytorch/engine.py | 23 +++++++++++----------- src/jaeger/training/pytorch/trainer.py | 6 ------ tests/unit/training/pytorch/test_engine.py | 6 ++---- 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py index ac6e8e7..db58b2b 100644 --- a/src/jaeger/training/pytorch/engine.py +++ b/src/jaeger/training/pytorch/engine.py @@ -1,4 +1,3 @@ -from collections import deque from typing import Any, Callable, Dict, Optional import torch @@ -18,13 +17,13 @@ def train_one_epoch( branch: str = "classifier", progress: Optional[Any] = None, task_id: Optional[Any] = None, - loss_window_size: int = 50, ) -> Dict[str, float]: """Train for one epoch and return averaged loss and metrics.""" model.train() total_loss = 0.0 total_samples = 0 - recent_losses: deque[float] = deque(maxlen=loss_window_size) + progress_loss = 0.0 + progress_batches = 0 if metrics: for m in metrics.values(): if hasattr(m, "reset"): @@ -56,12 +55,13 @@ def train_one_epoch( loss_value = loss.item() total_loss += loss_value * batch_size total_samples += batch_size - recent_losses.append(loss_value) + progress_loss += loss_value + progress_batches += 1 if progress is not None and task_id is not None: - moving_avg = sum(recent_losses) / len(recent_losses) + mean_loss = progress_loss / progress_batches progress.advance(task_id, 1) - progress.update(task_id, description=f"loss={moving_avg:.4f}") + progress.update(task_id, description=f"loss={mean_loss:.4f}") if metrics: for metric in metrics.values(): @@ -90,13 +90,13 @@ def evaluate( branch: str = "classifier", progress: Optional[Any] = None, task_id: Optional[Any] = None, - loss_window_size: int = 50, ) -> Dict[str, float]: """Evaluate and return averaged loss and metrics.""" model.eval() total_loss = 0.0 total_samples = 0 - recent_losses: deque[float] = deque(maxlen=loss_window_size) + progress_loss = 0.0 + progress_batches = 0 if metrics: for m in metrics.values(): if hasattr(m, "reset"): @@ -116,12 +116,13 @@ def evaluate( loss_value = loss.item() total_loss += loss_value * batch_size total_samples += batch_size - recent_losses.append(loss_value) + progress_loss += loss_value + progress_batches += 1 if progress is not None and task_id is not None: - moving_avg = sum(recent_losses) / len(recent_losses) + mean_loss = progress_loss / progress_batches progress.advance(task_id, 1) - progress.update(task_id, description=f"loss={moving_avg:.4f}") + progress.update(task_id, description=f"loss={mean_loss:.4f}") if metrics: for metric in metrics.values(): diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 0bb1b84..06ef7d8 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -47,7 +47,6 @@ def __init__( history_path: Optional[str] = None, branch: str = "classifier", progress_bar: bool = False, - loss_window_size: int = 50, ): self.model = model self.train_loader = train_loader @@ -62,7 +61,6 @@ def __init__( self.history_path = Path(history_path) if history_path else None self.branch = branch self.progress_bar = progress_bar - self.loss_window_size = loss_window_size self.history: List[Dict[str, float]] = [] self.should_stop = False @@ -104,7 +102,6 @@ def fit(self) -> List[Dict[str, float]]: branch=self.branch, progress=progress, task_id=train_task, - loss_window_size=self.loss_window_size, ) val_metrics = evaluate( self.model, @@ -115,7 +112,6 @@ def fit(self) -> List[Dict[str, float]]: branch=self.branch, progress=progress, task_id=val_task, - loss_window_size=self.loss_window_size, ) else: train_metrics = train_one_epoch( @@ -126,7 +122,6 @@ def fit(self) -> List[Dict[str, float]]: self.device, metrics=self.metrics, branch=self.branch, - loss_window_size=self.loss_window_size, ) val_metrics = evaluate( self.model, @@ -135,7 +130,6 @@ def fit(self) -> List[Dict[str, float]]: self.device, metrics=self.metrics, branch=self.branch, - loss_window_size=self.loss_window_size, ) epoch_log = {"epoch": epoch} diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py index 35b90bc..bdc1f83 100644 --- a/tests/unit/training/pytorch/test_engine.py +++ b/tests/unit/training/pytorch/test_engine.py @@ -81,8 +81,8 @@ def test_evaluate_with_progress(): assert "loss" in history -def test_train_one_epoch_progress_uses_moving_average(): - """The progress bar description should show the moving average, not a single batch.""" +def test_train_one_epoch_progress_shows_cumulative_mean(): + """The progress bar description should show the cumulative mean loss so far.""" model = _make_dummy_classifier() loader = _make_dummy_loader(n=8) loss_fn = torch.nn.CrossEntropyLoss() @@ -97,10 +97,8 @@ def test_train_one_epoch_progress_uses_moving_average(): torch.device("cpu"), progress=progress, task_id=task_id, - loss_window_size=4, ) description = progress.tasks[task_id].description assert description.startswith("loss=") - # With a window of 4 over 8 identical batches the average should be stable. loss_value = float(description.split("=")[1]) assert loss_value > 0 From b63bbacfd86805f1b8a1f752510df471f73d421d Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:53:28 +0200 Subject: [PATCH 45/64] feat(training): add per-batch timing profiling to train/eval loops - Instrument train_one_epoch and evaluate with time.perf_counter() to measure data loading, forward+loss, backward, optimizer step, and metrics update. - Add profile flag to Trainer and a --profile CLI flag for jaeger train. - When profiling is enabled, the progress bar shows per-section timings and the returned history includes average timing keys (time_*_ms). --- src/jaeger/commands/train.py | 10 ++ src/jaeger/training/pytorch/engine.py | 124 ++++++++++++++++++--- src/jaeger/training/pytorch/trainer.py | 6 + tests/unit/training/pytorch/test_engine.py | 29 +++++ 4 files changed, 155 insertions(+), 14 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index a781192..0d9da78 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -250,6 +250,12 @@ def forward( help="Enable XLA JIT compilation for training.", ) @click.option("--meta", type=click.Path(), default=None) +@click.option( + "--profile", + is_flag=True, + default=False, + help="Show per-section timing in the progress bar.", +) def train_fragment( config, mixed_precision, @@ -263,6 +269,7 @@ def train_fragment( self_supervised_pretraining, xla, meta, + profile, ): """Train a fragment-level Jaeger model.""" train_fragment_core( @@ -278,6 +285,7 @@ def train_fragment( self_supervised_pretraining=self_supervised_pretraining, xla=xla, meta=meta, + profile=profile, ) @@ -393,6 +401,7 @@ def train_fragment_core(**kwargs): history_path=str(classifier_dir / "history.json"), branch="classifier", progress_bar=kwargs.get("progress_bar", False), + profile=kwargs.get("profile", False), ) trainer.fit() @@ -450,6 +459,7 @@ def train_fragment_core(**kwargs): history_path=str(reliability_dir / "history.json"), branch="reliability", progress_bar=kwargs.get("progress_bar", False), + profile=kwargs.get("profile", False), ) rel_trainer.fit() diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py index db58b2b..0909b14 100644 --- a/src/jaeger/training/pytorch/engine.py +++ b/src/jaeger/training/pytorch/engine.py @@ -1,3 +1,4 @@ +import time from typing import Any, Callable, Dict, Optional import torch @@ -5,6 +6,13 @@ from torch.utils.data import DataLoader +def _format_timing(data_ms: float, forward_ms: float, backward_ms: float = 0.0, optim_ms: float = 0.0, metrics_ms: float = 0.0) -> str: + """Return a compact timing string for the progress bar.""" + if backward_ms > 0 or optim_ms > 0: + return f" d={data_ms:.1f}ms f={forward_ms:.1f}ms b={backward_ms:.1f}ms o={optim_ms:.1f}ms m={metrics_ms:.1f}ms" + return f" d={data_ms:.1f}ms f={forward_ms:.1f}ms m={metrics_ms:.1f}ms" + + def train_one_epoch( model: nn.Module, dataloader: DataLoader, @@ -17,39 +25,68 @@ def train_one_epoch( branch: str = "classifier", progress: Optional[Any] = None, task_id: Optional[Any] = None, + profile: bool = False, ) -> Dict[str, float]: - """Train for one epoch and return averaged loss and metrics.""" + """Train for one epoch and return averaged loss and metrics. + + When ``profile=True``, per-section timings are included in the progress + bar description and returned under keys prefixed with ``time_``. + """ model.train() total_loss = 0.0 total_samples = 0 progress_loss = 0.0 progress_batches = 0 + timing = { + "data": 0.0, + "forward": 0.0, + "backward": 0.0, + "optim": 0.0, + "metrics": 0.0, + } if metrics: for m in metrics.values(): if hasattr(m, "reset"): m.reset() + data_start = time.perf_counter() for batch in dataloader: + data_time = time.perf_counter() - data_start + timing["data"] += data_time + if branch == "classifier": x, y, mask = [b.to(device) for b in batch] optimizer.zero_grad() + forward_start = time.perf_counter() outputs = model(x, mask=mask) if isinstance(outputs, dict): preds = outputs[forward_key] else: preds = outputs loss = loss_fn(preds, y) + forward_time = time.perf_counter() - forward_start + timing["forward"] += forward_time elif branch == "reliability": x, y, mask = [b.to(device) for b in batch] optimizer.zero_grad() + forward_start = time.perf_counter() outputs = model(x, mask=mask) preds = outputs[forward_key] loss = loss_fn(preds, y) + forward_time = time.perf_counter() - forward_start + timing["forward"] += forward_time else: raise ValueError(f"Unknown branch: {branch}") + backward_start = time.perf_counter() loss.backward() + backward_time = time.perf_counter() - backward_start + timing["backward"] += backward_time + + optim_start = time.perf_counter() optimizer.step() + optim_time = time.perf_counter() - optim_start + timing["optim"] += optim_time batch_size = y.size(0) loss_value = loss.item() @@ -58,18 +95,32 @@ def train_one_epoch( progress_loss += loss_value progress_batches += 1 - if progress is not None and task_id is not None: - mean_loss = progress_loss / progress_batches - progress.advance(task_id, 1) - progress.update(task_id, description=f"loss={mean_loss:.4f}") - + metrics_start = time.perf_counter() if metrics: for metric in metrics.values(): if hasattr(metric, "update"): metric.update(preds.detach().cpu(), y.detach().cpu()) + metrics_time = time.perf_counter() - metrics_start + timing["metrics"] += metrics_time + + if progress is not None and task_id is not None: + mean_loss = progress_loss / progress_batches + description = f"loss={mean_loss:.4f}" + if profile: + description += _format_timing( + data_time * 1000, + forward_time * 1000, + backward_time * 1000, + optim_time * 1000, + metrics_time * 1000, + ) + progress.advance(task_id, 1) + progress.update(task_id, description=description) + + data_start = time.perf_counter() avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 - result = {"loss": avg_loss} + result: Dict[str, float] = {"loss": avg_loss} if metrics: result.update( { @@ -77,6 +128,17 @@ def train_one_epoch( for name, metric in metrics.items() } ) + if profile: + n = progress_batches if progress_batches > 0 else 1 + result.update( + { + "time_data_ms": timing["data"] / n * 1000, + "time_forward_ms": timing["forward"] / n * 1000, + "time_backward_ms": timing["backward"] / n * 1000, + "time_optim_ms": timing["optim"] / n * 1000, + "time_metrics_ms": timing["metrics"] / n * 1000, + } + ) return result @@ -90,27 +152,40 @@ def evaluate( branch: str = "classifier", progress: Optional[Any] = None, task_id: Optional[Any] = None, + profile: bool = False, ) -> Dict[str, float]: - """Evaluate and return averaged loss and metrics.""" + """Evaluate and return averaged loss and metrics. + + When ``profile=True``, per-section timings are included in the progress + bar description and returned under keys prefixed with ``time_``. + """ model.eval() total_loss = 0.0 total_samples = 0 progress_loss = 0.0 progress_batches = 0 + timing = {"data": 0.0, "forward": 0.0, "metrics": 0.0} if metrics: for m in metrics.values(): if hasattr(m, "reset"): m.reset() with torch.no_grad(): + data_start = time.perf_counter() for batch in dataloader: + data_time = time.perf_counter() - data_start + timing["data"] += data_time + x, y, mask = [b.to(device) for b in batch] + forward_start = time.perf_counter() outputs = model(x, mask=mask) if isinstance(outputs, dict): preds = outputs[forward_key] else: preds = outputs loss = loss_fn(preds, y) + forward_time = time.perf_counter() - forward_start + timing["forward"] += forward_time batch_size = y.size(0) loss_value = loss.item() @@ -119,18 +194,30 @@ def evaluate( progress_loss += loss_value progress_batches += 1 - if progress is not None and task_id is not None: - mean_loss = progress_loss / progress_batches - progress.advance(task_id, 1) - progress.update(task_id, description=f"loss={mean_loss:.4f}") - + metrics_start = time.perf_counter() if metrics: for metric in metrics.values(): if hasattr(metric, "update"): metric.update(preds.cpu(), y.cpu()) + metrics_time = time.perf_counter() - metrics_start + timing["metrics"] += metrics_time + + if progress is not None and task_id is not None: + mean_loss = progress_loss / progress_batches + description = f"loss={mean_loss:.4f}" + if profile: + description += _format_timing( + data_time * 1000, + forward_time * 1000, + metrics_ms=metrics_time * 1000, + ) + progress.advance(task_id, 1) + progress.update(task_id, description=description) + + data_start = time.perf_counter() avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 - result = {"loss": avg_loss} + result: Dict[str, float] = {"loss": avg_loss} if metrics: result.update( { @@ -138,4 +225,13 @@ def evaluate( for name, metric in metrics.items() } ) + if profile: + n = progress_batches if progress_batches > 0 else 1 + result.update( + { + "time_data_ms": timing["data"] / n * 1000, + "time_forward_ms": timing["forward"] / n * 1000, + "time_metrics_ms": timing["metrics"] / n * 1000, + } + ) return result diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 06ef7d8..2a52767 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -47,6 +47,7 @@ def __init__( history_path: Optional[str] = None, branch: str = "classifier", progress_bar: bool = False, + profile: bool = False, ): self.model = model self.train_loader = train_loader @@ -61,6 +62,7 @@ def __init__( self.history_path = Path(history_path) if history_path else None self.branch = branch self.progress_bar = progress_bar + self.profile = profile self.history: List[Dict[str, float]] = [] self.should_stop = False @@ -102,6 +104,7 @@ def fit(self) -> List[Dict[str, float]]: branch=self.branch, progress=progress, task_id=train_task, + profile=self.profile, ) val_metrics = evaluate( self.model, @@ -112,6 +115,7 @@ def fit(self) -> List[Dict[str, float]]: branch=self.branch, progress=progress, task_id=val_task, + profile=self.profile, ) else: train_metrics = train_one_epoch( @@ -122,6 +126,7 @@ def fit(self) -> List[Dict[str, float]]: self.device, metrics=self.metrics, branch=self.branch, + profile=self.profile, ) val_metrics = evaluate( self.model, @@ -130,6 +135,7 @@ def fit(self) -> List[Dict[str, float]]: self.device, metrics=self.metrics, branch=self.branch, + profile=self.profile, ) epoch_log = {"epoch": epoch} diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py index bdc1f83..5940f59 100644 --- a/tests/unit/training/pytorch/test_engine.py +++ b/tests/unit/training/pytorch/test_engine.py @@ -102,3 +102,32 @@ def test_train_one_epoch_progress_shows_cumulative_mean(): assert description.startswith("loss=") loss_value = float(description.split("=")[1]) assert loss_value > 0 + + +def test_train_one_epoch_profile_returns_timings(): + """Profiling should return per-section timing averages.""" + model = _make_dummy_classifier() + loader = _make_dummy_loader() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + history = train_one_epoch( + model, loader, loss_fn, optimizer, torch.device("cpu"), profile=True + ) + assert "time_data_ms" in history + assert "time_forward_ms" in history + assert "time_backward_ms" in history + assert "time_optim_ms" in history + assert "time_metrics_ms" in history + assert all(history[k] >= 0 for k in history if k.startswith("time_")) + + +def test_evaluate_profile_returns_timings(): + """Profiling should return per-section timing averages for evaluation.""" + model = _make_dummy_classifier() + loader = _make_dummy_loader() + loss_fn = torch.nn.CrossEntropyLoss() + history = evaluate(model, loader, loss_fn, torch.device("cpu"), profile=True) + assert "time_data_ms" in history + assert "time_forward_ms" in history + assert "time_metrics_ms" in history + assert "time_backward_ms" not in history From 2bad2c21b68c8fc456ae6a5a7b4d78cfa945268c Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 20:57:44 +0200 Subject: [PATCH 46/64] fix(cli): restore jaeger train command and add --profile flag The train command decorators for --progress-bar, -v/--verbose, and the def train(**kwargs): definition were accidentally removed, leaving the docstring dangling under --meta and breaking the CLI. - Restore --progress-bar and -v/--verbose decorators. - Add --profile flag to expose per-section timing in the progress bar. - All 262 tests pass. --- src/jaeger/cli.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/jaeger/cli.py b/src/jaeger/cli.py index be37699..21cb2ff 100644 --- a/src/jaeger/cli.py +++ b/src/jaeger/cli.py @@ -309,6 +309,13 @@ def tune(**kwargs): default=False, help="Show a per-batch Rich progress bar with loss and speed", ) +@click.option( + "--profile", + is_flag=True, + required=False, + default=False, + help="Show per-section timing in the progress bar", +) @click.option( "-v", "--verbose", From 7f504e4eba74e0526baf1d2ed8bb5cbb9c7b04c2 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 21:04:41 +0200 Subject: [PATCH 47/64] fix(training): respect classifier/reliability train and validation steps from config The PyTorch trainer was ignoring classifier_train_steps, classifier_validation_steps, reliability_train_steps, and reliability_validation_steps from the YAML config and always iterated through the entire dataloader each epoch. - Add train_steps and validation_steps parameters to train_one_epoch and evaluate; negative values or None keep the legacy unlimited behavior. - Plumb the same parameters through Trainer.fit, including progress-bar totals. - Read the step counts from the config in commands/train.py and pass them to Trainer for both classifier and reliability branches. - Add unit tests verifying step limiting, -1/unlimited behavior, and config-to-Trainer propagation. 267 tests pass; ruff clean. --- src/jaeger/commands/train.py | 25 +++++++++++ src/jaeger/training/pytorch/engine.py | 22 ++++++++- src/jaeger/training/pytorch/trainer.py | 23 +++++++++- tests/unit/commands/test_train_pytorch.py | 45 +++++++++++++++++++ tests/unit/training/pytorch/test_engine.py | 50 +++++++++++++++++++++ tests/unit/training/pytorch/test_trainer.py | 40 +++++++++++++++++ 6 files changed, 201 insertions(+), 4 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 0d9da78..8714566 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -59,6 +59,17 @@ def _count_parameters(model: nn.Module) -> Tuple[int, int]: return total, trainable +def _resolve_steps(value: Optional[int]) -> Optional[int]: + """Convert a config step count to a value understood by the engine. + + ``None`` or negative values mean "run until the dataloader is exhausted". + Zero or positive values limit the loop to that many batches. + """ + if value is None or value < 0: + return None + return int(value) + + def _initialize_lazy_layers( models: Dict[str, nn.Module], config: Dict[str, Any] ) -> None: @@ -383,6 +394,10 @@ def train_fragment_core(**kwargs): callbacks = _build_callbacks(train_cfg) epochs = int(train_cfg.get("classifier_epochs", 1)) + train_steps = _resolve_steps(train_cfg.get("classifier_train_steps")) + validation_steps = _resolve_steps( + train_cfg.get("classifier_validation_steps") + ) if kwargs.get("only_save", False): logger.info("Skipping classifier training (--only_save)") @@ -402,6 +417,8 @@ def train_fragment_core(**kwargs): branch="classifier", progress_bar=kwargs.get("progress_bar", False), profile=kwargs.get("profile", False), + train_steps=train_steps, + validation_steps=validation_steps, ) trainer.fit() @@ -444,6 +461,12 @@ def train_fragment_core(**kwargs): rel_callbacks = _build_callbacks(train_cfg) rel_epochs = int(train_cfg.get("reliability_epochs", 1)) + rel_train_steps = _resolve_steps( + train_cfg.get("reliability_train_steps") + ) + rel_validation_steps = _resolve_steps( + train_cfg.get("reliability_validation_steps") + ) rel_trainer = Trainer( model=rel_pipeline, @@ -460,6 +483,8 @@ def train_fragment_core(**kwargs): branch="reliability", progress_bar=kwargs.get("progress_bar", False), profile=kwargs.get("profile", False), + train_steps=rel_train_steps, + validation_steps=rel_validation_steps, ) rel_trainer.fit() diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py index 0909b14..8a46e7c 100644 --- a/src/jaeger/training/pytorch/engine.py +++ b/src/jaeger/training/pytorch/engine.py @@ -26,11 +26,16 @@ def train_one_epoch( progress: Optional[Any] = None, task_id: Optional[Any] = None, profile: bool = False, + train_steps: Optional[int] = None, ) -> Dict[str, float]: """Train for one epoch and return averaged loss and metrics. When ``profile=True``, per-section timings are included in the progress bar description and returned under keys prefixed with ``time_``. + + If ``train_steps`` is a non-negative integer, training stops after that + many batches even if the dataloader has more data. Negative values or + ``None`` mean iterate until the dataloader is exhausted. """ model.train() total_loss = 0.0 @@ -50,7 +55,9 @@ def train_one_epoch( m.reset() data_start = time.perf_counter() - for batch in dataloader: + for batch_idx, batch in enumerate(dataloader): + if train_steps is not None and train_steps >= 0 and batch_idx >= train_steps: + break data_time = time.perf_counter() - data_start timing["data"] += data_time @@ -153,11 +160,16 @@ def evaluate( progress: Optional[Any] = None, task_id: Optional[Any] = None, profile: bool = False, + validation_steps: Optional[int] = None, ) -> Dict[str, float]: """Evaluate and return averaged loss and metrics. When ``profile=True``, per-section timings are included in the progress bar description and returned under keys prefixed with ``time_``. + + If ``validation_steps`` is a non-negative integer, evaluation stops after + that many batches even if the dataloader has more data. Negative values or + ``None`` mean iterate until the dataloader is exhausted. """ model.eval() total_loss = 0.0 @@ -172,7 +184,13 @@ def evaluate( with torch.no_grad(): data_start = time.perf_counter() - for batch in dataloader: + for batch_idx, batch in enumerate(dataloader): + if ( + validation_steps is not None + and validation_steps >= 0 + and batch_idx >= validation_steps + ): + break data_time = time.perf_counter() - data_start timing["data"] += data_time diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 2a52767..6bffbc9 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -48,6 +48,8 @@ def __init__( branch: str = "classifier", progress_bar: bool = False, profile: bool = False, + train_steps: Optional[int] = None, + validation_steps: Optional[int] = None, ): self.model = model self.train_loader = train_loader @@ -63,6 +65,8 @@ def __init__( self.branch = branch self.progress_bar = progress_bar self.profile = profile + self.train_steps = train_steps + self.validation_steps = validation_steps self.history: List[Dict[str, float]] = [] self.should_stop = False @@ -77,6 +81,17 @@ def fit(self) -> List[Dict[str, float]]: if hasattr(callback, "on_epoch_begin"): callback.on_epoch_begin(self, epoch) + train_total = ( + self.train_steps + if self.train_steps is not None + else self._loader_length(self.train_loader) + ) + val_total = ( + self.validation_steps + if self.validation_steps is not None + else self._loader_length(self.val_loader) + ) + if self.progress_bar: with Progress( TextColumn("[progress.description]{task.description}"), @@ -88,11 +103,11 @@ def fit(self) -> List[Dict[str, float]]: ) as progress: train_task = progress.add_task( f"train epoch {epoch}/{self.epochs}", - total=self._loader_length(self.train_loader), + total=train_total, ) val_task = progress.add_task( f"val epoch {epoch}/{self.epochs}", - total=self._loader_length(self.val_loader), + total=val_total, ) train_metrics = train_one_epoch( self.model, @@ -105,6 +120,7 @@ def fit(self) -> List[Dict[str, float]]: progress=progress, task_id=train_task, profile=self.profile, + train_steps=self.train_steps, ) val_metrics = evaluate( self.model, @@ -116,6 +132,7 @@ def fit(self) -> List[Dict[str, float]]: progress=progress, task_id=val_task, profile=self.profile, + validation_steps=self.validation_steps, ) else: train_metrics = train_one_epoch( @@ -127,6 +144,7 @@ def fit(self) -> List[Dict[str, float]]: metrics=self.metrics, branch=self.branch, profile=self.profile, + train_steps=self.train_steps, ) val_metrics = evaluate( self.model, @@ -136,6 +154,7 @@ def fit(self) -> List[Dict[str, float]]: metrics=self.metrics, branch=self.branch, profile=self.profile, + validation_steps=self.validation_steps, ) epoch_log = {"epoch": epoch} diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index 9463f86..4c8d8f9 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -173,3 +173,48 @@ def test_train_fragment_core_progress_bar(tmp_path): meta=None, progress_bar=True, ) + + +def test_train_fragment_core_respects_config_steps(tmp_path, monkeypatch): + """``train_fragment_core`` should pass config train/validation steps to Trainer.""" + from jaeger.training.pytorch import trainer as trainer_module + + captured = {} + + original_fit = trainer_module.Trainer.fit + + def _capturing_fit(self): + captured["train_steps"] = self.train_steps + captured["validation_steps"] = self.validation_steps + return original_fit(self) + + monkeypatch.setattr(trainer_module.Trainer, "fit", _capturing_fit) + + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config["training"]["classifier_train_steps"] = 2 + config["training"]["classifier_validation_steps"] = 1 + config_path.write_text(yaml.safe_dump(config)) + + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + assert captured["train_steps"] == 2 + assert captured["validation_steps"] == 1 diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py index 5940f59..cdd3e92 100644 --- a/tests/unit/training/pytorch/test_engine.py +++ b/tests/unit/training/pytorch/test_engine.py @@ -131,3 +131,53 @@ def test_evaluate_profile_returns_timings(): assert "time_forward_ms" in history assert "time_metrics_ms" in history assert "time_backward_ms" not in history + + +class _CounterClassifier(DummyClassifier): + """Classifier that counts forward calls.""" + + def __init__(self): + super().__init__() + self.forward_count = 0 + + def forward(self, x, mask=None): + self.forward_count += 1 + return super().forward(x, mask) + + +def test_train_one_epoch_respects_train_steps(): + """train_one_epoch should stop after train_steps batches.""" + model = _CounterClassifier() + loader = _make_dummy_loader(n=20) # 5 batches of 4 + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + history = train_one_epoch( + model, loader, loss_fn, optimizer, torch.device("cpu"), train_steps=2 + ) + assert "loss" in history + assert model.forward_count == 2 + + +def test_train_one_epoch_negative_train_steps_runs_all(): + """Negative train_steps should run through the whole dataloader.""" + model = _CounterClassifier() + loader = _make_dummy_loader(n=20) # 5 batches + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + history = train_one_epoch( + model, loader, loss_fn, optimizer, torch.device("cpu"), train_steps=-1 + ) + assert "loss" in history + assert model.forward_count == 5 + + +def test_evaluate_respects_validation_steps(): + """evaluate should stop after validation_steps batches.""" + model = _CounterClassifier() + loader = _make_dummy_loader(n=20) # 5 batches of 4 + loss_fn = torch.nn.CrossEntropyLoss() + history = evaluate( + model, loader, loss_fn, torch.device("cpu"), validation_steps=2 + ) + assert "loss" in history + assert model.forward_count == 2 diff --git a/tests/unit/training/pytorch/test_trainer.py b/tests/unit/training/pytorch/test_trainer.py index 01f2c14..b08d891 100644 --- a/tests/unit/training/pytorch/test_trainer.py +++ b/tests/unit/training/pytorch/test_trainer.py @@ -87,3 +87,43 @@ def test_trainer_fit_with_progress_bar(): assert len(history) == 1 assert "train_loss" in history[0] assert "val_loss" in history[0] + + +class _CounterModel(_DummyClassifier): + """Model that counts forward calls.""" + + def __init__(self): + super().__init__() + self.forward_count = 0 + + def forward(self, x, mask=None): + self.forward_count += 1 + return super().forward(x, mask) + + +def test_trainer_respects_train_and_validation_steps(): + """Trainer should pass train_steps/validation_steps to the engine loops.""" + model = _CounterModel() + train_loader = _make_dummy_data(n=20) # 5 batches + val_loader = _make_dummy_data(n=12) # 3 batches + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + checkpoint_dir=tmpdir, + history_path=Path(tmpdir) / "history.json", + train_steps=2, + validation_steps=1, + ) + history = trainer.fit() + assert len(history) == 2 + # Each epoch runs train_steps + validation_steps forward passes. + assert model.forward_count == 2 * (2 + 1) From 9eeea43572f40bdba5d418a082d3d67dfb1ce603 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 21:14:42 +0200 Subject: [PATCH 48/64] fix(training): synchronize CUDA during profiling for accurate timings The per-section timings shown with --profile measured only CPU-side kernel-launch latency, not actual GPU execution time. Because CUDA ops are asynchronous, the real forward/backward/optimizer work completed later at implicit sync points (e.g. loss.item()), leaving most wall-clock time unaccounted (e.g. validation reported ~2 ms of section time vs ~40 ms/batch real average). - Add _maybe_synchronize() helper that calls torch.cuda.synchronize() only when profile=True and device.type == 'cuda'. - Place sync boundaries around data, forward, backward, optimizer, and metrics sections in train_one_epoch and evaluate so each measured interval includes the real GPU work. - Non-profiling runs and CPU runs are unaffected. - Add unit tests for _maybe_synchronize behavior. 270 tests pass; ruff clean. --- src/jaeger/training/pytorch/engine.py | 17 +++++++++ tests/unit/training/pytorch/test_engine.py | 43 +++++++++++++++++++++- 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/jaeger/training/pytorch/engine.py b/src/jaeger/training/pytorch/engine.py index 8a46e7c..0f1ec1d 100644 --- a/src/jaeger/training/pytorch/engine.py +++ b/src/jaeger/training/pytorch/engine.py @@ -13,6 +13,12 @@ def _format_timing(data_ms: float, forward_ms: float, backward_ms: float = 0.0, return f" d={data_ms:.1f}ms f={forward_ms:.1f}ms m={metrics_ms:.1f}ms" +def _maybe_synchronize(device: torch.device, profile: bool) -> None: + """Synchronize CUDA when profiling so timings measure real GPU work.""" + if profile and device.type == "cuda": + torch.cuda.synchronize() + + def train_one_epoch( model: nn.Module, dataloader: DataLoader, @@ -54,6 +60,7 @@ def train_one_epoch( if hasattr(m, "reset"): m.reset() + _maybe_synchronize(device, profile) data_start = time.perf_counter() for batch_idx, batch in enumerate(dataloader): if train_steps is not None and train_steps >= 0 and batch_idx >= train_steps: @@ -71,6 +78,7 @@ def train_one_epoch( else: preds = outputs loss = loss_fn(preds, y) + _maybe_synchronize(device, profile) forward_time = time.perf_counter() - forward_start timing["forward"] += forward_time elif branch == "reliability": @@ -80,6 +88,7 @@ def train_one_epoch( outputs = model(x, mask=mask) preds = outputs[forward_key] loss = loss_fn(preds, y) + _maybe_synchronize(device, profile) forward_time = time.perf_counter() - forward_start timing["forward"] += forward_time else: @@ -87,11 +96,13 @@ def train_one_epoch( backward_start = time.perf_counter() loss.backward() + _maybe_synchronize(device, profile) backward_time = time.perf_counter() - backward_start timing["backward"] += backward_time optim_start = time.perf_counter() optimizer.step() + _maybe_synchronize(device, profile) optim_time = time.perf_counter() - optim_start timing["optim"] += optim_time @@ -107,6 +118,7 @@ def train_one_epoch( for metric in metrics.values(): if hasattr(metric, "update"): metric.update(preds.detach().cpu(), y.detach().cpu()) + _maybe_synchronize(device, profile) metrics_time = time.perf_counter() - metrics_start timing["metrics"] += metrics_time @@ -124,6 +136,7 @@ def train_one_epoch( progress.advance(task_id, 1) progress.update(task_id, description=description) + _maybe_synchronize(device, profile) data_start = time.perf_counter() avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 @@ -183,6 +196,7 @@ def evaluate( m.reset() with torch.no_grad(): + _maybe_synchronize(device, profile) data_start = time.perf_counter() for batch_idx, batch in enumerate(dataloader): if ( @@ -202,6 +216,7 @@ def evaluate( else: preds = outputs loss = loss_fn(preds, y) + _maybe_synchronize(device, profile) forward_time = time.perf_counter() - forward_start timing["forward"] += forward_time @@ -217,6 +232,7 @@ def evaluate( for metric in metrics.values(): if hasattr(metric, "update"): metric.update(preds.cpu(), y.cpu()) + _maybe_synchronize(device, profile) metrics_time = time.perf_counter() - metrics_start timing["metrics"] += metrics_time @@ -232,6 +248,7 @@ def evaluate( progress.advance(task_id, 1) progress.update(task_id, description=description) + _maybe_synchronize(device, profile) data_start = time.perf_counter() avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 diff --git a/tests/unit/training/pytorch/test_engine.py b/tests/unit/training/pytorch/test_engine.py index cdd3e92..ef8c782 100644 --- a/tests/unit/training/pytorch/test_engine.py +++ b/tests/unit/training/pytorch/test_engine.py @@ -1,7 +1,7 @@ import torch from rich.progress import Progress -from jaeger.training.pytorch.engine import evaluate, train_one_epoch +from jaeger.training.pytorch.engine import _maybe_synchronize, evaluate, train_one_epoch class DummyClassifier(torch.nn.Module): @@ -181,3 +181,44 @@ def test_evaluate_respects_validation_steps(): ) assert "loss" in history assert model.forward_count == 2 + + +def test_maybe_synchronize_no_op_on_cpu(): + """_maybe_synchronize should not call torch.cuda.synchronize on CPU.""" + _maybe_synchronize(torch.device("cpu"), profile=True) + + +def test_maybe_synchronize_no_op_when_not_profiling(): + """_maybe_synchronize should not call torch.cuda.synchronize when not profiling.""" + class _FakeCuda: + called = False + + @staticmethod + def synchronize(): + _FakeCuda.called = True + + original = torch.cuda.synchronize + try: + torch.cuda.synchronize = _FakeCuda.synchronize + _maybe_synchronize(torch.device("cuda:0"), profile=False) + assert not _FakeCuda.called + finally: + torch.cuda.synchronize = original + + +def test_maybe_synchronize_syncs_when_profiling_on_cuda(): + """_maybe_synchronize should call torch.cuda.synchronize when profiling on CUDA.""" + class _FakeCuda: + called = False + + @staticmethod + def synchronize(): + _FakeCuda.called = True + + original = torch.cuda.synchronize + try: + torch.cuda.synchronize = _FakeCuda.synchronize + _maybe_synchronize(torch.device("cuda:0"), profile=True) + assert _FakeCuda.called + finally: + torch.cuda.synchronize = original From 84afcf762534371f387da53aa698aec3d6dcc1b7 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 21:25:59 +0200 Subject: [PATCH 49/64] fix(train): guard against overwriting existing checkpoints Training would silently write into an existing checkpoint directory, potentially overwriting prior work. Users must now explicitly choose to overwrite with --force or resume with --from_last_checkpoint. - Add _has_checkpoints() and _guard_against_overwrite() helpers. - Before classifier/reliability training, raise a ClickException if the checkpoint directory already contains .pt files and neither --force nor --from_last_checkpoint is set. - --force continues to delete existing checkpoint directories first. - --from_last_checkpoint continues to load the latest .pt checkpoint and train. - Add tests for refusal, --force overwrite, and --from_last_checkpoint resume paths. 273 tests pass; ruff clean. --- src/jaeger/commands/train.py | 38 +++++++ tests/unit/commands/test_train_pytorch.py | 120 ++++++++++++++++++++++ 2 files changed, 158 insertions(+) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 8714566..b4107cd 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -199,6 +199,32 @@ def _find_latest_checkpoint(directory: Path) -> Optional[Path]: return checkpoints[-1] if checkpoints else None +def _has_checkpoints(directory: Path) -> bool: + """Return True if *directory* exists and contains any ``.pt`` files.""" + if not directory.exists(): + return False + return any(directory.glob("*.pt")) + + +def _guard_against_overwrite( + directory: Path, + force: bool, + from_last_checkpoint: bool, + branch_name: str, +) -> None: + """Raise an error if existing checkpoints would be overwritten. + + Existing checkpoints are allowed when resuming via ``from_last_checkpoint`` + or when explicitly overwriting via ``force``. + """ + if force or from_last_checkpoint or not _has_checkpoints(directory): + return + raise click.ClickException( + f"{branch_name} checkpoint directory already exists: {directory}. " + "Use --force to overwrite or --from_last_checkpoint to resume training." + ) + + def _save_model_checkpoint( model: nn.Module, save_dir: Path, @@ -402,6 +428,12 @@ def train_fragment_core(**kwargs): if kwargs.get("only_save", False): logger.info("Skipping classifier training (--only_save)") else: + _guard_against_overwrite( + classifier_dir, + force=kwargs.get("force", False), + from_last_checkpoint=kwargs.get("from_last_checkpoint", False), + branch_name="Classifier", + ) trainer = Trainer( model=model, train_loader=classifier_loaders["train"], @@ -438,6 +470,12 @@ def train_fragment_core(**kwargs): reliability_loaders = None if reliability_loaders is not None and not kwargs.get("only_save", False): + _guard_against_overwrite( + reliability_dir, + force=kwargs.get("force", False), + from_last_checkpoint=kwargs.get("from_last_checkpoint", False), + branch_name="Reliability", + ) rel_pipeline = _ReliabilityPipeline( rep_model, models["reliability_head"] ) diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index 4c8d8f9..aecf3c4 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -2,7 +2,9 @@ from __future__ import annotations +import click import numpy as np +import pytest import yaml from jaeger.commands.train import train_fragment_core @@ -218,3 +220,121 @@ def _capturing_fit(self): assert captured["train_steps"] == 2 assert captured["validation_steps"] == 1 + + +def test_train_fragment_core_refuses_overwrite_without_force_or_resume(tmp_path): + """``train_fragment_core`` should stop if checkpoints exist and neither + ``--force`` nor ``--from_last_checkpoint`` is given.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + classifier_dir = tmp_path / "checkpoints" / "classifier" + classifier_dir.mkdir(parents=True, exist_ok=True) + (classifier_dir / "checkpoint_epoch_1.pt").write_bytes(b"") + + with pytest.raises(click.ClickException): + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + +def test_train_fragment_core_force_overwrites_existing_checkpoints(tmp_path): + """``--force`` should remove existing checkpoints and train.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + classifier_dir = tmp_path / "checkpoints" / "classifier" + classifier_dir.mkdir(parents=True, exist_ok=True) + (classifier_dir / "checkpoint_epoch_1.pt").write_bytes(b"") + + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=True, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + # Directory was removed and a fresh checkpoint was written. + assert len(list(classifier_dir.glob("checkpoint_epoch_*.pt"))) == 1 + + +def test_train_fragment_core_resume_from_last_checkpoint(tmp_path): + """``--from_last_checkpoint`` should load the latest checkpoint and train.""" + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + # First training run creates a checkpoint. + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + checkpoints = list((tmp_path / "checkpoints" / "classifier").glob("*.pt")) + assert len(checkpoints) == 1 + + # Second run with from_last_checkpoint should not raise and should load + # the existing checkpoint before training. + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=True, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + checkpoints = list((tmp_path / "checkpoints" / "classifier").glob("*.pt")) + assert len(checkpoints) >= 1 From 035e9e750e7260b0cad71101122d821829e92782 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 21:31:23 +0200 Subject: [PATCH 50/64] fix(train): resume from checkpoint without device mismatch Loading a checkpoint on CPU then moving the model to CUDA left the optimizer's param_groups pointing at stale CPU tensors and the optimizer state on CPU. optimizer.step() then crashed with: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! - Move rep_model, jaeger_classifier, and reliability_head to the target device immediately after lazy-layer initialization and before building optimizers. - Pass the target device into _load_checkpoint_if_requested() and use map_location=device so checkpoint tensors land directly on the target device. - Add a unit test verifying resumed optimizer state is mapped to the requested device. 274 tests pass; ruff clean. --- src/jaeger/commands/train.py | 21 ++++++++++--- tests/unit/commands/test_train_pytorch.py | 38 ++++++++++++++++++++++- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index b4107cd..501ffa1 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -110,8 +110,13 @@ def _load_checkpoint_if_requested( model: nn.Module, optimizer: Optional[torch.optim.Optimizer], checkpoint_path: Optional[str | Path], + device: torch.device, ) -> None: - """Load model and optimizer state dicts from *checkpoint_path* if provided.""" + """Load model and optimizer state dicts from *checkpoint_path* if provided. + + The checkpoint is mapped to *device* so that both model parameters and + optimizer state land on the target device and stay in sync. + """ if checkpoint_path is None: return path = Path(checkpoint_path) @@ -119,7 +124,7 @@ def _load_checkpoint_if_requested( logger.warning("Checkpoint path does not exist: %s", path) return logger.info("Loading checkpoint from %s", path) - checkpoint = torch.load(path, map_location="cpu", weights_only=True) + checkpoint = torch.load(path, map_location=device, weights_only=True) model.load_state_dict(checkpoint["model_state_dict"]) if optimizer is not None and "optimizer_state_dict" in checkpoint: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) @@ -370,6 +375,14 @@ def train_fragment_core(**kwargs): ) model_num_params = numerize(total_params, decimal=1) + # Move models to the target device before building optimizers so that + # optimizer param_groups and any resumed optimizer state live on the + # same device as the model. + rep_model.to(device) + models["jaeger_classifier"].to(device) + if models.get("reliability_head") is not None: + models["reliability_head"].to(device) + train_cfg = config.get("training", {}) classifier_dir = Path(train_cfg.get("classifier_dir", "checkpoints/classifier")) reliability_dir = Path( @@ -410,7 +423,7 @@ def train_fragment_core(**kwargs): checkpoint_path = None if kwargs.get("from_last_checkpoint", False): checkpoint_path = _find_latest_checkpoint(classifier_dir) - _load_checkpoint_if_requested(model, optimizer, checkpoint_path) + _load_checkpoint_if_requested(model, optimizer, checkpoint_path, device) if kwargs.get("only_classification_head", False) or kwargs.get( "only_heads", False @@ -494,7 +507,7 @@ def train_fragment_core(**kwargs): if kwargs.get("from_last_checkpoint", False): rel_checkpoint_path = _find_latest_checkpoint(reliability_dir) _load_checkpoint_if_requested( - rel_pipeline, rel_optimizer, rel_checkpoint_path + rel_pipeline, rel_optimizer, rel_checkpoint_path, device ) rel_callbacks = _build_callbacks(train_cfg) diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index aecf3c4..9308c8c 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -5,9 +5,10 @@ import click import numpy as np import pytest +import torch import yaml -from jaeger.commands.train import train_fragment_core +from jaeger.commands.train import _load_checkpoint_if_requested, train_fragment_core def _make_raw_npz(path, n_samples, seq_length=50): @@ -338,3 +339,38 @@ def test_train_fragment_core_resume_from_last_checkpoint(tmp_path): checkpoints = list((tmp_path / "checkpoints" / "classifier").glob("*.pt")) assert len(checkpoints) >= 1 + + +def test_load_checkpoint_maps_optimizer_state_to_device(tmp_path): + """Resumed optimizer state should live on the target device.""" + model = torch.nn.Linear(4, 2) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + # Run one step so the optimizer has state. + x = torch.randn(2, 4) + loss = model(x).sum() + loss.backward() + optimizer.step() + + checkpoint_path = tmp_path / "checkpoint.pt" + torch.save( + { + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + }, + checkpoint_path, + ) + + # Load onto CPU (tests in CI may not have CUDA). + target_device = torch.device("cpu") + fresh_model = torch.nn.Linear(4, 2) + fresh_model.to(target_device) + fresh_optimizer = torch.optim.Adam(fresh_model.parameters(), lr=1e-3) + _load_checkpoint_if_requested( + fresh_model, fresh_optimizer, checkpoint_path, target_device + ) + + # Verify optimizer state tensors are on the target device. + for state in fresh_optimizer.state.values(): + for value in state.values(): + if isinstance(value, torch.Tensor): + assert value.device == target_device From 13f49a737f6eea08f91e5b92ef7289f5c914c81b Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 21:40:00 +0200 Subject: [PATCH 51/64] feat(nnlib): support any torch.optim optimizer by name Replace the hardcoded adam/sgd/rmsprop map with a case-insensitive lookup over torch.optim classes. This makes AdamW, NAdam, RAdam, and any future PyTorch optimizer available via training.optimizer. - Empty/missing optimizer names still default to Adam. - Unknown names raise a ValueError listing available optimizers. - Add unit tests for AdamW resolution, case-insensitivity, default, and unknown-name handling. --- src/jaeger/nnlib/pytorch/builder.py | 23 +++++++++++++------ tests/unit/nnlib/pytorch/test_builder.py | 29 ++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py index 27f0a07..3a5a186 100644 --- a/src/jaeger/nnlib/pytorch/builder.py +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -195,11 +195,20 @@ def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: return metrics def _get_optimizer(self, name: str) -> torch.optim.Optimizer: - optimizers = { - "adam": torch.optim.Adam, - "sgd": torch.optim.SGD, - "rmsprop": torch.optim.RMSprop, + """Resolve a PyTorch optimizer class by case-insensitive name.""" + if not name: + name = "adam" + + available = { + cls_name.lower(): cls + for cls_name, cls in vars(torch.optim).items() + if isinstance(cls, type) and issubclass(cls, torch.optim.Optimizer) } - if name not in optimizers: - raise ValueError(f"Unsupported optimizer: {name}") - return optimizers[name] + + opt_class = available.get(name.lower()) + if opt_class is None: + raise ValueError( + f"Unsupported optimizer: {name}. " + f"Available torch.optim optimizers include: {', '.join(sorted(available))}" + ) + return opt_class diff --git a/tests/unit/nnlib/pytorch/test_builder.py b/tests/unit/nnlib/pytorch/test_builder.py index a07901d..465a64e 100644 --- a/tests/unit/nnlib/pytorch/test_builder.py +++ b/tests/unit/nnlib/pytorch/test_builder.py @@ -70,6 +70,35 @@ def test_builder_compile_classifier(): assert isinstance(loss, torch.nn.CrossEntropyLoss) +def test_builder_resolve_adamw_optimizer(): + """Generic torch.optim lookup should resolve AdamW.""" + builder = ModelBuilder(_minimal_config()) + opt_class = builder._get_optimizer("AdamW") + assert opt_class is torch.optim.AdamW + + +def test_builder_resolve_optimizer_case_insensitive(): + """Optimizer names should be case-insensitive.""" + builder = ModelBuilder(_minimal_config()) + assert builder._get_optimizer("adam") is torch.optim.Adam + assert builder._get_optimizer("ADAM") is torch.optim.Adam + assert builder._get_optimizer("Adam") is torch.optim.Adam + + +def test_builder_default_optimizer_is_adam(): + """Empty or missing optimizer name should default to Adam.""" + builder = ModelBuilder(_minimal_config()) + assert builder._get_optimizer("") is torch.optim.Adam + assert builder._get_optimizer(None) is torch.optim.Adam + + +def test_builder_unknown_optimizer_raises(): + """An unknown optimizer name should raise a clear ValueError.""" + builder = ModelBuilder(_minimal_config()) + with pytest.raises(ValueError, match="Unsupported optimizer: not_an_optimizer"): + builder._get_optimizer("not_an_optimizer") + + def test_builder_missing_classifier_raises(): config = _minimal_config() del config["model"]["classifier"] From 75aa876f357f3e2b01d20913ad459b724d1708bf Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 21:48:07 +0200 Subject: [PATCH 52/64] fix(train): resume from checkpoint respects last completed epoch --from_last_checkpoint previously always started training from epoch 1, ignoring the epoch stored in the checkpoint. This caused it to re-train already-completed epochs or overwrite previous checkpoints. - _load_checkpoint_if_requested() now returns the epoch stored in the checkpoint (default 0). - Trainer accepts a start_epoch argument and trains from start_epoch + 1 through classifier_epochs. - If start_epoch >= epochs, training is skipped and a message is logged. - Add tests for resumed epoch ranges and skipping when no epochs remain. 281 tests pass; ruff clean. --- src/jaeger/commands/train.py | 18 +++++-- src/jaeger/training/pytorch/trainer.py | 19 ++++++- tests/unit/commands/test_train_pytorch.py | 59 +++++++++++++++++++++ tests/unit/training/pytorch/test_trainer.py | 55 +++++++++++++++++++ 4 files changed, 145 insertions(+), 6 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index 501ffa1..b76e26d 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -111,23 +111,27 @@ def _load_checkpoint_if_requested( optimizer: Optional[torch.optim.Optimizer], checkpoint_path: Optional[str | Path], device: torch.device, -) -> None: +) -> int: """Load model and optimizer state dicts from *checkpoint_path* if provided. The checkpoint is mapped to *device* so that both model parameters and optimizer state land on the target device and stay in sync. + + Returns the epoch stored in the checkpoint, or ``0`` if no checkpoint was + loaded or no epoch was recorded. """ if checkpoint_path is None: - return + return 0 path = Path(checkpoint_path) if not path.exists(): logger.warning("Checkpoint path does not exist: %s", path) - return + return 0 logger.info("Loading checkpoint from %s", path) checkpoint = torch.load(path, map_location=device, weights_only=True) model.load_state_dict(checkpoint["model_state_dict"]) if optimizer is not None and "optimizer_state_dict" in checkpoint: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + return int(checkpoint.get("epoch", 0)) def _normalize_optimizer_params(params: Optional[Dict[str, Any]]) -> Dict[str, Any]: @@ -423,7 +427,9 @@ def train_fragment_core(**kwargs): checkpoint_path = None if kwargs.get("from_last_checkpoint", False): checkpoint_path = _find_latest_checkpoint(classifier_dir) - _load_checkpoint_if_requested(model, optimizer, checkpoint_path, device) + start_epoch = _load_checkpoint_if_requested( + model, optimizer, checkpoint_path, device + ) if kwargs.get("only_classification_head", False) or kwargs.get( "only_heads", False @@ -464,6 +470,7 @@ def train_fragment_core(**kwargs): profile=kwargs.get("profile", False), train_steps=train_steps, validation_steps=validation_steps, + start_epoch=start_epoch, ) trainer.fit() @@ -506,7 +513,7 @@ def train_fragment_core(**kwargs): rel_checkpoint_path = None if kwargs.get("from_last_checkpoint", False): rel_checkpoint_path = _find_latest_checkpoint(reliability_dir) - _load_checkpoint_if_requested( + rel_start_epoch = _load_checkpoint_if_requested( rel_pipeline, rel_optimizer, rel_checkpoint_path, device ) @@ -536,6 +543,7 @@ def train_fragment_core(**kwargs): profile=kwargs.get("profile", False), train_steps=rel_train_steps, validation_steps=rel_validation_steps, + start_epoch=rel_start_epoch, ) rel_trainer.fit() diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 6bffbc9..a7f321b 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -17,6 +17,10 @@ from torch.utils.data import DataLoader from jaeger.training.pytorch.engine import evaluate, train_one_epoch +from jaeger.utils.logging import get_logger + + +logger = get_logger(log_file=None, log_path=None, level=3) class _SpeedMsColumn(ProgressColumn): @@ -50,6 +54,7 @@ def __init__( profile: bool = False, train_steps: Optional[int] = None, validation_steps: Optional[int] = None, + start_epoch: int = 0, ): self.model = model self.train_loader = train_loader @@ -67,6 +72,7 @@ def __init__( self.profile = profile self.train_steps = train_steps self.validation_steps = validation_steps + self.start_epoch = start_epoch self.history: List[Dict[str, float]] = [] self.should_stop = False @@ -76,7 +82,18 @@ def fit(self) -> List[Dict[str, float]]: if hasattr(callback, "on_train_begin"): callback.on_train_begin(self) - for epoch in range(1, self.epochs + 1): + if self.start_epoch >= self.epochs: + logger.info( + "Checkpoint epoch %d is >= requested epochs %d; nothing to train", + self.start_epoch, + self.epochs, + ) + for callback in self.callbacks: + if hasattr(callback, "on_train_end"): + callback.on_train_end(self) + return self.history + + for epoch in range(self.start_epoch + 1, self.epochs + 1): for callback in self.callbacks: if hasattr(callback, "on_epoch_begin"): callback.on_epoch_begin(self, epoch) diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index 9308c8c..f8f5c2f 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -341,6 +341,65 @@ def test_train_fragment_core_resume_from_last_checkpoint(tmp_path): assert len(checkpoints) >= 1 +def test_train_fragment_core_resume_respects_checkpoint_epoch(tmp_path, monkeypatch): + """``--from_last_checkpoint`` should pass the checkpoint epoch to Trainer.""" + from jaeger.training.pytorch import trainer as trainer_module + + captured = {} + + original_fit = trainer_module.Trainer.fit + + def _capturing_fit(self): + captured["start_epoch"] = self.start_epoch + return original_fit(self) + + monkeypatch.setattr(trainer_module.Trainer, "fit", _capturing_fit) + + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config["training"]["classifier_epochs"] = 2 + config_path.write_text(yaml.safe_dump(config)) + + # First run creates a checkpoint at epoch 2. + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + # Resume and capture the start_epoch passed to Trainer. + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=True, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=False, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + assert captured["start_epoch"] == 2 + + def test_load_checkpoint_maps_optimizer_state_to_device(tmp_path): """Resumed optimizer state should live on the target device.""" model = torch.nn.Linear(4, 2) diff --git a/tests/unit/training/pytorch/test_trainer.py b/tests/unit/training/pytorch/test_trainer.py index b08d891..4ff29e5 100644 --- a/tests/unit/training/pytorch/test_trainer.py +++ b/tests/unit/training/pytorch/test_trainer.py @@ -127,3 +127,58 @@ def test_trainer_respects_train_and_validation_steps(): assert len(history) == 2 # Each epoch runs train_steps + validation_steps forward passes. assert model.forward_count == 2 * (2 + 1) + + +def test_trainer_resumes_from_start_epoch(): + """Trainer should start training from start_epoch + 1.""" + model = _CounterModel() + train_loader = _make_dummy_data() + val_loader = _make_dummy_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=5, + device=torch.device("cpu"), + checkpoint_dir=tmpdir, + history_path=Path(tmpdir) / "history.json", + start_epoch=2, + ) + history = trainer.fit() + # start_epoch=2 means epochs 3, 4, 5 are trained. + assert len(history) == 3 + assert [entry["epoch"] for entry in history] == [3, 4, 5] + + +def test_trainer_skips_training_when_start_epoch_reaches_epochs(): + """Trainer should not train if start_epoch >= epochs.""" + model = _CounterModel() + train_loader = _make_dummy_data() + val_loader = _make_dummy_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=3, + device=torch.device("cpu"), + checkpoint_dir=tmpdir, + history_path=Path(tmpdir) / "history.json", + start_epoch=3, + ) + history = trainer.fit() + assert len(history) == 0 + assert model.forward_count == 0 + # No new checkpoint should be saved. + assert len(list(Path(tmpdir).glob("checkpoint_epoch_*.pt"))) == 0 From 6dc756ea2b82e0f5e4c215bb9cfce62717bd9afd Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 22:44:16 +0200 Subject: [PATCH 53/64] feat(training): add ReduceLROnPlateau and TerminateOnNaN callbacks Add two commonly requested PyTorch-style callbacks and wire them into the config-based callback builder. - ReduceLROnPlateau: lowers LR by factor after patience epochs of no improvement on monitor, down to min_lr. - TerminateOnNaN: stops training when the monitored metric becomes NaN or Inf. - Update _SUPPORTED_CALLBACKS and _build_callbacks in commands/train.py to instantiate them from YAML config entries. - Add unit tests for both callbacks. 283 tests pass; ruff clean. --- src/jaeger/commands/train.py | 24 ++++- src/jaeger/training/pytorch/callbacks.py | 90 +++++++++++++++++++ tests/unit/training/pytorch/test_callbacks.py | 75 +++++++++++++++- 3 files changed, 186 insertions(+), 3 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index b76e26d..daafa9e 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -24,7 +24,12 @@ from jaeger.dataops.pytorch.builders import build_datasets from jaeger.nnlib.pytorch.builder import ModelBuilder -from jaeger.training.pytorch.callbacks import EarlyStopping, ModelCheckpoint +from jaeger.training.pytorch.callbacks import ( + EarlyStopping, + ModelCheckpoint, + ReduceLROnPlateau, + TerminateOnNaN, +) from jaeger.training.pytorch.distributed import ( cleanup_distributed, get_device, @@ -158,7 +163,7 @@ def _callback_path_from_config( return None -_SUPPORTED_CALLBACKS = {"EarlyStopping", "ModelCheckpoint"} +_SUPPORTED_CALLBACKS = {"EarlyStopping", "ModelCheckpoint", "ReduceLROnPlateau", "TerminateOnNaN"} def _build_callbacks(train_cfg: Dict[str, Any]) -> List[Any]: @@ -192,6 +197,21 @@ def _build_callbacks(train_cfg: Dict[str, Any]) -> List[Any]: verbose=int(params.get("verbose", 1) or 0), ) ) + elif name == "ReduceLROnPlateau": + callbacks.append( + ReduceLROnPlateau( + monitor=params.get("monitor", "val_loss"), + mode=params.get("mode", "min"), + patience=int(params.get("patience", 2)), + factor=float(params.get("factor", 0.5)), + min_lr=float(params.get("min_lr", 1e-6)), + verbose=int(params.get("verbose", 0) or 0), + ) + ) + elif name == "TerminateOnNaN": + callbacks.append( + TerminateOnNaN(monitor=params.get("monitor", "train_loss")) + ) elif name in _SUPPORTED_CALLBACKS: continue else: diff --git a/src/jaeger/training/pytorch/callbacks.py b/src/jaeger/training/pytorch/callbacks.py index 7741cd8..0951a70 100644 --- a/src/jaeger/training/pytorch/callbacks.py +++ b/src/jaeger/training/pytorch/callbacks.py @@ -1,8 +1,14 @@ +import math from pathlib import Path from typing import Any, Dict, Optional import torch +from jaeger.utils.logging import get_logger + + +logger = get_logger(log_file=None, log_path=None, level=3) + class EarlyStopping: """Stop training when a monitored metric has stopped improving.""" @@ -118,3 +124,87 @@ def _save(self, trainer, epoch: int, logs: Optional[Dict[str, float]] = None): def on_train_end(self, trainer): pass + + +class ReduceLROnPlateau: + """Reduce learning rate when a metric has stopped improving.""" + + def __init__( + self, + monitor: str = "val_loss", + mode: str = "min", + patience: int = 2, + factor: float = 0.5, + min_lr: float = 1e-6, + verbose: int = 0, + ): + self.monitor = monitor + self.mode = mode + self.patience = patience + self.factor = factor + self.min_lr = min_lr + self.verbose = verbose + self.best_value = float("inf") if mode == "min" else float("-inf") + self.wait = 0 + + def on_train_begin(self, trainer): + pass + + def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): + current = logs.get(self.monitor) + if current is None: + return + + improved = ( + self.mode == "min" and current < self.best_value + ) or (self.mode == "max" and current > self.best_value) + + if improved: + self.best_value = current + self.wait = 0 + return + + self.wait += 1 + if self.wait >= self.patience: + self.wait = 0 + self._reduce_lr(trainer.optimizer) + + def _reduce_lr(self, optimizer: torch.optim.Optimizer): + for param_group in optimizer.param_groups: + old_lr = param_group["lr"] + new_lr = max(old_lr * self.factor, self.min_lr) + if old_lr > new_lr: + param_group["lr"] = new_lr + if self.verbose: + logger.info( + "Reducing learning rate: %.6f -> %.6f", old_lr, new_lr + ) + + def on_train_end(self, trainer): + pass + + +class TerminateOnNaN: + """Stop training if a monitored metric becomes NaN or Inf.""" + + def __init__(self, monitor: str = "train_loss"): + self.monitor = monitor + + def on_train_begin(self, trainer): + pass + + def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): + current = logs.get(self.monitor) + if current is None: + return + if math.isnan(current) or math.isinf(current): + logger.warning( + "Stopping training: %s is %s at epoch %d", + self.monitor, + current, + epoch, + ) + trainer.should_stop = True + + def on_train_end(self, trainer): + pass diff --git a/tests/unit/training/pytorch/test_callbacks.py b/tests/unit/training/pytorch/test_callbacks.py index 321da63..0892679 100644 --- a/tests/unit/training/pytorch/test_callbacks.py +++ b/tests/unit/training/pytorch/test_callbacks.py @@ -4,7 +4,12 @@ import torch from torch.utils.data import DataLoader, TensorDataset -from jaeger.training.pytorch.callbacks import EarlyStopping, ModelCheckpoint +from jaeger.training.pytorch.callbacks import ( + EarlyStopping, + ModelCheckpoint, + ReduceLROnPlateau, + TerminateOnNaN, +) from jaeger.training.pytorch.trainer import Trainer @@ -83,3 +88,71 @@ def test_model_checkpoint(): ) trainer.fit() assert ckpt_path.exists() + + +def test_reduce_lr_on_plateau_lowers_lr(): + """ReduceLROnPlateau should lower LR after patience epochs of no improvement.""" + param = torch.nn.Parameter(torch.zeros(1)) + optimizer = torch.optim.SGD([param], lr=0.1) + reduce_lr = ReduceLROnPlateau( + monitor="val_loss", mode="min", patience=1, factor=0.1, min_lr=1e-6 + ) + + class _FakeTrainer: + pass + + fake_trainer = _FakeTrainer() + fake_trainer.optimizer = optimizer + + # Improve once to set baseline. + reduce_lr.on_epoch_end(fake_trainer, 1, {"val_loss": 1.0}) + assert optimizer.param_groups[0]["lr"] == 0.1 + + # One epoch of no improvement triggers reduction (patience=1). + reduce_lr.on_epoch_end(fake_trainer, 2, {"val_loss": 1.0}) + assert abs(optimizer.param_groups[0]["lr"] - 0.01) < 1e-10 + + # Further plateaus reduce again until min_lr is reached. + reduce_lr.on_epoch_end(fake_trainer, 3, {"val_loss": 1.0}) + assert abs(optimizer.param_groups[0]["lr"] - 0.001) < 1e-10 + + +def test_terminate_on_nan_stops_training(): + """TerminateOnNaN should stop training when the monitored metric is NaN.""" + model = _make_model() + train_loader = _make_data() + val_loader = _make_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + terminate = TerminateOnNaN(monitor="train_loss") + + # Monkeypatch train_one_epoch to return NaN loss on epoch 2. + from jaeger.training.pytorch import trainer as trainer_module + from jaeger.training.pytorch.engine import train_one_epoch as original_train + + def _nan_train(*args, **kwargs): + return {"loss": float("nan")} + + original_evaluate = trainer_module.evaluate + + def _nan_evaluate(*args, **kwargs): + return {"loss": float("nan")} + + trainer_module.train_one_epoch = _nan_train + trainer_module.evaluate = _nan_evaluate + try: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=5, + device=torch.device("cpu"), + callbacks=[terminate], + ) + history = trainer.fit() + assert len(history) == 1 + finally: + trainer_module.train_one_epoch = original_train + trainer_module.evaluate = original_evaluate From 019643f7b8e8379ec8e24e397a197513f2eca731 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 22:53:28 +0200 Subject: [PATCH 54/64] refactor(training): remove Trainer's built-in checkpoint save The Trainer was unconditionally writing checkpoint_epoch_{epoch}.pt files in addition to the optional ModelCheckpoint callback. This produced redundant checkpoint files and duplicated responsibilities. - Remove Trainer.checkpoint_dir and Trainer._save_checkpoint(). - Remove checkpoint_dir argument from Trainer instantiation in commands/train.py. - Rely on the ModelCheckpoint callback for all checkpoint saving. - Fix ModelCheckpoint so save_best_only=False also saves the first epoch (previously it skipped the first epoch because best_value always improved from the initial value). - Update tests to remove checkpoint_dir usage and add ModelCheckpoint to the default test config so resume/overwrite tests still work. 283 tests pass; ruff clean. --- src/jaeger/commands/train.py | 2 -- src/jaeger/training/pytorch/callbacks.py | 10 ++++++++-- src/jaeger/training/pytorch/trainer.py | 18 ------------------ tests/unit/commands/test_train_pytorch.py | 16 ++++++++++++++++ tests/unit/training/pytorch/test_trainer.py | 8 -------- 5 files changed, 24 insertions(+), 30 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index daafa9e..d227349 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -483,7 +483,6 @@ def train_fragment_core(**kwargs): device=device, metrics=builder.get_metrics(branch="classifier"), callbacks=callbacks, - checkpoint_dir=str(classifier_dir), history_path=str(classifier_dir / "history.json"), branch="classifier", progress_bar=kwargs.get("progress_bar", False), @@ -556,7 +555,6 @@ def train_fragment_core(**kwargs): device=device, metrics=builder.get_metrics(branch="reliability"), callbacks=rel_callbacks, - checkpoint_dir=str(reliability_dir), history_path=str(reliability_dir / "history.json"), branch="reliability", progress_bar=kwargs.get("progress_bar", False), diff --git a/src/jaeger/training/pytorch/callbacks.py b/src/jaeger/training/pytorch/callbacks.py index 0951a70..56fa21a 100644 --- a/src/jaeger/training/pytorch/callbacks.py +++ b/src/jaeger/training/pytorch/callbacks.py @@ -95,14 +95,20 @@ def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): if improved: self.best_value = current - if self.save_best_only: + + if self.save_best_only: + if improved: self._save(trainer, epoch, logs) if self.verbose: print( f"Epoch {epoch}: {self.monitor} improved to {current:.4f}; saving model" ) - elif not self.save_best_only: + else: self._save(trainer, epoch, logs) + if self.verbose and improved: + print( + f"Epoch {epoch}: {self.monitor} improved to {current:.4f}; saving model" + ) def _save(self, trainer, epoch: int, logs: Optional[Dict[str, float]] = None): logs = dict(logs or {}) diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index a7f321b..85b0bbc 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -47,7 +47,6 @@ def __init__( device: torch.device, metrics: Optional[Dict[str, Any]] = None, callbacks: Optional[List[Any]] = None, - checkpoint_dir: Optional[str] = None, history_path: Optional[str] = None, branch: str = "classifier", progress_bar: bool = False, @@ -65,7 +64,6 @@ def __init__( self.device = device self.metrics = metrics or {} self.callbacks = callbacks or [] - self.checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir else None self.history_path = Path(history_path) if history_path else None self.branch = branch self.progress_bar = progress_bar @@ -188,8 +186,6 @@ def fit(self) -> List[Dict[str, float]]: if getattr(self, "should_stop", False): break - if self.checkpoint_dir: - self._save_checkpoint(epoch) if self.history_path: self._save_history() @@ -206,20 +202,6 @@ def _loader_length(loader: DataLoader) -> Optional[int]: except TypeError: return None - def _save_checkpoint(self, epoch: int): - if self.checkpoint_dir is None: - return - self.checkpoint_dir.mkdir(parents=True, exist_ok=True) - path = self.checkpoint_dir / f"checkpoint_epoch_{epoch}.pt" - torch.save( - { - "epoch": epoch, - "model_state_dict": self.model.state_dict(), - "optimizer_state_dict": self.optimizer.state_dict(), - }, - path, - ) - def _save_history(self): if self.history_path is None: return diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index f8f5c2f..e3ba807 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -52,6 +52,22 @@ def _build_config(tmp_path, train_path, val_path, data_format="numpy_raw"): "classifier_dir": str(tmp_path / "checkpoints" / "classifier"), "optimizer": "adam", "optimizer_params": {"lr": 1e-3}, + "callbacks": { + "classifier": [ + { + "name": "ModelCheckpoint", + "params": { + "filepath": str( + tmp_path / "checkpoints" / "classifier" / "checkpoint_epoch_{epoch:02d}.pt" + ), + "monitor": "val_loss", + "mode": "min", + "save_best_only": False, + "verbose": 0, + }, + } + ] + }, "fragment_classifier_data": { "train": [{"path": [str(train_path)]}], "validation": [{"path": [str(val_path)]}], diff --git a/tests/unit/training/pytorch/test_trainer.py b/tests/unit/training/pytorch/test_trainer.py index 4ff29e5..50bc0d1 100644 --- a/tests/unit/training/pytorch/test_trainer.py +++ b/tests/unit/training/pytorch/test_trainer.py @@ -52,7 +52,6 @@ def test_trainer_fit(): optimizer=optimizer, epochs=2, device=torch.device("cpu"), - checkpoint_dir=tmpdir, history_path=Path(tmpdir) / "history.json", ) history = trainer.fit() @@ -60,7 +59,6 @@ def test_trainer_fit(): assert "train_loss" in history[0] assert "val_loss" in history[0] assert (Path(tmpdir) / "history.json").exists() - assert len(list(Path(tmpdir).glob("checkpoint_epoch_*.pt"))) == 2 def test_trainer_fit_with_progress_bar(): @@ -79,7 +77,6 @@ def test_trainer_fit_with_progress_bar(): optimizer=optimizer, epochs=1, device=torch.device("cpu"), - checkpoint_dir=tmpdir, history_path=Path(tmpdir) / "history.json", progress_bar=True, ) @@ -118,7 +115,6 @@ def test_trainer_respects_train_and_validation_steps(): optimizer=optimizer, epochs=2, device=torch.device("cpu"), - checkpoint_dir=tmpdir, history_path=Path(tmpdir) / "history.json", train_steps=2, validation_steps=1, @@ -146,7 +142,6 @@ def test_trainer_resumes_from_start_epoch(): optimizer=optimizer, epochs=5, device=torch.device("cpu"), - checkpoint_dir=tmpdir, history_path=Path(tmpdir) / "history.json", start_epoch=2, ) @@ -173,12 +168,9 @@ def test_trainer_skips_training_when_start_epoch_reaches_epochs(): optimizer=optimizer, epochs=3, device=torch.device("cpu"), - checkpoint_dir=tmpdir, history_path=Path(tmpdir) / "history.json", start_epoch=3, ) history = trainer.fit() assert len(history) == 0 assert model.forward_count == 0 - # No new checkpoint should be saved. - assert len(list(Path(tmpdir).glob("checkpoint_epoch_*.pt"))) == 0 From c8c632a1b600f1cfdb2245ab6f173ca1309e3e98 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:00:14 +0200 Subject: [PATCH 55/64] feat(training): add JsonLogger callback and remove built-in history save Add a config-driven JsonLogger callback for saving per-epoch metrics and remove the Trainer's built-in history.json saving for consistency with the ModelCheckpoint-only approach. - JsonLogger writes a JSON array of epoch logs; append=True loads existing history and extends it (useful for resumed runs). - Wire JsonLogger into commands/train.py _build_callbacks. - Remove Trainer.history_path and Trainer._save_history(). - Update commands/train.py to stop passing history_path. - Update tests and add JsonLogger-specific tests. 286 tests pass; ruff clean. --- src/jaeger/commands/train.py | 14 +- src/jaeger/training/pytorch/callbacks.py | 37 +++- src/jaeger/training/pytorch/trainer.py | 13 -- tests/unit/training/pytorch/test_callbacks.py | 64 +++++++ tests/unit/training/pytorch/test_trainer.py | 163 ++++++++++-------- 5 files changed, 201 insertions(+), 90 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index d227349..c8d53e3 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -26,6 +26,7 @@ from jaeger.nnlib.pytorch.builder import ModelBuilder from jaeger.training.pytorch.callbacks import ( EarlyStopping, + JsonLogger, ModelCheckpoint, ReduceLROnPlateau, TerminateOnNaN, @@ -163,7 +164,7 @@ def _callback_path_from_config( return None -_SUPPORTED_CALLBACKS = {"EarlyStopping", "ModelCheckpoint", "ReduceLROnPlateau", "TerminateOnNaN"} +_SUPPORTED_CALLBACKS = {"EarlyStopping", "JsonLogger", "ModelCheckpoint", "ReduceLROnPlateau", "TerminateOnNaN"} def _build_callbacks(train_cfg: Dict[str, Any]) -> List[Any]: @@ -212,6 +213,15 @@ def _build_callbacks(train_cfg: Dict[str, Any]) -> List[Any]: callbacks.append( TerminateOnNaN(monitor=params.get("monitor", "train_loss")) ) + elif name == "JsonLogger": + filename = params.get("filename") + if filename is not None: + callbacks.append( + JsonLogger( + filename=filename, + append=bool(params.get("append", False)), + ) + ) elif name in _SUPPORTED_CALLBACKS: continue else: @@ -483,7 +493,6 @@ def train_fragment_core(**kwargs): device=device, metrics=builder.get_metrics(branch="classifier"), callbacks=callbacks, - history_path=str(classifier_dir / "history.json"), branch="classifier", progress_bar=kwargs.get("progress_bar", False), profile=kwargs.get("profile", False), @@ -555,7 +564,6 @@ def train_fragment_core(**kwargs): device=device, metrics=builder.get_metrics(branch="reliability"), callbacks=rel_callbacks, - history_path=str(reliability_dir / "history.json"), branch="reliability", progress_bar=kwargs.get("progress_bar", False), profile=kwargs.get("profile", False), diff --git a/src/jaeger/training/pytorch/callbacks.py b/src/jaeger/training/pytorch/callbacks.py index 56fa21a..6e29445 100644 --- a/src/jaeger/training/pytorch/callbacks.py +++ b/src/jaeger/training/pytorch/callbacks.py @@ -1,6 +1,7 @@ +import json import math from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional import torch @@ -214,3 +215,37 @@ def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): def on_train_end(self, trainer): pass + + +class JsonLogger: + """Write per-epoch metrics to a JSON file. + + The file is written as a JSON array. When ``append=True``, any existing + array is loaded at the start of training and new epochs are appended. + """ + + def __init__(self, filename: str, append: bool = False): + self.filename = Path(filename) + self.append = append + self.history: List[Dict[str, Any]] = [] + + def on_train_begin(self, trainer): + self.history = [] + if self.append and self.filename.exists(): + try: + with self.filename.open("r") as fh: + self.history = json.load(fh) + except (json.JSONDecodeError, OSError) as exc: + logger.warning("Could not load existing history: %s", exc) + + def on_epoch_end(self, trainer, epoch: int, logs: Dict[str, float]): + self.history.append(dict(logs)) + self._write() + + def on_train_end(self, trainer): + self._write() + + def _write(self): + self.filename.parent.mkdir(parents=True, exist_ok=True) + with self.filename.open("w") as fh: + json.dump(self.history, fh, indent=2) diff --git a/src/jaeger/training/pytorch/trainer.py b/src/jaeger/training/pytorch/trainer.py index 85b0bbc..8565d29 100644 --- a/src/jaeger/training/pytorch/trainer.py +++ b/src/jaeger/training/pytorch/trainer.py @@ -1,5 +1,3 @@ -import json -from pathlib import Path from typing import Any, Callable, Dict, List, Optional import torch @@ -47,7 +45,6 @@ def __init__( device: torch.device, metrics: Optional[Dict[str, Any]] = None, callbacks: Optional[List[Any]] = None, - history_path: Optional[str] = None, branch: str = "classifier", progress_bar: bool = False, profile: bool = False, @@ -64,7 +61,6 @@ def __init__( self.device = device self.metrics = metrics or {} self.callbacks = callbacks or [] - self.history_path = Path(history_path) if history_path else None self.branch = branch self.progress_bar = progress_bar self.profile = profile @@ -186,9 +182,6 @@ def fit(self) -> List[Dict[str, float]]: if getattr(self, "should_stop", False): break - if self.history_path: - self._save_history() - for callback in self.callbacks: if hasattr(callback, "on_train_end"): callback.on_train_end(self) @@ -202,9 +195,3 @@ def _loader_length(loader: DataLoader) -> Optional[int]: except TypeError: return None - def _save_history(self): - if self.history_path is None: - return - self.history_path.parent.mkdir(parents=True, exist_ok=True) - with self.history_path.open("w") as fh: - json.dump(self.history, fh, indent=2) diff --git a/tests/unit/training/pytorch/test_callbacks.py b/tests/unit/training/pytorch/test_callbacks.py index 0892679..5d02918 100644 --- a/tests/unit/training/pytorch/test_callbacks.py +++ b/tests/unit/training/pytorch/test_callbacks.py @@ -4,8 +4,11 @@ import torch from torch.utils.data import DataLoader, TensorDataset +import json + from jaeger.training.pytorch.callbacks import ( EarlyStopping, + JsonLogger, ModelCheckpoint, ReduceLROnPlateau, TerminateOnNaN, @@ -156,3 +159,64 @@ def _nan_evaluate(*args, **kwargs): finally: trainer_module.train_one_epoch = original_train trainer_module.evaluate = original_evaluate + + +def test_json_logger_writes_history(): + """JsonLogger should write per-epoch metrics to a JSON file.""" + model = _make_model() + train_loader = _make_data() + val_loader = _make_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + history_path = Path(tmpdir) / "history.json" + json_logger = JsonLogger(filename=str(history_path)) + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + callbacks=[json_logger], + ) + trainer.fit() + assert history_path.exists() + with history_path.open() as fh: + saved = json.load(fh) + assert len(saved) == 2 + assert "train_loss" in saved[0] + assert "val_loss" in saved[0] + + +def test_json_logger_appends_to_existing_history(): + """JsonLogger with append=True should extend an existing history file.""" + model = _make_model() + train_loader = _make_data() + val_loader = _make_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + with tempfile.TemporaryDirectory() as tmpdir: + history_path = Path(tmpdir) / "history.json" + history_path.write_text(json.dumps([{"epoch": 1, "train_loss": 0.9}])) + + json_logger = JsonLogger(filename=str(history_path), append=True) + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + callbacks=[json_logger], + ) + trainer.fit() + + with history_path.open() as fh: + saved = json.load(fh) + assert len(saved) == 3 + assert saved[0]["epoch"] == 1 diff --git a/tests/unit/training/pytorch/test_trainer.py b/tests/unit/training/pytorch/test_trainer.py index 50bc0d1..ed22d13 100644 --- a/tests/unit/training/pytorch/test_trainer.py +++ b/tests/unit/training/pytorch/test_trainer.py @@ -1,9 +1,9 @@ -import tempfile from pathlib import Path import torch from torch.utils.data import DataLoader, TensorDataset +from jaeger.training.pytorch.callbacks import JsonLogger from jaeger.training.pytorch.trainer import Trainer @@ -43,22 +43,19 @@ def test_trainer_fit(): loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - with tempfile.TemporaryDirectory() as tmpdir: - trainer = Trainer( - model=model, - train_loader=train_loader, - val_loader=val_loader, - loss_fn=loss_fn, - optimizer=optimizer, - epochs=2, - device=torch.device("cpu"), - history_path=Path(tmpdir) / "history.json", - ) - history = trainer.fit() - assert len(history) == 2 - assert "train_loss" in history[0] - assert "val_loss" in history[0] - assert (Path(tmpdir) / "history.json").exists() + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + ) + history = trainer.fit() + assert len(history) == 2 + assert "train_loss" in history[0] + assert "val_loss" in history[0] def test_trainer_fit_with_progress_bar(): @@ -68,22 +65,20 @@ def test_trainer_fit_with_progress_bar(): loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - with tempfile.TemporaryDirectory() as tmpdir: - trainer = Trainer( - model=model, - train_loader=train_loader, - val_loader=val_loader, - loss_fn=loss_fn, - optimizer=optimizer, - epochs=1, - device=torch.device("cpu"), - history_path=Path(tmpdir) / "history.json", - progress_bar=True, - ) - history = trainer.fit() - assert len(history) == 1 - assert "train_loss" in history[0] - assert "val_loss" in history[0] + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=1, + device=torch.device("cpu"), + progress_bar=True, + ) + history = trainer.fit() + assert len(history) == 1 + assert "train_loss" in history[0] + assert "val_loss" in history[0] class _CounterModel(_DummyClassifier): @@ -106,23 +101,21 @@ def test_trainer_respects_train_and_validation_steps(): loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - with tempfile.TemporaryDirectory() as tmpdir: - trainer = Trainer( - model=model, - train_loader=train_loader, - val_loader=val_loader, - loss_fn=loss_fn, - optimizer=optimizer, - epochs=2, - device=torch.device("cpu"), - history_path=Path(tmpdir) / "history.json", - train_steps=2, - validation_steps=1, - ) - history = trainer.fit() - assert len(history) == 2 - # Each epoch runs train_steps + validation_steps forward passes. - assert model.forward_count == 2 * (2 + 1) + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=2, + device=torch.device("cpu"), + train_steps=2, + validation_steps=1, + ) + history = trainer.fit() + assert len(history) == 2 + # Each epoch runs train_steps + validation_steps forward passes. + assert model.forward_count == 2 * (2 + 1) def test_trainer_resumes_from_start_epoch(): @@ -133,22 +126,20 @@ def test_trainer_resumes_from_start_epoch(): loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - with tempfile.TemporaryDirectory() as tmpdir: - trainer = Trainer( - model=model, - train_loader=train_loader, - val_loader=val_loader, - loss_fn=loss_fn, - optimizer=optimizer, - epochs=5, - device=torch.device("cpu"), - history_path=Path(tmpdir) / "history.json", - start_epoch=2, - ) - history = trainer.fit() - # start_epoch=2 means epochs 3, 4, 5 are trained. - assert len(history) == 3 - assert [entry["epoch"] for entry in history] == [3, 4, 5] + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=5, + device=torch.device("cpu"), + start_epoch=2, + ) + history = trainer.fit() + # start_epoch=2 means epochs 3, 4, 5 are trained. + assert len(history) == 3 + assert [entry["epoch"] for entry in history] == [3, 4, 5] def test_trainer_skips_training_when_start_epoch_reaches_epochs(): @@ -159,18 +150,44 @@ def test_trainer_skips_training_when_start_epoch_reaches_epochs(): loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - with tempfile.TemporaryDirectory() as tmpdir: + trainer = Trainer( + model=model, + train_loader=train_loader, + val_loader=val_loader, + loss_fn=loss_fn, + optimizer=optimizer, + epochs=3, + device=torch.device("cpu"), + start_epoch=3, + ) + history = trainer.fit() + assert len(history) == 0 + assert model.forward_count == 0 + + +def test_trainer_json_logger_callback(): + """Trainer should work with JsonLogger callback to save history.""" + model = _make_dummy_model() + train_loader = _make_dummy_data() + val_loader = _make_dummy_data() + loss_fn = torch.nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + history_path = Path("history.json") + json_logger = JsonLogger(filename=str(history_path)) + try: trainer = Trainer( model=model, train_loader=train_loader, val_loader=val_loader, loss_fn=loss_fn, optimizer=optimizer, - epochs=3, + epochs=2, device=torch.device("cpu"), - history_path=Path(tmpdir) / "history.json", - start_epoch=3, + callbacks=[json_logger], ) - history = trainer.fit() - assert len(history) == 0 - assert model.forward_count == 0 + trainer.fit() + assert history_path.exists() + finally: + if history_path.exists(): + history_path.unlink() From c8c16f3af24381660183ff15ec377ddc211c1a48 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:04:32 +0200 Subject: [PATCH 56/64] feat(metrics): make classifier/reliability metrics configurable Replace hardcoded per-class precision/recall with a config-driven metric builder supporting categorical_accuracy, binary_accuracy, precision, recall, per_class_precision, per_class_recall, auc, and per_class_auc. - Add metric classes: CategoricalAccuracy, BinaryAccuracy, Precision, Recall, AUC, AUCForClass. - Update PrecisionForClass/RecallForClass to handle both one-hot and index labels via _as_class_indices(). - Add build_metrics() dispatcher that parses metrics_classifier and metrics_reliability config lists. - ModelBuilder.get_metrics() now delegates to build_metrics(). - Add unit tests for new metrics and config-driven builder. 291 tests pass; ruff clean. --- src/jaeger/nnlib/pytorch/builder.py | 9 +- src/jaeger/nnlib/pytorch/metrics.py | 202 ++++++++++++++++++++++- tests/unit/nnlib/pytorch/test_metrics.py | 61 ++++++- 3 files changed, 263 insertions(+), 9 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/builder.py b/src/jaeger/nnlib/pytorch/builder.py index 3a5a186..a300f07 100644 --- a/src/jaeger/nnlib/pytorch/builder.py +++ b/src/jaeger/nnlib/pytorch/builder.py @@ -181,18 +181,15 @@ def compile_model( raise ValueError(f"Unknown train_branch: {train_branch}") def get_metrics(self, branch: str = "classifier") -> Dict[str, Any]: - from jaeger.nnlib.pytorch.metrics import PrecisionForClass, RecallForClass + from jaeger.nnlib.pytorch.metrics import build_metrics - metrics = {} out_dim = ( self.classifier_out_dim if branch == "classifier" else self.reliability_out_dim ) - for cls in range(out_dim): - metrics[f"precision_class_{cls}"] = PrecisionForClass(class_id=cls) - metrics[f"recall_class_{cls}"] = RecallForClass(class_id=cls) - return metrics + metrics_cfg = self.train_cfg.get(f"metrics_{branch}", []) + return build_metrics(metrics_cfg, out_dim) def _get_optimizer(self, name: str) -> torch.optim.Optimizer: """Resolve a PyTorch optimizer class by case-insensitive name.""" diff --git a/src/jaeger/nnlib/pytorch/metrics.py b/src/jaeger/nnlib/pytorch/metrics.py index 6e20e09..9f8b64e 100644 --- a/src/jaeger/nnlib/pytorch/metrics.py +++ b/src/jaeger/nnlib/pytorch/metrics.py @@ -1,6 +1,55 @@ import torch +def _as_class_indices(y_true: torch.Tensor) -> torch.Tensor: + """Return class-index labels from one-hot or index labels.""" + if y_true.dim() == 2 and y_true.shape[1] > 1: + return y_true.argmax(dim=-1) + return y_true.view(-1).long() + + +class CategoricalAccuracy: + """Multi-class accuracy.""" + + def __init__(self): + self.correct = 0 + self.total = 0 + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + pred_labels = y_pred.argmax(dim=-1) + true_labels = _as_class_indices(y_true) + self.correct += int((pred_labels == true_labels).sum().item()) + self.total += y_true.size(0) + + def compute(self) -> float: + if self.total == 0: + return 0.0 + return self.correct / self.total + + +class BinaryAccuracy: + """Binary accuracy.""" + + def __init__(self, threshold: float = 0.5): + self.threshold = threshold + self.correct = 0 + self.total = 0 + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + if y_pred.shape[-1] == 1: + pred_labels = (torch.sigmoid(y_pred).squeeze(-1) >= self.threshold).long() + else: + pred_labels = y_pred.argmax(dim=-1) + true_labels = _as_class_indices(y_true) + self.correct += int((pred_labels == true_labels).sum().item()) + self.total += y_true.size(0) + + def compute(self) -> float: + if self.total == 0: + return 0.0 + return self.correct / self.total + + class PrecisionForClass: def __init__(self, class_id: int): self.class_id = class_id @@ -9,7 +58,7 @@ def __init__(self, class_id: int): def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): pred_labels = y_pred.argmax(dim=-1) - true_labels = y_true.argmax(dim=-1) + true_labels = _as_class_indices(y_true) cls = self.class_id self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) self.fp += int(((pred_labels == cls) & (true_labels != cls)).sum().item()) @@ -27,7 +76,7 @@ def __init__(self, class_id: int): def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): pred_labels = y_pred.argmax(dim=-1) - true_labels = y_true.argmax(dim=-1) + true_labels = _as_class_indices(y_true) cls = self.class_id self.tp += int(((pred_labels == cls) & (true_labels == cls)).sum().item()) self.fn += int(((pred_labels != cls) & (true_labels == cls)).sum().item()) @@ -36,3 +85,152 @@ def compute(self) -> float: if self.tp + self.fn == 0: return 0.0 return self.tp / (self.tp + self.fn) + + +class _SklearnMetric: + """Base for metrics that accumulate predictions and use scikit-learn.""" + + def __init__(self): + self.y_preds: list[torch.Tensor] = [] + self.y_trues: list[torch.Tensor] = [] + + def update(self, y_pred: torch.Tensor, y_true: torch.Tensor): + self.y_preds.append(y_pred.detach().cpu()) + self.y_trues.append(y_true.detach().cpu()) + + def _labels(self) -> tuple[torch.Tensor, torch.Tensor]: + y_pred = torch.cat(self.y_preds) + y_true = torch.cat(self.y_trues) + return y_pred, y_true + + +class Precision(_SklearnMetric): + """Overall precision (macro-averaged by default).""" + + def __init__(self, average: str = "macro"): + super().__init__() + self.average = average + + def compute(self) -> float: + from sklearn.metrics import precision_score + + y_pred, y_true = self._labels() + pred_labels = y_pred.argmax(dim=-1).numpy() + true_labels = _as_class_indices(y_true).numpy() + return float( + precision_score(true_labels, pred_labels, average=self.average, zero_division=0) + ) + + +class Recall(_SklearnMetric): + """Overall recall (macro-averaged by default).""" + + def __init__(self, average: str = "macro"): + super().__init__() + self.average = average + + def compute(self) -> float: + from sklearn.metrics import recall_score + + y_pred, y_true = self._labels() + pred_labels = y_pred.argmax(dim=-1).numpy() + true_labels = _as_class_indices(y_true).numpy() + return float( + recall_score(true_labels, pred_labels, average=self.average, zero_division=0) + ) + + +class AUC(_SklearnMetric): + """ROC AUC. Multi-class problems use one-vs-rest macro averaging.""" + + def __init__(self, multi_class: str = "ovr", average: str = "macro"): + super().__init__() + self.multi_class = multi_class + self.average = average + + def compute(self) -> float: + from sklearn.metrics import roc_auc_score + + y_pred, y_true = self._labels() + y_score = torch.softmax(y_pred, dim=-1).numpy() + true_labels = _as_class_indices(y_true).numpy() + num_classes = y_score.shape[-1] + if num_classes == 2: + y_score = y_score[:, 1] + return float( + roc_auc_score( + true_labels, + y_score, + multi_class=self.multi_class if num_classes > 2 else "raise", + average=self.average, + ) + ) + + +class AUCForClass(_SklearnMetric): + """One-vs-rest ROC AUC for a single class.""" + + def __init__(self, class_id: int): + super().__init__() + self.class_id = class_id + + def compute(self) -> float: + from sklearn.metrics import roc_auc_score + + y_pred, y_true = self._labels() + y_score = torch.softmax(y_pred, dim=-1)[:, self.class_id].numpy() + true_labels = (_as_class_indices(y_true) == self.class_id).long().numpy() + return float(roc_auc_score(true_labels, y_score)) + + +_METRIC_BUILDERS = { + "categorical_accuracy": lambda params, num_classes: CategoricalAccuracy(), + "binary_accuracy": lambda params, num_classes: BinaryAccuracy(), + "precision": lambda params, num_classes: Precision(average=params.get("average", "macro")), + "recall": lambda params, num_classes: Recall(average=params.get("average", "macro")), + "per_class_precision": lambda params, num_classes: { + f"precision_class_{cls}": PrecisionForClass(class_id=cls) + for cls in range(num_classes) + }, + "per_class_recall": lambda params, num_classes: { + f"recall_class_{cls}": RecallForClass(class_id=cls) + for cls in range(num_classes) + }, + "auc": lambda params, num_classes: AUC( + multi_class=params.get("multi_class", "ovr"), + average=params.get("average", "macro"), + ), + "per_class_auc": lambda params, num_classes: { + f"auc_class_{cls}": AUCForClass(class_id=cls) + for cls in range(num_classes) + }, +} + + +def build_metrics(metrics_cfg: list[dict], num_classes: int) -> dict[str, object]: + """Build a dictionary of metric objects from a config list. + + Parameters + ---------- + metrics_cfg: + List of dicts with ``name`` and optional ``params`` keys. + num_classes: + Number of output classes for the branch. + + Returns + ------- + Mapping from metric name to metric instance. + """ + metrics: dict[str, object] = {} + for entry in metrics_cfg or []: + name = entry.get("name") + params = entry.get("params") or {} + builder = _METRIC_BUILDERS.get(name) + if builder is None: + raise ValueError(f"Unsupported metric: {name}") + result = builder(params, num_classes) + if isinstance(result, dict): + metrics.update(result) + else: + metrics[name] = result + return metrics diff --git a/tests/unit/nnlib/pytorch/test_metrics.py b/tests/unit/nnlib/pytorch/test_metrics.py index 5573dea..8f15075 100644 --- a/tests/unit/nnlib/pytorch/test_metrics.py +++ b/tests/unit/nnlib/pytorch/test_metrics.py @@ -1,5 +1,13 @@ import torch -from jaeger.nnlib.pytorch.metrics import PrecisionForClass, RecallForClass +from jaeger.nnlib.pytorch.metrics import ( + BinaryAccuracy, + CategoricalAccuracy, + Precision, + PrecisionForClass, + Recall, + RecallForClass, + build_metrics, +) def test_precision_for_class(): @@ -39,3 +47,54 @@ def test_metrics_multiple_updates(): metric.update(preds2, labels2) result = metric.compute() assert 0.0 <= result <= 1.0 + + +def test_categorical_accuracy(): + metric = CategoricalAccuracy() + preds = torch.tensor([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]) + labels = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.float32) + metric.update(preds, labels) + result = metric.compute() + # Predictions: class 1, class 0, class 1. Labels: class 1, class 0, class 1. + assert result == 1.0 + + +def test_binary_accuracy(): + metric = BinaryAccuracy() + preds = torch.tensor([[2.0], [-2.0], [1.0], [-1.0]]) + labels = torch.tensor([[1], [0], [0], [1]], dtype=torch.float32) + metric.update(preds, labels) + result = metric.compute() + # Predictions: 1, 0, 1, 0. Labels: 1, 0, 0, 1. 2/4 correct. + assert result == 0.5 + + +def test_overall_precision_and_recall(): + metric_precision = Precision(average="macro") + metric_recall = Recall(average="macro") + preds = torch.tensor([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]) + labels = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.float32) + metric_precision.update(preds, labels) + metric_recall.update(preds, labels) + assert 0.0 <= metric_precision.compute() <= 1.0 + assert 0.0 <= metric_recall.compute() <= 1.0 + + +def test_build_metrics_from_config(): + metrics_cfg = [ + {"name": "categorical_accuracy", "params": None}, + {"name": "per_class_precision", "params": None}, + ] + metrics = build_metrics(metrics_cfg, num_classes=3) + assert "categorical_accuracy" in metrics + assert "precision_class_0" in metrics + assert "precision_class_1" in metrics + assert "precision_class_2" in metrics + + +def test_build_metrics_unknown_metric_raises(): + metrics_cfg = [{"name": "not_a_metric", "params": None}] + try: + build_metrics(metrics_cfg, num_classes=2) + except ValueError as exc: + assert "Unsupported metric" in str(exc) From bda5e1d9225f667ab06906ec9adee7e5cb10d08a Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:22:01 +0200 Subject: [PATCH 57/64] build: add torchinfo dependency for model summary --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index f4ce25b..c8465a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "requests >= 2.0", "rich >= 13.0", "scipy >= 1.10", + "torchinfo >=1.8.0", "pyfastx >= 2", "ruptures >= 1.1.9", "scikit-learn >= 1.6.0", From 8e659a699f9a0a147c9fb2a7774369b212bb39f6 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:23:33 +0200 Subject: [PATCH 58/64] feat(nnlib): add ModelSummary wrapper around torchinfo --- src/jaeger/nnlib/pytorch/summary.py | 47 ++++++++++++++++++++++ tests/unit/nnlib/pytorch/test_summary.py | 50 ++++++++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 src/jaeger/nnlib/pytorch/summary.py create mode 100644 tests/unit/nnlib/pytorch/test_summary.py diff --git a/src/jaeger/nnlib/pytorch/summary.py b/src/jaeger/nnlib/pytorch/summary.py new file mode 100644 index 0000000..133785f --- /dev/null +++ b/src/jaeger/nnlib/pytorch/summary.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import logging +from typing import Tuple + +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +class ModelSummary: + """Keras-style model summary using torchinfo.""" + + def __init__(self, model: nn.Module, input_data: Tuple[torch.Tensor, ...]): + """Initialize summary wrapper with a model and representative input tensors.""" + self.model = model + self.input_data = input_data + + def summary(self, branch_label: str = "") -> str: + """Return a Keras-style summary string, or an empty string on failure.""" + try: + from torchinfo import summary as torchinfo_summary + except ImportError: + logger.warning( + "torchinfo is not installed; skipping %s model summary", + branch_label or "model", + ) + return "" + + try: + result = torchinfo_summary( + self.model, + input_data=self.input_data, + col_names=["input_size", "output_size", "num_params"], + row_settings=["var_names"], + depth=8, + verbose=0, + ) + except Exception as exc: # Intentionally broad: model summary must never block training. + logger.warning("Failed to generate %s model summary: %s", branch_label or "model", exc) + return "" + + text = str(result) + if branch_label: + text = f"=== {branch_label} model summary ===\n{text}" + return text diff --git a/tests/unit/nnlib/pytorch/test_summary.py b/tests/unit/nnlib/pytorch/test_summary.py new file mode 100644 index 0000000..0a3ce44 --- /dev/null +++ b/tests/unit/nnlib/pytorch/test_summary.py @@ -0,0 +1,50 @@ +import pytest +import torch +from jaeger.nnlib.pytorch.models import ( + ClassificationHead, + Embedding, + JaegerModel, + RepresentationModel, +) +from jaeger.nnlib.pytorch.summary import ModelSummary + +try: + import torchinfo +except ImportError: + torchinfo = None + + +def _make_tiny_jaeger_model(): + """Build a minimal JaegerModel for lightweight summary tests.""" + embedding = Embedding( + input_type="nucleotide", + vocab_size=5, + embedding_size=8, + use_embedding_layer=False, + onehot_dim=4, + ) + rep_model = RepresentationModel( + embedding=embedding, + hidden_layers=[ + {"name": "masked_conv1d", "config": {"filters": 8, "kernel_size": 3}}, + {"name": "masked_batchnorm", "config": {}}, + ], + pooling="average", + ) + classifier = ClassificationHead( + input_dim=rep_model.output_dim, + num_classes=2, + hidden_units=[], + ) + return JaegerModel(rep_model=rep_model, classification_head=classifier) + + +@pytest.mark.skipif(torchinfo is None, reason="torchinfo not installed") +def test_model_summary_returns_non_empty_string(): + model = _make_tiny_jaeger_model() + dummy = (torch.zeros((1, 2, 32), dtype=torch.long), torch.ones((1, 2, 32), dtype=torch.bool)) + summary = ModelSummary(model, input_data=dummy) + text = summary.summary(branch_label="classifier") + assert isinstance(text, str) + assert "Total params" in text + assert len(text) > 0 From 0a1c4064ebd7a949ee33078effef61592ebb7794 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:26:47 +0200 Subject: [PATCH 59/64] test(summary): cover missing and failing torchinfo --- src/jaeger/nnlib/pytorch/summary.py | 11 ++++--- tests/unit/nnlib/pytorch/test_summary.py | 40 ++++++++++++++++++++++-- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/src/jaeger/nnlib/pytorch/summary.py b/src/jaeger/nnlib/pytorch/summary.py index 133785f..3cb4d54 100644 --- a/src/jaeger/nnlib/pytorch/summary.py +++ b/src/jaeger/nnlib/pytorch/summary.py @@ -6,6 +6,11 @@ import torch import torch.nn as nn +try: + import torchinfo +except ImportError: + torchinfo = None + logger = logging.getLogger(__name__) @@ -19,9 +24,7 @@ def __init__(self, model: nn.Module, input_data: Tuple[torch.Tensor, ...]): def summary(self, branch_label: str = "") -> str: """Return a Keras-style summary string, or an empty string on failure.""" - try: - from torchinfo import summary as torchinfo_summary - except ImportError: + if torchinfo is None: logger.warning( "torchinfo is not installed; skipping %s model summary", branch_label or "model", @@ -29,7 +32,7 @@ def summary(self, branch_label: str = "") -> str: return "" try: - result = torchinfo_summary( + result = torchinfo.summary( self.model, input_data=self.input_data, col_names=["input_size", "output_size", "num_params"], diff --git a/tests/unit/nnlib/pytorch/test_summary.py b/tests/unit/nnlib/pytorch/test_summary.py index 0a3ce44..d872143 100644 --- a/tests/unit/nnlib/pytorch/test_summary.py +++ b/tests/unit/nnlib/pytorch/test_summary.py @@ -39,12 +39,46 @@ def _make_tiny_jaeger_model(): return JaegerModel(rep_model=rep_model, classification_head=classifier) +@pytest.fixture +def dummy_input(): + return ( + torch.zeros((1, 2, 32), dtype=torch.long), + torch.ones((1, 2, 32), dtype=torch.bool), + ) + + @pytest.mark.skipif(torchinfo is None, reason="torchinfo not installed") -def test_model_summary_returns_non_empty_string(): +def test_model_summary_returns_non_empty_string(dummy_input): model = _make_tiny_jaeger_model() - dummy = (torch.zeros((1, 2, 32), dtype=torch.long), torch.ones((1, 2, 32), dtype=torch.bool)) - summary = ModelSummary(model, input_data=dummy) + summary = ModelSummary(model, input_data=dummy_input) text = summary.summary(branch_label="classifier") assert isinstance(text, str) assert "Total params" in text assert len(text) > 0 + + +def test_model_summary_warns_when_torchinfo_missing(monkeypatch, caplog, dummy_input): + import jaeger.nnlib.pytorch.summary as summary_mod + + monkeypatch.setattr(summary_mod, "torchinfo", None) + model = _make_tiny_jaeger_model() + with caplog.at_level("WARNING"): + text = ModelSummary(model, input_data=dummy_input).summary(branch_label="classifier") + assert text == "" + assert "torchinfo is not installed" in caplog.text + + +def test_model_summary_warns_when_torchinfo_raises(monkeypatch, caplog, dummy_input): + import jaeger.nnlib.pytorch.summary as summary_mod + + def _broken_summary(*args, **kwargs): + raise RuntimeError("boom") + + fake_torchinfo = type("FakeTorchinfo", (), {"summary": staticmethod(_broken_summary)})() + monkeypatch.setattr(summary_mod, "torchinfo", fake_torchinfo) + + model = _make_tiny_jaeger_model() + with caplog.at_level("WARNING"): + text = ModelSummary(model, input_data=dummy_input).summary(branch_label="classifier") + assert text == "" + assert "Failed to generate classifier model summary" in caplog.text From 92029b5ec67f1edff7e7033910b714cae5c003b9 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:30:06 +0200 Subject: [PATCH 60/64] refactor(train): extract _make_dummy_input helper --- src/jaeger/commands/train.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index c8d53e3..a3c8cd4 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -76,20 +76,19 @@ def _resolve_steps(value: Optional[int]) -> Optional[int]: return int(value) -def _initialize_lazy_layers( - models: Dict[str, nn.Module], config: Dict[str, Any] -) -> None: - """Run a dummy forward pass to materialize any lazy layers.""" +def _make_dummy_input( + config: Dict[str, Any], device: torch.device, length: Optional[int] = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build a dummy (input, mask) tuple matching the model config.""" model_cfg = config.get("model", {}) embedding_cfg = model_cfg.get("embedding", {}) sp_cfg = model_cfg.get("string_processor", {}) input_type = embedding_cfg.get("input_type", "translated") input_shape = embedding_cfg.get("input_shape") crop_size = int(sp_cfg.get("crop_size", 500)) - device = next(models["jaeger_classifier"].parameters()).device - # Use a short dummy length; lazy layers only need to infer channel dims. - length = 64 if input_type == "nucleotide" else max(1, crop_size // 3 - 1) + if length is None: + length = 64 if input_type == "nucleotide" else max(1, crop_size // 3 - 1) if input_shape is not None and len(input_shape) == 3: # One-hot input, e.g. [2, null, 4] or [6, null, vocab_size]. @@ -108,6 +107,15 @@ def _initialize_lazy_layers( dtype=torch.bool, device=device, ) + return dummy, mask + + +def _initialize_lazy_layers( + models: Dict[str, nn.Module], config: Dict[str, Any] +) -> None: + """Run a dummy forward pass to materialize any lazy layers.""" + device = next(models["jaeger_classifier"].parameters()).device + dummy, mask = _make_dummy_input(config, device) with torch.no_grad(): models["jaeger_classifier"](dummy, mask) From 9b8ff6eea21661b396ed5dbdd5a678968cd1fcb0 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:34:54 +0200 Subject: [PATCH 61/64] feat(train): print torchinfo model summary before training --- src/jaeger/commands/train.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/jaeger/commands/train.py b/src/jaeger/commands/train.py index a3c8cd4..c349037 100644 --- a/src/jaeger/commands/train.py +++ b/src/jaeger/commands/train.py @@ -24,6 +24,7 @@ from jaeger.dataops.pytorch.builders import build_datasets from jaeger.nnlib.pytorch.builder import ModelBuilder +from jaeger.nnlib.pytorch.summary import ModelSummary from jaeger.training.pytorch.callbacks import ( EarlyStopping, JsonLogger, @@ -110,6 +111,19 @@ def _make_dummy_input( return dummy, mask +def _print_model_summary( + model: nn.Module, config: Dict[str, Any], branch_label: str +) -> None: + """Log a Keras-style model summary for *model* using ModelSummary.""" + device = next(model.parameters()).device + dummy, mask = _make_dummy_input(config, device) + summary_text = ModelSummary(model, input_data=(dummy, mask)).summary( + branch_label=branch_label + ) + if summary_text: + logger.info("%s", summary_text) + + def _initialize_lazy_layers( models: Dict[str, nn.Module], config: Dict[str, Any] ) -> None: @@ -425,6 +439,8 @@ def train_fragment_core(**kwargs): if models.get("reliability_head") is not None: models["reliability_head"].to(device) + _print_model_summary(models["jaeger_classifier"], config, branch_label="classifier") + train_cfg = config.get("training", {}) classifier_dir = Path(train_cfg.get("classifier_dir", "checkpoints/classifier")) reliability_dir = Path( @@ -535,6 +551,7 @@ def train_fragment_core(**kwargs): rel_pipeline = _ReliabilityPipeline( rep_model, models["reliability_head"] ) + _print_model_summary(rel_pipeline, config, branch_label="reliability") rel_optimizer = torch.optim.Adam( rel_pipeline.parameters(), lr=(train_cfg.get("optimizer_params", {}) or {}).get("lr", 1e-3), From d4bda436db637ad8dd574f8e2e78ca743c75d6ff Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:39:56 +0200 Subject: [PATCH 62/64] test(train): assert model summary is logged --- tests/unit/commands/test_train_pytorch.py | 37 +++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index e3ba807..83c363b 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -166,6 +166,43 @@ def test_train_fragment_core_only_classification_head(tmp_path): ) +def test_train_fragment_core_logs_model_summary(tmp_path, monkeypatch, caplog): + """``train_fragment_core`` should log a Keras-style model summary.""" + from jaeger.commands import train as train_module + + # The jaeger logger disables propagation; re-enable it so caplog can capture + # the summary emitted by ``_print_model_summary``. + monkeypatch.setattr(train_module.logger, "propagate", True) + + train_path = tmp_path / "train.npz" + val_path = tmp_path / "val.npz" + _make_raw_npz(train_path, n_samples=8, seq_length=seq_length) + _make_raw_npz(val_path, n_samples=4, seq_length=seq_length) + + config_path = tmp_path / "config.yaml" + config = _build_config(tmp_path, train_path, val_path) + config_path.write_text(yaml.safe_dump(config)) + + with caplog.at_level("INFO"): + train_fragment_core( + config=str(config_path), + mixed_precision=False, + from_last_checkpoint=False, + force=False, + only_classification_head=False, + only_reliability_head=False, + only_heads=False, + only_save=True, + save_model=False, + self_supervised_pretraining=False, + xla=False, + meta=None, + ) + + assert "classifier model summary" in caplog.text + assert "Total params" in caplog.text + + def test_train_fragment_core_progress_bar(tmp_path): """``train_fragment_core`` should support the ``--progress-bar`` path.""" train_path = tmp_path / "train.npz" From 94cf8814e43c9a3c74988da1cb48f6cb57da9720 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:43:08 +0200 Subject: [PATCH 63/64] test: fix caplog capture for model summary warning tests --- tests/conftest.py | 8 ++++++++ tests/unit/commands/test_train_pytorch.py | 7 +------ tests/unit/nnlib/pytorch/test_summary.py | 8 ++++++-- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index c6748a8..85c4eb8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import shutil from pathlib import Path @@ -9,6 +10,13 @@ import pytest +@pytest.fixture +def propagate_jaeger_logger(monkeypatch): + """Re-enable propagation on the ``jaeger`` logger so caplog can capture.""" + logger = logging.getLogger("jaeger") + monkeypatch.setattr(logger, "propagate", True) + + @pytest.fixture def tmp_path_with_cleanup(tmp_path: Path): """Yield a temporary path and clean it up after the test.""" diff --git a/tests/unit/commands/test_train_pytorch.py b/tests/unit/commands/test_train_pytorch.py index 83c363b..ee701ec 100644 --- a/tests/unit/commands/test_train_pytorch.py +++ b/tests/unit/commands/test_train_pytorch.py @@ -166,13 +166,8 @@ def test_train_fragment_core_only_classification_head(tmp_path): ) -def test_train_fragment_core_logs_model_summary(tmp_path, monkeypatch, caplog): +def test_train_fragment_core_logs_model_summary(tmp_path, caplog, propagate_jaeger_logger): """``train_fragment_core`` should log a Keras-style model summary.""" - from jaeger.commands import train as train_module - - # The jaeger logger disables propagation; re-enable it so caplog can capture - # the summary emitted by ``_print_model_summary``. - monkeypatch.setattr(train_module.logger, "propagate", True) train_path = tmp_path / "train.npz" val_path = tmp_path / "val.npz" diff --git a/tests/unit/nnlib/pytorch/test_summary.py b/tests/unit/nnlib/pytorch/test_summary.py index d872143..ee3d8c2 100644 --- a/tests/unit/nnlib/pytorch/test_summary.py +++ b/tests/unit/nnlib/pytorch/test_summary.py @@ -57,7 +57,9 @@ def test_model_summary_returns_non_empty_string(dummy_input): assert len(text) > 0 -def test_model_summary_warns_when_torchinfo_missing(monkeypatch, caplog, dummy_input): +def test_model_summary_warns_when_torchinfo_missing( + monkeypatch, caplog, dummy_input, propagate_jaeger_logger +): import jaeger.nnlib.pytorch.summary as summary_mod monkeypatch.setattr(summary_mod, "torchinfo", None) @@ -68,7 +70,9 @@ def test_model_summary_warns_when_torchinfo_missing(monkeypatch, caplog, dummy_i assert "torchinfo is not installed" in caplog.text -def test_model_summary_warns_when_torchinfo_raises(monkeypatch, caplog, dummy_input): +def test_model_summary_warns_when_torchinfo_raises( + monkeypatch, caplog, dummy_input, propagate_jaeger_logger +): import jaeger.nnlib.pytorch.summary as summary_mod def _broken_summary(*args, **kwargs): From 0ce69ba079ac687db7edd38b69d6678d058f7741 Mon Sep 17 00:00:00 2001 From: Yasas1994 Date: Mon, 15 Jun 2026 23:46:39 +0200 Subject: [PATCH 64/64] docs: note implementation deviations in model summary design --- ...2026-06-15-pytorch-model-summary-design.md | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-15-pytorch-model-summary-design.md diff --git a/docs/superpowers/specs/2026-06-15-pytorch-model-summary-design.md b/docs/superpowers/specs/2026-06-15-pytorch-model-summary-design.md new file mode 100644 index 0000000..8429a2f --- /dev/null +++ b/docs/superpowers/specs/2026-06-15-pytorch-model-summary-design.md @@ -0,0 +1,123 @@ +# PyTorch Model Summary Design + +## Goal + +Provide a Keras `model.summary()`-style display of the Jaeger PyTorch model before training starts. The summary is printed automatically when `jaeger train` runs and is not suppressible. + +## Decision log + +- **Where it appears:** automatic print before `jaeger train` starts, for both classifier and reliability stages. +- **Content:** layers, output shapes, params per layer, total/trainable params. +- **Layout:** one table per training stage, with branch prefixes on layer names (`rep_model/...`, `classification_head/...`, `reliability_head/...`). +- **Implementation approach:** use the `torchinfo` third-party library. +- **Missing-library behavior:** warn and continue training. + +## Context + +Jaeger's PyTorch model is built by `jaeger.nnlib.pytorch.builder.ModelBuilder`. The top-level object used for classifier training is `JaegerModel`, which contains: + +- `rep_model` — `RepresentationModel` with embedding, blocks, and pooler +- `classification_head` — `ClassificationHead` +- `reliability_head` — optional `ReliabilityHead` + +Reliability training reuses `rep_model` through `_ReliabilityPipeline(rep_model, reliability_head)`. + +`src/jaeger/commands/train.py` already runs `_initialize_lazy_layers(models, config)` to materialize lazy layers before training. + +## Architecture + +### New module: `jaeger.nnlib.pytorch.summary` + +Create `src/jaeger/nnlib/pytorch/summary.py` containing a `ModelSummary` class. + +```python +class ModelSummary: + def __init__(self, model: nn.Module, input_data: Tuple[torch.Tensor, ...]): + self.model = model + self.input_data = input_data + + def summary(self, branch_label: str = "") -> str: + ... +``` + +Responsibilities: + +1. Import `torchinfo` at module level with a `None` fallback so the dependency remains optional and the module can be monkeypatched in tests. +2. Build a `torchinfo.summary(...)` call with: + - `input_data=self.input_data` + - `col_names=["input_size", "output_size", "num_params"]` + - `row_settings=["var_names"]` + - `depth=8` +3. Prepend a branch label header (e.g., `"=== Classifier model summary ==="`). +4. Return the formatted string. + +If `torchinfo` is unavailable or `torchinfo.summary` raises, log a warning and return an empty string so training continues. + +### Integration point + +Modify `src/jaeger/commands/train.py`: + +1. After `_initialize_lazy_layers(models, config)` succeeds and models are on the target device, construct a fresh dummy input matching the config's `input_type` and `input_shape`. To avoid duplication, factor the dummy-input construction out of `_initialize_lazy_layers` into a shared helper (e.g., `_make_dummy_input(config, device)`) that both `_initialize_lazy_layers` and `ModelSummary` use. +2. **Classifier stage:** before creating the `Trainer`, print the summary for `models["jaeger_classifier"]`. +3. **Reliability stage:** before creating the reliability `Trainer`, print the summary for `_ReliabilityPipeline(rep_model, models["reliability_head"])`. + +Both summaries are logged with `logger.info` so they appear in stdout and log files. + +### Data flow + +```text +config -> ModelBuilder -> models["jaeger_classifier"] + | + v + _initialize_lazy_layers (dummy forward pass) + | + v + ModelSummary(dummy input) + | + v + torchinfo.summary(model, input_data) + | + v + logger.info(output) + | + v + Trainer.fit() +``` + +## Error handling + +| Situation | Behavior | +|---|---| +| `torchinfo` not installed | Log warning, skip summary, continue training | +| `torchinfo.summary` raises (e.g., unsupported layer) | Catch exception, log warning with message, continue training | +| Dummy input shape mismatch | Already guarded by `_initialize_lazy_layers`; reuse same shape-building logic | +| Model has no parameters | `torchinfo` will show zero params; still printed | + +The summary step must not mutate model parameters, buffers, or optimizer state. + +## Dependencies + +Add `torchinfo>=1.8.0` to the main dependencies in `pyproject.toml`. + +## Testing + +### Unit tests + +1. **Happy path:** Build a minimal `JaegerModel`, call `ModelSummary(...).summary()`, assert the returned string contains `"Total params"` and is non-empty. +2. **Missing `torchinfo`:** Patch `jaeger.nnlib.pytorch.summary.torchinfo` to `None`, call `summary()`, assert a warning is logged and no exception is raised. +3. **Failing `torchinfo.summary`:** Patch `torchinfo.summary` to raise, assert a warning is logged and no exception propagates. + +### Smoke test + +Run a short `jaeger train` invocation (e.g., `--only_save` or `--epochs 0` if supported) with a tiny config and assert the summary header appears in the captured log output. + +## Open questions / future work + +- Should the summary also be exposed as a standalone CLI command (e.g., `jaeger summary --config ...`)? Out of scope for this design; can be added later without changing the API. +- Should the summary depth be configurable? For now fixed at `depth=8`; can become a config key if users request it. + +## Implementation notes + +- `torchinfo` is imported at module level in `src/jaeger/nnlib/pytorch/summary.py` (with a `None` fallback) rather than purely lazily inside `summary()`. This makes the module testable via monkeypatching while keeping the public API unchanged. +- `_initialize_lazy_layers` calls `_make_dummy_input(config, device)` without an explicit `length` argument, preserving the original per-input-type default lengths (64 for nucleotide, `crop_size // 3 - 1` for translated). +- Tests that assert on Jaeger log output use a shared `propagate_jaeger_logger` fixture in `tests/conftest.py`, because the Jaeger logger disables propagation by default.