From e6e53c4e232c6e89257f5095553d4fadaf8690dc Mon Sep 17 00:00:00 2001 From: magehrig Date: Thu, 16 Sep 2021 16:34:37 +0200 Subject: [PATCH 1/2] create tensors on device --- core/corr.py | 6 +++--- core/raft.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/core/corr.py b/core/corr.py index 645ca55e..cffcbc82 100644 --- a/core/corr.py +++ b/core/corr.py @@ -34,9 +34,9 @@ def __call__(self, coords): out_pyramid = [] for i in range(self.num_levels): corr = self.corr_pyramid[i] - dx = torch.linspace(-r, r, 2*r+1) - dy = torch.linspace(-r, r, 2*r+1) - delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device) + dx = torch.linspace(-r, r, 2*r+1, device=coords.device) + dy = torch.linspace(-r, r, 2*r+1, device=coords.device) + delta = torch.stack(torch.meshgrid(dy, dx), axis=-1) centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2) diff --git a/core/raft.py b/core/raft.py index e0519ed6..652b81a3 100644 --- a/core/raft.py +++ b/core/raft.py @@ -63,8 +63,8 @@ def freeze_bn(self): def initialize_flow(self, img): """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0""" N, C, H, W = img.shape - coords0 = coords_grid(N, H//8, W//8).to(img.device) - coords1 = coords_grid(N, H//8, W//8).to(img.device) + coords0 = coords_grid(N, H//8, W//8, device=img.device) + coords1 = coords_grid(N, H//8, W//8, device=img.device) # optical flow computed as difference: flow = coords1 - coords0 return coords0, coords1 From 0d123fda29f57f787b21f2cea0e7bf39309e03b6 Mon Sep 17 00:00:00 2001 From: magehrig Date: Tue, 28 Sep 2021 10:55:56 +0200 Subject: [PATCH 2/2] coords_grid on device --- core/utils/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/utils/utils.py b/core/utils/utils.py index 5f32d281..741ccfe4 100644 --- a/core/utils/utils.py +++ b/core/utils/utils.py @@ -71,8 +71,8 @@ def bilinear_sampler(img, coords, mode='bilinear', mask=False): return img -def coords_grid(batch, ht, wd): - coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) +def coords_grid(batch, ht, wd, device): + coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device)) coords = torch.stack(coords[::-1], dim=0).float() return coords[None].repeat(batch, 1, 1, 1)