add explicit CPU for persistent buffer

Author: gau-nernst

timm/layers/lambda_layer.py

@@ -31,7 +31,7 @@
 
 def rel_pos_indices(size):
     size = to_2tuple(size)
-    pos = torch.stack(ndgrid(torch.arange(size[0]), torch.arange(size[1]))).flatten(1)
+    pos = torch.stack(ndgrid(torch.arange(size[0], device="cpu"), torch.arange(size[1], device="cpu"))).flatten(1)
     rel_pos = pos[:, None, :] - pos[:, :, None]
     rel_pos[0] += size[0] - 1
     rel_pos[1] += size[1] - 1

timm/layers/pos_embed_rel.py

@@ -27,7 +27,7 @@ def gen_relative_position_index(
     # get pair-wise relative position index for each token inside the window
     assert k_size is None, 'Different q & k sizes not currently supported'  # FIXME
 
-    coords = torch.stack(ndgrid(torch.arange(q_size[0]), torch.arange(q_size[1]))).flatten(1)  # 2, Wh, Ww
+    coords = torch.stack(ndgrid(torch.arange(q_size[0], device="cpu"), torch.arange(q_size[1], device="cpu"))).flatten(1)  # 2, Wh, Ww
     relative_coords = coords[:, :, None] - coords[:, None, :]  # 2, Wh*Ww, Wh*Ww
     relative_coords = relative_coords.permute(1, 2, 0)  # Qh*Qw, Kh*Kw, 2
     relative_coords[:, :, 0] += q_size[0] - 1  # shift to start from 0
@@ -307,8 +307,8 @@ def gen_relative_log_coords(
 ):
     assert mode in ('swin', 'cr')
     # as per official swin-v2 impl, supporting timm specific 'cr' log coords as well
-    relative_coords_h = torch.arange(-(win_size[0] - 1), win_size[0]).to(torch.float32)
-    relative_coords_w = torch.arange(-(win_size[1] - 1), win_size[1]).to(torch.float32)
+    relative_coords_h = torch.arange(-(win_size[0] - 1), win_size[0], device="cpu").to(torch.float32)
+    relative_coords_w = torch.arange(-(win_size[1] - 1), win_size[1], device="cpu").to(torch.float32)
     relative_coords_table = torch.stack(ndgrid(relative_coords_h, relative_coords_w))
     relative_coords_table = relative_coords_table.permute(1, 2, 0).contiguous()  # 2*Wh-1, 2*Ww-1, 2
     if mode == 'swin':
@@ -415,7 +415,7 @@ def generate_lookup_tensor(
         max_relative_position = length - 1
     # Return the cached lookup tensor, otherwise compute it and cache it.
     vocab_size = 2 * max_relative_position + 1
-    ret = torch.zeros(length, length, vocab_size)
+    ret = torch.zeros(length, length, vocab_size, device="cpu")
     for i in range(length):
         for x in range(length):
             v = x - i + max_relative_position

timm/layers/pos_embed_sincos.py

@@ -163,7 +163,7 @@ def __init__(
         self.keep_spatial = keep_spatial
         self.register_buffer(
             'bands',
-            pixel_freq_bands(max_res, num_bands),
+            pixel_freq_bands(max_res, num_bands, device="cpu"),
             persistent=False,
         )
 
@@ -305,12 +305,14 @@ def __init__(
                     dim // 4,
                     float(max_res),
                     linear_bands=linear_bands,
+                    device="cpu",
                 )
             else:
                 bands = freq_bands(
                     dim // 4,
                     temperature=temperature,
                     step=1,
+                    device="cpu",
                 )
             self.register_buffer(
                 'bands',
@@ -328,6 +330,7 @@ def __init__(
                 linear_bands=linear_bands,
                 in_pixels=in_pixels,
                 ref_feat_shape=self.ref_feat_shape,
+                device="cpu",
             )
             self.bands = None
             self.register_buffer(
@@ -392,12 +395,14 @@ def __init__(
                     dim // 4,
                     float(max_res),
                     linear_bands=linear_bands,
+                    device="cpu",
                 )
             else:
                 bands = freq_bands(
                     dim // 4,
                     temperature=temperature,
                     step=1,
+                    device="cpu",
                 )
             self.register_buffer(
                 'bands',
@@ -414,6 +419,7 @@ def __init__(
                 linear_bands=linear_bands,
                 in_pixels=in_pixels,
                 ref_feat_shape=self.ref_feat_shape,
+                device="cpu",
             )
             self.bands = None
             self.register_buffer(

timm/models/beit.py

@@ -63,7 +63,7 @@ def gen_relative_position_index(window_size: Tuple[int, int]) -> torch.Tensor:
     # cls to token & token 2 cls & cls to cls
     # get pair-wise relative position index for each token inside the window
     window_area = window_size[0] * window_size[1]
-    coords = torch.stack(ndgrid(torch.arange(window_size[0]), torch.arange(window_size[1])))  # 2, Wh, Ww
+    coords = torch.stack(ndgrid(torch.arange(window_size[0], device="cpu"), torch.arange(window_size[1], device="cpu")))  # 2, Wh, Ww
     coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
     relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
     relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
@@ -105,7 +105,7 @@ def __init__(
         self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
         if qkv_bias:
             self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
-            self.register_buffer('k_bias', torch.zeros(all_head_dim), persistent=False)
+            self.register_buffer('k_bias', torch.zeros(all_head_dim, device="cpu"), persistent=False)
             self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
         else:
             self.q_bias = None

timm/models/efficientformer_v2.py

@@ -131,7 +131,7 @@ def __init__(
         self.act = act_layer()
         self.proj = ConvNorm(self.dh, dim, 1)
 
-        pos = torch.stack(ndgrid(torch.arange(self.resolution[0]), torch.arange(self.resolution[1]))).flatten(1)
+        pos = torch.stack(ndgrid(torch.arange(self.resolution[0], device="cpu"), torch.arange(self.resolution[1], device="cpu"))).flatten(1)
         rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
         rel_pos = (rel_pos[0] * self.resolution[1]) + rel_pos[1]
         self.attention_biases = torch.nn.Parameter(torch.zeros(num_heads, self.N))
@@ -233,10 +233,10 @@ def __init__(
         self.proj = ConvNorm(self.dh, self.out_dim, 1)
 
         self.attention_biases = nn.Parameter(torch.zeros(num_heads, self.N))
-        k_pos = torch.stack(ndgrid(torch.arange(self.resolution[0]), torch.arange(self.resolution[1]))).flatten(1)
+        k_pos = torch.stack(ndgrid(torch.arange(self.resolution[0], device="cpu"), torch.arange(self.resolution[1], device="cpu"))).flatten(1)
         q_pos = torch.stack(ndgrid(
-            torch.arange(0, self.resolution[0], step=2),
-            torch.arange(0, self.resolution[1], step=2)
+            torch.arange(0, self.resolution[0], step=2, device="cpu"),
+            torch.arange(0, self.resolution[1], step=2, device="cpu")
         )).flatten(1)
         rel_pos = (q_pos[..., :, None] - k_pos[..., None, :]).abs()
         rel_pos = (rel_pos[0] * self.resolution[1]) + rel_pos[1]

timm/models/eva.py

@@ -88,7 +88,7 @@ def __init__(
             self.q_proj = self.k_proj = self.v_proj = None
             if qkv_bias:
                 self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
-                self.register_buffer('k_bias', torch.zeros(all_head_dim), persistent=False)
+                self.register_buffer('k_bias', torch.zeros(all_head_dim, device="cpu"), persistent=False)
                 self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
             else:
                 self.q_bias = self.k_bias = self.v_bias = None

timm/models/levit.py

@@ -195,7 +195,7 @@ def __init__(
         ]))
 
         self.attention_biases = nn.Parameter(torch.zeros(num_heads, resolution[0] * resolution[1]))
-        pos = torch.stack(ndgrid(torch.arange(resolution[0]), torch.arange(resolution[1]))).flatten(1)
+        pos = torch.stack(ndgrid(torch.arange(resolution[0], device="cpu"), torch.arange(resolution[1], device="cpu"))).flatten(1)
         rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
         rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
         self.register_buffer('attention_bias_idxs', rel_pos, persistent=False)
@@ -291,10 +291,10 @@ def __init__(
         ]))
 
         self.attention_biases = nn.Parameter(torch.zeros(num_heads, resolution[0] * resolution[1]))
-        k_pos = torch.stack(ndgrid(torch.arange(resolution[0]), torch.arange(resolution[1]))).flatten(1)
+        k_pos = torch.stack(ndgrid(torch.arange(resolution[0], device="cpu"), torch.arange(resolution[1], device="cpu"))).flatten(1)
         q_pos = torch.stack(ndgrid(
-            torch.arange(0, resolution[0], step=stride),
-            torch.arange(0, resolution[1], step=stride)
+            torch.arange(0, resolution[0], step=stride, device="cpu"),
+            torch.arange(0, resolution[1], step=stride, device="cpu"),
         )).flatten(1)
         rel_pos = (q_pos[..., :, None] - k_pos[..., None, :]).abs()
         rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]

timm/models/swin_transformer.py

@@ -79,7 +79,7 @@ def window_reverse(windows, window_size: Tuple[int, int], H: int, W: int):
 
 def get_relative_position_index(win_h: int, win_w: int):
     # get pair-wise relative position index for each token inside the window
-    coords = torch.stack(ndgrid(torch.arange(win_h), torch.arange(win_w)))  # 2, Wh, Ww
+    coords = torch.stack(ndgrid(torch.arange(win_h, device="cpu"), torch.arange(win_w, device="cpu")))  # 2, Wh, Ww
     coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
     relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
     relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
@@ -291,7 +291,7 @@ def get_attn_mask(self, x: Optional[torch.Tensor] = None) -> Optional[torch.Tens
                 dtype = x.dtype
             else:
                 H, W = self.input_resolution
-                device = None
+                device = "cpu"
                 dtype = None
             H = math.ceil(H / self.window_size[0]) * self.window_size[0]
             W = math.ceil(W / self.window_size[1]) * self.window_size[1]

timm/models/swin_transformer_v2.py

@@ -110,7 +110,7 @@ def __init__(
         self.qkv = nn.Linear(dim, dim * 3, bias=False)
         if qkv_bias:
             self.q_bias = nn.Parameter(torch.zeros(dim))
-            self.register_buffer('k_bias', torch.zeros(dim), persistent=False)
+            self.register_buffer('k_bias', torch.zeros(dim, device="cpu"), persistent=False)
             self.v_bias = nn.Parameter(torch.zeros(dim))
         else:
             self.q_bias = None
@@ -125,8 +125,8 @@ def __init__(
 
     def _make_pair_wise_relative_positions(self):
         # get relative_coords_table
-        relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0]).to(torch.float32)
-        relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1]).to(torch.float32)
+        relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], device="cpu").to(torch.float32)
+        relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], device="cpu").to(torch.float32)
         relative_coords_table = torch.stack(ndgrid(relative_coords_h, relative_coords_w))
         relative_coords_table = relative_coords_table.permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, 2
         if self.pretrained_window_size[0] > 0:
@@ -141,8 +141,8 @@ def _make_pair_wise_relative_positions(self):
         self.register_buffer("relative_coords_table", relative_coords_table, persistent=False)
 
         # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(self.window_size[0])
-        coords_w = torch.arange(self.window_size[1])
+        coords_h = torch.arange(self.window_size[0], device="cpu")
+        coords_w = torch.arange(self.window_size[1], device="cpu")
         coords = torch.stack(ndgrid(coords_h, coords_w))  # 2, Wh, Ww
         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
@@ -293,7 +293,7 @@ def get_attn_mask(self, x: Optional[torch.Tensor] = None) -> Optional[torch.Tens
         if any(self.shift_size):
             # calculate attention mask for SW-MSA
             if x is None:
-                img_mask = torch.zeros((1, *self.input_resolution, 1))  # 1 H W 1
+                img_mask = torch.zeros((1, *self.input_resolution, 1), device="cpu")  # 1 H W 1
             else:
                 img_mask = torch.zeros((1, x.shape[1], x.shape[2], 1), dtype=x.dtype, device=x.device)  # 1 H W 1
             cnt = 0

timm/models/swin_transformer_v2_cr.py

@@ -141,7 +141,8 @@ def __init__(
 
     def _make_pair_wise_relative_positions(self) -> None:
         """Method initializes the pair-wise relative positions to compute the positional biases."""
-        device = self.logit_scale.device
+        # device = self.logit_scale.device
+        device = "cpu"
         coordinates = torch.stack(ndgrid(
             torch.arange(self.window_size[0], device=device),
             torch.arange(self.window_size[1], device=device)
@@ -314,7 +315,7 @@ def get_attn_mask(self, x: Optional[torch.Tensor] = None) -> Optional[torch.Tens
         if any(self.shift_size):
             # calculate attention mask for SW-MSA
             if x is None:
-                img_mask = torch.zeros((1, *self.feat_size, 1))  # 1 H W 1
+                img_mask = torch.zeros((1, *self.feat_size, 1), device="cpu")  # 1 H W 1
             else:
                 img_mask = torch.zeros((1, x.shape[1], x.shape[2], 1), dtype=x.dtype, device=x.device)  # 1 H W 1
             cnt = 0