From 55222540d8033c9596f3a3c728cf0dd347aef63f Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 17 Jul 2026 12:02:52 +0200 Subject: [PATCH 1/6] Fix top-down pose postprocessing device Update PyTorch top-down pose postprocessing to keep tensors on the same device/dtype by using `new_zeros` and `new_tensor` instead of creating CPU tensors. Also replace `torch.cat(poses)` with `torch.stack(..., dim=0)` so batched pose outputs preserve the expected batch dimension. --- dlclive/pose_estimation_pytorch/runner.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/dlclive/pose_estimation_pytorch/runner.py b/dlclive/pose_estimation_pytorch/runner.py index 0c6b169..2befe25 100644 --- a/dlclive/pose_estimation_pytorch/runner.py +++ b/dlclive/pose_estimation_pytorch/runner.py @@ -357,21 +357,24 @@ def _postprocess_top_down( """Post-processes pose for top-down models.""" if len(batch_pose) == 0: bodyparts, coords = batch_pose.shape[-2:] - return torch.zeros((0, bodyparts, coords)) + return batch_pose.new_zeros((0, bodyparts, coords)) poses = [] + for pose, (offset, scale) in zip(batch_pose, offsets_and_scales, strict=False): - poses.append( - torch.cat( - [ - pose[..., :2] * torch.tensor(scale) + torch.tensor(offset), - pose[..., 2:3], - ], - dim=-1, - ) + scale_tensor = pose.new_tensor(scale) + offset_tensor = pose.new_tensor(offset) + + transformed_pose = torch.cat( + [ + pose[..., :2] * scale_tensor + offset_tensor, + pose[..., 2:3], + ], + dim=-1, ) + poses.append(transformed_pose) - return torch.cat(poses) + return torch.stack(poses, dim=0) def _parse_device(device: str | None) -> str: From 216bdb22146a7fb4f9289fb4465f2413485fb7e8 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 17 Jul 2026 12:04:26 +0200 Subject: [PATCH 2/6] Make SkipFrames bbox tensors device-safe Refactor bbox creation in `SkipFrames.__call__` to build tensors from `pose` (`new_zeros/new_tensor/new_ones`) so boxes and scores stay on the same device and dtype. The NaN handling for min/max coordinates is now explicit via `nan=` arguments, and bbox limiting uses `torch.clamp` with typed min/max tensors. --- dlclive/pose_estimation_pytorch/runner.py | 32 +++++++++++++++++++---- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/dlclive/pose_estimation_pytorch/runner.py b/dlclive/pose_estimation_pytorch/runner.py index 2befe25..e13b8fb 100644 --- a/dlclive/pose_estimation_pytorch/runner.py +++ b/dlclive/pose_estimation_pytorch/runner.py @@ -72,11 +72,33 @@ def update(self, pose: torch.Tensor, w: int, h: int) -> None: num_det, num_kpts = pose.shape[:2] size = max(w, h) - bboxes = torch.zeros((num_det, 4)) - bboxes[:, :2] = torch.min(torch.nan_to_num(pose, size)[..., :2], dim=1)[0] - self.margin - bboxes[:, 2:4] = torch.max(torch.nan_to_num(pose, 0)[..., :2], dim=1)[0] + self.margin - bboxes = torch.clip(bboxes, min=torch.zeros(4), max=torch.tensor([w, h, w, h])) - self._detections = dict(boxes=bboxes, scores=torch.ones(num_det)) + bboxes = pose.new_zeros((num_det, 4)) + + bboxes[:, :2] = ( + torch.min( + torch.nan_to_num(pose[..., :2], nan=float(size)), + dim=1, + )[0] + - self.margin + ) + + bboxes[:, 2:4] = ( + torch.max( + torch.nan_to_num(pose[..., :2], nan=0.0), + dim=1, + )[0] + + self.margin + ) + + minimum = pose.new_zeros(4) + maximum = pose.new_tensor([w, h, w, h]) + + bboxes = torch.clamp(bboxes, min=minimum, max=maximum) + + self._detections = { + "boxes": bboxes, + "scores": pose.new_ones(num_det), + } self._age += 1 From d35b17f125995b331ce94ddbece84f5baa8eca04 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 17 Jul 2026 12:04:45 +0200 Subject: [PATCH 3/6] Fix empty top-down offsets/scales default Update the PyTorch top-down path to use a correctly structured fallback `offsets_and_scales` value when no detections are found. This keeps the data shape consistent with post-processing expectations and avoids downstream tuple unpacking/type mismatches. --- dlclive/pose_estimation_pytorch/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dlclive/pose_estimation_pytorch/runner.py b/dlclive/pose_estimation_pytorch/runner.py index e13b8fb..98bab95 100644 --- a/dlclive/pose_estimation_pytorch/runner.py +++ b/dlclive/pose_estimation_pytorch/runner.py @@ -210,7 +210,7 @@ def get_pose(self, frame: np.ndarray) -> np.ndarray: frame_batch, offsets_and_scales = self._prepare_top_down(tensor, detections) if len(frame_batch) == 0: - offsets_and_scales = [(0, 0), 1] + offsets_and_scales = [((0, 0), (1.0, 1.0))] tensor = frame_batch # still CHW, batched if self.dynamic is not None: From 305b86393d8e7b7dd9178aa6007f311a555c410e Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 17 Jul 2026 13:50:56 +0200 Subject: [PATCH 4/6] Validate frame input and fix scale defaults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Improve the PyTorch runner’s frame handling by validating HWC input shape and allowed channel counts, correctly unpacking dimensions, and dropping alpha for 4-channel frames. Also fix the top-down fallback `offsets_and_scales` format to use paired float scales instead of a scalar, preventing downstream shape/format issues. --- dlclive/pose_estimation_pytorch/runner.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/dlclive/pose_estimation_pytorch/runner.py b/dlclive/pose_estimation_pytorch/runner.py index 98bab95..5f06d27 100644 --- a/dlclive/pose_estimation_pytorch/runner.py +++ b/dlclive/pose_estimation_pytorch/runner.py @@ -11,6 +11,7 @@ """PyTorch and ONNX runners for DeepLabCut-Live""" import copy +import logging from dataclasses import dataclass from pathlib import Path from typing import Literal @@ -25,6 +26,8 @@ from dlclive.core.runner import BaseRunner from dlclive.pose_estimation_pytorch.data.image import AutoPadToDivisor +logger = logging.getLogger(__name__) + @dataclass class SkipFrames: @@ -192,8 +195,18 @@ def close(self) -> None: @torch.inference_mode() def get_pose(self, frame: np.ndarray) -> np.ndarray: - c, h, w = frame.shape - tensor = torch.from_numpy(frame).permute(2, 0, 1) # CHW, still on CPU + if frame.ndim != 3: + raise ValueError(f"Expected HWC frame, got shape {frame.shape}") + + h, w, c = frame.shape + + if c not in (1, 3, 4): + raise ValueError(f"Expected 1, 3, or 4 channels in the last dimension, got shape {frame.shape}") + + if c == 4: + frame = frame[..., :3] + + tensor = torch.from_numpy(frame).permute(2, 0, 1) offsets_and_scales = None if self.detector is not None: @@ -367,7 +380,7 @@ def _prepare_top_down(self, frame: torch.Tensor, detections: dict[str, torch.Ten else: crop_w, crop_h = self.top_down_config.crop_size frame_batch = torch.zeros((1, 3, crop_h, crop_w), device=frame.device) - offsets_and_scales = [(0, 0), 1] + offsets_and_scales = [(0, 0), (1.0, 1.0)] return frame_batch, offsets_and_scales From 8cdb1810fe6f49e59a4bfb207b7eca37b501077d Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 17 Jul 2026 13:51:56 +0200 Subject: [PATCH 5/6] Switch pose aggregation from stack to cat In `PyTorchRunner`, transformed pose tensors are now merged with `torch.cat(..., dim=0)` instead of `torch.stack`. This aligns aggregation with tensors that already carry a leading batch dimension and avoids introducing an extra axis. --- dlclive/pose_estimation_pytorch/runner.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dlclive/pose_estimation_pytorch/runner.py b/dlclive/pose_estimation_pytorch/runner.py index 5f06d27..ed85e7a 100644 --- a/dlclive/pose_estimation_pytorch/runner.py +++ b/dlclive/pose_estimation_pytorch/runner.py @@ -409,7 +409,10 @@ def _postprocess_top_down( ) poses.append(transformed_pose) - return torch.stack(poses, dim=0) + result = torch.cat(poses, dim=0) + + if result.ndim != 3: + raise ValueError(f"Expected result of shape (detections, bodyparts, coords), got {tuple(result.shape)}") def _parse_device(device: str | None) -> str: From f34d980f24bac5a9cea964b79449f1e34a30d411 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 17 Jul 2026 14:35:00 +0200 Subject: [PATCH 6/6] Add timing diagnostics to PyTorch runner Introduce a reusable `DebugTimingStats` utility in `dlclive.utils` and instrument `PyTorchRunner.get_pose` with stage-level timing (input conversion, detector, pose transform/inference/decode, dynamic crop/update, postprocess, and CPU transfer). The runner now synchronizes CUDA around timed inference sections for accurate measurements and periodically logs aggregated performance stats. This also fixes top-down fallback metadata shape (`offsets_and_scales`) and restores the missing return in `_postprocess_top_down`. --- dlclive/pose_estimation_pytorch/runner.py | 169 +++++++++++++++------- dlclive/utils.py | 103 +++++++++++++ 2 files changed, 219 insertions(+), 53 deletions(-) diff --git a/dlclive/pose_estimation_pytorch/runner.py b/dlclive/pose_estimation_pytorch/runner.py index ed85e7a..932dd36 100644 --- a/dlclive/pose_estimation_pytorch/runner.py +++ b/dlclive/pose_estimation_pytorch/runner.py @@ -25,6 +25,7 @@ import dlclive.pose_estimation_pytorch.models as models from dlclive.core.runner import BaseRunner from dlclive.pose_estimation_pytorch.data.image import AutoPadToDivisor +from dlclive.utils import DebugTimingStats logger = logging.getLogger(__name__) @@ -168,6 +169,8 @@ def __init__( self.detector_transform = None self.pose_transform = None + self._timing = DebugTimingStats("PyTorchRunner", logger=logger, enabled=True, log_interval=2) + # Parse Dynamic Cropping parameters if isinstance(dynamic, dict): dynamic_type = dynamic.get("type", "DynamicCropper") @@ -189,77 +192,135 @@ def __init__( self.dynamic = dynamic self.top_down_config = top_down_config + def _sync_cuda_for_timing(self) -> None: + if self._timing.enabled and str(self.device).startswith("cuda"): + torch.cuda.synchronize() + def close(self) -> None: """Clears any resources used by the runner.""" pass @torch.inference_mode() def get_pose(self, frame: np.ndarray) -> np.ndarray: - if frame.ndim != 3: - raise ValueError(f"Expected HWC frame, got shape {frame.shape}") + with self._timing.measure("get_pose.total"): + if frame.ndim != 3: + raise ValueError(f"Expected HWC frame, got shape {frame.shape}") + + h, w, c = frame.shape + + if c not in (1, 3, 4): + raise ValueError(f"Expected 1, 3, or 4 channels in the last dimension, got shape {frame.shape}") + + if c == 4: + frame = frame[..., :3] + + with self._timing.measure("input.numpy_to_torch"): + tensor = torch.from_numpy(frame).permute(2, 0, 1) + + offsets_and_scales = None + + if self.detector is not None: + detections = None + + if self.top_down_config.skip_frames is not None: + with self._timing.measure("skip_frames.get_detections"): + detections = self.top_down_config.skip_frames.get_detections() + + if detections is None: + with self._timing.measure("detector.transform"): + detector_input = self.detector_transform(tensor).unsqueeze(0) + + with self._timing.measure("detector.to_device"): + detector_input = detector_input.to(self.device) + + if self.precision == "FP16": + detector_input = detector_input.half() - h, w, c = frame.shape + self._sync_cuda_for_timing() - if c not in (1, 3, 4): - raise ValueError(f"Expected 1, 3, or 4 channels in the last dimension, got shape {frame.shape}") + with self._timing.measure("detector.inference"): + detections = self.detector(detector_input)[0] - if c == 4: - frame = frame[..., :3] + self._sync_cuda_for_timing() - tensor = torch.from_numpy(frame).permute(2, 0, 1) + with self._timing.measure("top_down.prepare"): + frame_batch, offsets_and_scales = self._prepare_top_down( + tensor, + detections, + ) - offsets_and_scales = None - if self.detector is not None: - detections = None - if self.top_down_config.skip_frames is not None: - detections = self.top_down_config.skip_frames.get_detections() + tensor = frame_batch + + if self.dynamic is not None: + with self._timing.measure("dynamic.crop"): + tensor = self.dynamic.crop(tensor) + + with self._timing.measure("pose.transform"): + model_input = self.pose_transform(tensor) + + if model_input.dim() == 3: + model_input = model_input.unsqueeze(0) + + with self._timing.measure("pose.to_device"): + model_input = model_input.to(self.device) - if detections is None: - # Apply detector transform before inference - detector_input = self.detector_transform(tensor).unsqueeze(0).to(self.device) if self.precision == "FP16": - detector_input = detector_input.half() - detections = self.detector(detector_input)[0] - - frame_batch, offsets_and_scales = self._prepare_top_down(tensor, detections) - if len(frame_batch) == 0: - offsets_and_scales = [((0, 0), (1.0, 1.0))] - tensor = frame_batch # still CHW, batched - - if self.dynamic is not None: - tensor = self.dynamic.crop(tensor) - - # Apply pose transform - model_input = self.pose_transform(tensor) - # Ensure 4D input: (N, C, H, W) - if model_input.dim() == 3: - model_input = model_input.unsqueeze(0) - # Send to device - model_input = model_input.to(self.device) - if self.precision == "FP16": - model_input = model_input.half() + model_input = model_input.half() - outputs = self.model(model_input) - batch_pose = self.model.get_predictions(outputs)["bodypart"]["poses"] + self._sync_cuda_for_timing() - if self.dynamic is not None: - batch_pose = self.dynamic.update(batch_pose) + with self._timing.measure("pose.inference"): + outputs = self.model(model_input) + self._sync_cuda_for_timing() - if self.detector is None: - pose = batch_pose[0] - else: - pose = self._postprocess_top_down(batch_pose, offsets_and_scales) - if self.top_down_config.skip_frames is not None: - self.top_down_config.skip_frames.update(pose, w, h) + with self._timing.measure("pose.decode"): + batch_pose = self.model.get_predictions(outputs)["bodypart"]["poses"] - if self.single_animal: - if len(pose) == 0: - bodyparts, coords = pose.shape[-2:] - return np.zeros((bodyparts, coords)) + if self.dynamic is not None: + with self._timing.measure("dynamic.update"): + batch_pose = self.dynamic.update(batch_pose) - pose = pose[0] + if self.detector is None: + pose = batch_pose[0] + else: + with self._timing.measure("top_down.postprocess"): + pose = self._postprocess_top_down( + batch_pose, + offsets_and_scales, + ) + + if self.top_down_config.skip_frames is not None: + with self._timing.measure("skip_frames.update"): + self.top_down_config.skip_frames.update( + pose, + w, + h, + ) + + if self.single_animal: + if len(pose) == 0: + bodyparts, coords = pose.shape[-2:] + result = np.zeros( + (bodyparts, coords), + dtype=np.float32, + ) + else: + pose = pose[0] + + self._sync_cuda_for_timing() + + with self._timing.measure("output.to_cpu"): + result = pose.cpu().numpy() + else: + self._sync_cuda_for_timing() + + with self._timing.measure("output.to_cpu"): + result = pose.cpu().numpy() - return pose.cpu().numpy() + self._timing.note_frame() + self._timing.maybe_log() + + return result def init_inference(self, frame: np.ndarray, **kwargs) -> np.ndarray: """ @@ -380,7 +441,7 @@ def _prepare_top_down(self, frame: torch.Tensor, detections: dict[str, torch.Ten else: crop_w, crop_h = self.top_down_config.crop_size frame_batch = torch.zeros((1, 3, crop_h, crop_w), device=frame.device) - offsets_and_scales = [(0, 0), (1.0, 1.0)] + offsets_and_scales = [((0, 0), (1.0, 1.0))] return frame_batch, offsets_and_scales @@ -414,6 +475,8 @@ def _postprocess_top_down( if result.ndim != 3: raise ValueError(f"Expected result of shape (detections, bodyparts, coords), got {tuple(result.shape)}") + return result + def _parse_device(device: str | None) -> str: if device is None: diff --git a/dlclive/utils.py b/dlclive/utils.py index 98fe45f..3aff212 100644 --- a/dlclive/utils.py +++ b/dlclive/utils.py @@ -5,6 +5,8 @@ Licensed under GNU Lesser General Public License v3.0 """ +import logging +import time import urllib.error import urllib.request import warnings @@ -308,3 +310,104 @@ def download_file(url: str, filepath: str, chunk_size: int = 8192) -> None: raise urllib.error.URLError(f"Failed to download from {url}: {e.reason}") from e except OSError as e: raise OSError(f"Failed to write file to {filepath}") from e + + +class DebugTimingStats: + """Tiny timing accumulator for camera worker performance diagnostics. + + Usage: + with stats.measure("read"): + frame, ts = backend.read() + + Logs aggregate timings once per log_interval seconds. + """ + + def __init__( + self, info: str, *, logger: logging.Logger | None = None, log_interval: float = 1.0, enabled: bool = True + ): + self.info = info + self.log_interval = float(log_interval) + self.enabled = bool(enabled) + self.logger = logger or logging.getLogger(__name__) + if self.enabled: # force logger to proper level + if not self.logger.isEnabledFor(logging.DEBUG): + self.logger.setLevel(logging.DEBUG) + + self._last_log = time.perf_counter() + self._frames = 0 + self._timeouts = 0 + self._errors = 0 + self._totals: dict[str, float] = {} + self._counts: dict[str, int] = {} + + class _Measure: + def __init__(self, parent: "DebugTimingStats", name: str): + self.parent = parent + self.name = name + self.t0 = 0.0 + self.elapsed = 0.0 + + def __enter__(self): + if self.parent.enabled: + self.t0 = time.perf_counter() + return self + + def __exit__(self, exc_type, exc, tb): + if not self.parent.enabled: + return False + + self.elapsed = time.perf_counter() - self.t0 + self.parent._totals[self.name] = self.parent._totals.get(self.name, 0.0) + self.elapsed + self.parent._counts[self.name] = self.parent._counts.get(self.name, 0) + 1 + return False + + def measure(self, name: str): + return self._Measure(self, name) + + def note_frame(self) -> None: + if self.enabled: + self._frames += 1 + + def note_timeout(self) -> None: + if self.enabled: + self._timeouts += 1 + + def note_error(self) -> None: + if self.enabled: + self._errors += 1 + + def maybe_log(self) -> None: + if not self.enabled: + return + + now = time.perf_counter() + elapsed = now - self._last_log + if elapsed < self.log_interval: + return + + fps = self._frames / max(elapsed, 1e-9) + + parts = [ + f"[Worker {self.info}]", + f"fps={fps:.1f}", + f"frames={self._frames}", + ] + + if self._timeouts: + parts.append(f"timeouts={self._timeouts}") + if self._errors: + parts.append(f"errors={self._errors}") + + for name in sorted(self._totals): + count = max(self._counts.get(name, 0), 1) + avg_ms = 1000.0 * self._totals[name] / count + parts.append(f"avg_{name}_ms={avg_ms:.3f}") + + self.logger.debug(" ".join(parts)) + + self._last_log = now + self._frames = 0 + self._timeouts = 0 + self._errors = 0 + self._totals.clear() + self._counts.clear()