perf(unraid): enable Ampere mixed precision inference

2026-08-15 10:18:23 -07:00
parent cd29046592
commit 593186b325
2 changed files with 10 additions and 5 deletions
@@ -29,6 +29,9 @@ class Models:
    def load(self) -> None:
        from mmdet.apis import init_detector
        from mmpose.apis import init_model
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
        self.detector = init_detector(self.det_config, self.det_checkpoint, device="cuda:0")
        self.detector.test_cfg.nms.iou_threshold = float(os.getenv("NMS_IOU_THRESHOLD", ".60"))
        self.pose = init_model(self.pose_config, self.pose_checkpoint, device="cuda:0")
@@ -41,7 +44,8 @@ class Models:
        height, width = frame.shape[:2]
        t0 = time.perf_counter_ns()
        init_default_scope("mmdet")
-        result = inference_detector(self.detector, frame)
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            result = inference_detector(self.detector, frame)
        pred = result.pred_instances.cpu().numpy()
        keep = (pred.labels == 0) & (pred.scores >= threshold)
        boxes = pred.bboxes[keep]
@@ -52,7 +56,8 @@ class Models:
        det_ms = (time.perf_counter_ns() - t0) / 1e6
        t1 = time.perf_counter_ns()
        init_default_scope("mmpose")
-        pose_results = inference_topdown(self.pose, frame, bboxes=boxes) if len(boxes) else []
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            pose_results = inference_topdown(self.pose, frame, bboxes=boxes) if len(boxes) else []
        pose_ms = (time.perf_counter_ns() - t1) / 1e6
        people = []
        for idx, (box, det_score, sample) in enumerate(zip(boxes, scores, pose_results)):