Fix device handling

Author: isaacvock

src/transcriptml/devices.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import torch
+
+
+def resolve_device(device: str | torch.device = "cpu") -> torch.device:
+    """Resolve ``auto`` or explicit torch device names to a torch device."""
+
+    if isinstance(device, torch.device):
+        return device
+    name = str(device)
+    if name == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(name)

src/transcriptml/interpret/codon_ism.py

@@ -11,6 +11,7 @@
 
 from transcriptml.data.encoding import infer_valid_lengths
 from transcriptml.data.schemas import SequenceSchema, get_schema
+from transcriptml.devices import resolve_device
 from transcriptml.interpret.predictor import Predictor
 from transcriptml.progress import ProgressReporter, log_progress
 
@@ -844,7 +845,7 @@ def compute_codon_ism(
         raise ValueError("mutation_batch_size must be positive")
 
     if isinstance(predictor, torch.nn.Module):
-        predictor = Predictor(predictor, device=device, batch_size=reference_batch_size or 128)
+        predictor = Predictor(predictor, device=resolve_device(device), batch_size=reference_batch_size or 128)
 
     N, _, L = arr.shape
     analysis_indices = _resolve_analysis_indices(

src/transcriptml/interpret/predictor.py

@@ -6,6 +6,7 @@
 import numpy as np
 import torch
 
+from transcriptml.devices import resolve_device
 from transcriptml.models.common import squeeze_prediction
 from transcriptml.models.registry import load_checkpoint
 
@@ -30,7 +31,7 @@ def __init__(
         """
 
         self.model = model
-        self.device = torch.device(device)
+        self.device = resolve_device(device)
         self.batch_size = int(batch_size)
         if isinstance(model, torch.nn.Module):
             self.model.to(self.device)
@@ -53,8 +54,9 @@ def from_checkpoint(
             batch_size: Default batch size for prediction.
         """
 
-        model, _ = load_checkpoint(checkpoint_path, map_location=device)
-        return cls(model, device=device, batch_size=batch_size)
+        resolved = resolve_device(device)
+        model, _ = load_checkpoint(checkpoint_path, map_location=resolved)
+        return cls(model, device=resolved, batch_size=batch_size)
 
     @torch.no_grad()
     def predict(self, X: np.ndarray | torch.Tensor, *, batch_size: int | None = None) -> np.ndarray:

src/transcriptml/training/evaluation.py

@@ -8,6 +8,7 @@
 import torch
 
 from transcriptml.data.bundle import DatasetBundle, load_bundle
+from transcriptml.devices import resolve_device
 from transcriptml.models.common import squeeze_prediction
 from transcriptml.models.registry import load_checkpoint
 from transcriptml.training.metrics import mse, pearson_corr
@@ -33,6 +34,7 @@ def predict_array(
         progress: Whether to emit progress messages while predicting.
     """
 
+    device = resolve_device(device)
     model = model.to(device)
     model.eval()
     preds: list[np.ndarray] = []
@@ -75,6 +77,7 @@ def _predict_indexed_array(
         progress_label: Label shown in progress messages.
     """
 
+    device = resolve_device(device)
     model = model.to(device)
     model.eval()
     preds: list[np.ndarray] = []
@@ -115,6 +118,7 @@ def evaluate_model(
         progress: Whether to emit progress messages while evaluating.
     """
 
+    device = resolve_device(device)
     idx = np.arange(bundle.X.shape[0]) if indices is None else np.asarray(indices, dtype=int)
     preds = _predict_indexed_array(
         model,
@@ -193,6 +197,7 @@ def evaluate_checkpoint(
         progress: Whether to emit progress messages while evaluating.
     """
 
+    device = resolve_device(device)
     log_progress(f"evaluate: loading checkpoint {checkpoint_path}", enabled=progress)
     model, _ = load_checkpoint(checkpoint_path, map_location=device)
     log_progress(f"evaluate: loading dataset {dataset_path}", enabled=progress)

src/transcriptml/training/trainer.py

@@ -13,6 +13,7 @@
 
 from transcriptml.data.bundle import DatasetBundle, load_bundle
 from transcriptml.data.controls import apply_sequence_controls_to_bundle
+from transcriptml.devices import resolve_device
 from transcriptml.models.common import squeeze_prediction
 from transcriptml.models.registry import build_model, normalize_model_config, save_checkpoint
 from transcriptml.training.evaluation import evaluate_model, predict_to_csv
@@ -91,18 +92,6 @@ def _as_train_config(config: TrainConfig | Mapping[str, Any]) -> TrainConfig:
     return TrainConfig(**dict(config))
 
 
-def _select_device(name: str) -> torch.device:
-    """Resolve ``auto`` or explicit torch device names.
-
-    Args:
-        name: ``"auto"`` or any string accepted by ``torch.device``.
-    """
-
-    if name == "auto":
-        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    return torch.device(name)
-
-
 def _make_config_splits(bundle: DatasetBundle, cfg: TrainConfig) -> dict[str, list[int]]:
     """Choose dataset splits from the training config.
 
@@ -404,7 +393,7 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any])
 
     cfg = _as_train_config(config)
     _seed_everything(cfg.seed)
-    device = _select_device(cfg.device)
+    device = resolve_device(cfg.device)
     out = Path(cfg.output_dir)
     out.mkdir(parents=True, exist_ok=True)
     sequence_control_stats: dict[str, Any] | None = None

tests/test_device_auto.py

@@ -0,0 +1,73 @@
+import numpy as np
+
+from transcriptml.data.bundle import DatasetBundle, save_bundle
+from transcriptml.devices import resolve_device
+from transcriptml.interpret.predictor import Predictor
+from transcriptml.models.registry import build_model, save_checkpoint
+from transcriptml.training.evaluation import evaluate_checkpoint
+
+
+def _small_model_config() -> dict[str, object]:
+    return {
+        "name": "small_cnn",
+        "params": {
+            "in_ch": 4,
+            "n_filters": 4,
+            "kernel_size": 3,
+            "n_layers": 1,
+            "dropout": 0.0,
+            "head_hidden": 4,
+        },
+    }
+
+
+def _write_checkpoint_and_bundle(tmp_path):
+    model_config = _small_model_config()
+    model = build_model(model_config)
+    checkpoint = tmp_path / "model.pt"
+    save_checkpoint(checkpoint, model, model_config)
+
+    X = np.zeros((3, 4, 12), dtype=np.float32)
+    X[:, 0, :] = 1.0
+    bundle = DatasetBundle(
+        X=X,
+        y=np.array([0.0, 1.0, 2.0], dtype=np.float32),
+        ids=["a", "b", "c"],
+        schema="rna4",
+        splits={"train": [0], "val": [1], "test": [2]},
+    )
+    dataset = tmp_path / "dataset"
+    save_bundle(bundle, dataset)
+    return checkpoint, dataset
+
+
+def test_resolve_device_auto_cpu_when_cuda_unavailable(monkeypatch):
+    monkeypatch.setattr("torch.cuda.is_available", lambda: False)
+    assert str(resolve_device("auto")) == "cpu"
+
+
+def test_evaluate_checkpoint_accepts_auto_device(tmp_path, monkeypatch):
+    monkeypatch.setattr("torch.cuda.is_available", lambda: False)
+    checkpoint, dataset = _write_checkpoint_and_bundle(tmp_path)
+
+    result = evaluate_checkpoint(
+        checkpoint,
+        dataset,
+        tmp_path / "predictions.csv",
+        split="test",
+        device="auto",
+        progress=False,
+    )
+
+    assert len(result["predictions"]) == 1
+    assert (tmp_path / "predictions.csv").exists()
+
+
+def test_predictor_from_checkpoint_accepts_auto_device(tmp_path, monkeypatch):
+    monkeypatch.setattr("torch.cuda.is_available", lambda: False)
+    checkpoint, _ = _write_checkpoint_and_bundle(tmp_path)
+
+    predictor = Predictor.from_checkpoint(checkpoint, device="auto", batch_size=2)
+    preds = predictor.predict(np.zeros((2, 4, 12), dtype=np.float32))
+
+    assert preds.shape == (2,)