update: aspanformer

third_party/ASpanFormer/src/ASpanFormer/aspan_module/attention.py

@@ -4,6 +4,7 @@ import torch.nn as nn
 from itertools import product
 from torch.nn import functional as F
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class layernorm2d(nn.Module):
     def __init__(self, dim):
@@ -176,7 +177,7 @@ class HierachicalAttention(Module):
         offset_sample = self.sample_offset[None, None] * span_scale
         sample_pixel = offset[:, :, None] + offset_sample  # B*G*r^2*2
         sample_norm = (
-            sample_pixel / torch.tensor([wk / 2, hk / 2]).cuda()[None, None, None] - 1
+            sample_pixel / torch.tensor([wk / 2, hk / 2]).to(device)[None, None, None] - 1
         )
 
         q = (

third_party/DeDoDe/DeDoDe/utils.py

@@ -11,6 +11,7 @@ from einops import rearrange
 import torch
 from time import perf_counter
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def recover_pose(E, kpts0, kpts1, K0, K1, mask):
     best_num_inliers = 0
@@ -54,7 +55,7 @@ def estimate_pose(kpts0, kpts1, K0, K1, norm_thresh, conf=0.99999):
     return ret
 
 
-def get_grid(B, H, W, device="cuda"):
+def get_grid(B, H, W, device=device):
     x1_n = torch.meshgrid(
         *[torch.linspace(-1 + 1 / n, 1 - 1 / n, n, device=device) for n in (B, H, W)]
     )
@@ -63,7 +64,7 @@ def get_grid(B, H, W, device="cuda"):
 
 
 @torch.no_grad()
-def finite_diff_hessian(f: tuple(["B", "H", "W"]), device="cuda"):
+def finite_diff_hessian(f: tuple(["B", "H", "W"]), device=device):
     dxx = (
         torch.tensor([[0, 0, 0], [1, -2, 1], [0, 0, 0]], device=device)[None, None] / 2
     )
@@ -78,7 +79,7 @@ def finite_diff_hessian(f: tuple(["B", "H", "W"]), device="cuda"):
     return H
 
 
-def finite_diff_grad(f: tuple(["B", "H", "W"]), device="cuda"):
+def finite_diff_grad(f: tuple(["B", "H", "W"]), device=device):
     dx = torch.tensor([[0, 0, 0], [-1, 0, 1], [0, 0, 0]], device=device)[None, None] / 2
     dy = dx.mT
     gx = F.conv2d(f[:, None], dx, padding=1)
@@ -103,7 +104,7 @@ def fast_inv_2x2(matrix: tuple[..., 2, 2], eps=1e-10):
     )
 
 
-def newton_step(f: tuple["B", "H", "W"], inds, device="cuda"):
+def newton_step(f: tuple["B", "H", "W"], inds, device=device):
     B, H, W = f.shape
     Hess = finite_diff_hessian(f).reshape(B, H * W, 2, 2)
     Hess = torch.gather(Hess, dim=1, index=inds[..., None].expand(B, -1, 2, 2))
@@ -118,7 +119,7 @@ def newton_step(f: tuple["B", "H", "W"], inds, device="cuda"):
 def sample_keypoints(
     scoremap,
     num_samples=8192,
-    device="cuda",
+    device=device,
     use_nms=True,
     sample_topk=False,
     return_scoremap=False,
@@ -176,7 +177,7 @@ def sample_keypoints(
 
 
 @torch.no_grad()
-def jacobi_determinant(warp, certainty, R=3, device="cuda", dtype=torch.float32):
+def jacobi_determinant(warp, certainty, R=3, device=device, dtype=torch.float32):
     t = perf_counter()
     *dims, _ = warp.shape
     warp = warp.to(dtype)
@@ -831,7 +832,7 @@ def homog_transform(Homog, x):
     return y
 
 
-def get_homog_warp(Homog, H, W, device="cuda"):
+def get_homog_warp(Homog, H, W, device=device):
     grid = torch.meshgrid(
         torch.linspace(-1 + 1 / H, 1 - 1 / H, H, device=device),
         torch.linspace(-1 + 1 / W, 1 - 1 / W, W, device=device),