(debug) to torch operations in bfloat conversions.

configs/default.yaml

@@ -28,8 +28,8 @@ inference:
   stitch: avg_logits
 
 eval:
-  max_samples: 16
-  fine_batch: 32
+  max_samples: 12
+  fine_batch: 16
 
 optim:
   iters: 2000
@@ -38,14 +38,14 @@ optim:
   weight_decay: 0.01
   schedule: poly
   power: 1.0
-  precision: fp32  # one of: fp32, fp16, bf16
+  precision: bf16  # one of: fp32, fp16, bf16
 
 # training housekeeping
 seed: 42
 out_dir: runs/wireseghr
-eval_interval: 100
+eval_interval: 150
 ckpt_interval: 300
-resume: runs/wireseghr/ckpt_1800.pt  # optional
+# resume: runs/wireseghr/ckpt_1800.pt  # optional
 
 # dataset paths (placeholders)
 data:

infer.py

@@ -31,7 +31,7 @@ def _coarse_forward(
     device: torch.device,
     amp_flag: bool,
     amp_dtype,
-) -> Tuple[np.ndarray, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     # Convert to tensor on device
     t_img = (
         torch.from_numpy(np.transpose(img_rgb, (2, 0, 1)))
@@ -76,8 +76,8 @@ def _coarse_forward(
         F.interpolate(prob, size=(H, W), mode="bilinear", align_corners=False)[0, 0]
         .detach()
         .cpu()
-        .numpy()
-    )
+        .float()
+    )  # HxW torch.Tensor on CPU
     return prob_up, cond_map, t_img, y_min_full, y_max_full
 
 
@@ -94,7 +94,7 @@ def _tiled_fine_forward(
     device: torch.device,
     amp_flag: bool,
     amp_dtype,
-) -> np.ndarray:
+) -> torch.Tensor:
     H = int(t_img.shape[2])
     W = int(t_img.shape[3])
     P = patch_size
@@ -153,8 +153,8 @@ def _tiled_fine_forward(
             prob_sum_t[y0:y1, x0:x1] += prob_f_up[bi]
             weight_t[y0:y1, x0:x1] += 1.0
 
-    prob_full = (prob_sum_t / weight_t).detach().cpu().numpy()
-    return prob_full
+    prob_full = (prob_sum_t / weight_t).detach().cpu().float()
+    return prob_full  # HxW torch.Tensor on CPU
 
 
 def _build_model_from_cfg(cfg: dict, device: torch.device) -> WireSegHR:
@@ -216,7 +216,9 @@ def infer_image(
         amp_dtype,
     )
 
-    pred = (prob_f > prob_thresh).astype(np.uint8) * 255
+    # Threshold with torch on CPU; convert to numpy only for saving/returning
+    pred_t = (prob_f > prob_thresh).to(torch.uint8) * 255  # HxW uint8 torch
+    pred = pred_t.detach().cpu().numpy()
 
     if out_dir is not None:
         os.makedirs(out_dir, exist_ok=True)
@@ -225,9 +227,10 @@ def infer_image(
         cv2.imwrite(out_mask, pred)
         if save_prob:
             out_prob = os.path.join(out_dir, f"{stem}_prob.npy")
-            np.save(out_prob, prob_f.astype(np.float32))
+            np.save(out_prob, prob_f.detach().cpu().float().numpy())
 
-    return pred, prob_f
+    # Return numpy arrays for external consumers, computed via torch
+    return pred, prob_f.detach().cpu().numpy()
 
 
 def main():

train.py

@@ -635,7 +635,7 @@ def validate(
             amp_dtype,
         )
         # Coarse metrics
-        pred_coarse = (prob_up > prob_thresh).astype(np.uint8)
+        pred_coarse = (prob_up > prob_thresh).to(torch.uint8).cpu().numpy()
         m_c = compute_metrics(pred_coarse, mask)
         for k in coarse_sum:
             coarse_sum[k] += m_c[k]
@@ -664,7 +664,7 @@ def validate(
         if xs[-1] != (W - P):
             xs.append(W - P)
         total_tiles += len(ys) * len(xs)
-        pred_fine = (prob_full > prob_thresh).astype(np.uint8)
+        pred_fine = (prob_full > prob_thresh).to(torch.uint8).cpu().numpy()
         m_f = compute_metrics(pred_fine, mask)
         for k in metrics_sum:
             metrics_sum[k] += m_f[k]
@@ -721,7 +721,7 @@ def save_test_visuals(
             bool(amp_flag),
             None,
         )
-        pred = (prob_up > prob_thresh).astype(np.uint8) * 255
+        pred = ((prob_up > prob_thresh).to(torch.uint8) * 255).cpu().numpy()
         # Save input and prediction
         img_bgr = (img[..., ::-1] * 255.0).astype(np.uint8)
         cv2.imwrite(os.path.join(out_dir, f"{i:03d}_input.jpg"), img_bgr)