feat(pytorch): migrate training and inference to PyTorch (v1.27.1)#34
Open
Yasas1994 wants to merge 64 commits into
Open
feat(pytorch): migrate training and inference to PyTorch (v1.27.1)#34Yasas1994 wants to merge 64 commits into
Yasas1994 wants to merge 64 commits into
Conversation
… 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.
…eline 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.
- 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.
…s 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.
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.
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.
- 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).
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.
…eps 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.
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.
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.
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.
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.
--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.
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.
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.
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
src/jaeger/nnlib/pytorch/).src/jaeger/data/pytorch/) supportingnumpy_full,numpy_raw, andcsvformats.Trainer, callbacks, and distributed helper (src/jaeger/training/pytorch/).jaeger trainandjaeger predictCLI commands to use PyTorch.Test Plan
pytest tests/unit/ tests/integration/passes (218 tests)ruff check src/ tests/andruff format --check src/ tests/passpython -m build --wheel)Breaking Changes
--xla,--quantized,--onnx,--int8injaeger predictare no longer supported.