glenn-jocher
commited on
Commit
•
302b00b
1
Parent(s):
5f941a8
Update `_make_grid()` (#7346)
Browse files- models/yolo.py +6 -4
models/yolo.py
CHANGED
@@ -77,13 +77,15 @@ class Detect(nn.Module):
|
|
77 |
|
78 |
def _make_grid(self, nx=20, ny=20, i=0):
|
79 |
d = self.anchors[i].device
|
|
|
80 |
shape = 1, self.na, ny, nx, 2 # grid shape
|
|
|
81 |
if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
|
82 |
-
yv, xv = torch.meshgrid(
|
83 |
else:
|
84 |
-
yv, xv = torch.meshgrid(
|
85 |
-
grid = torch.stack((xv, yv), 2).expand(shape)
|
86 |
-
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
|
87 |
return grid, anchor_grid
|
88 |
|
89 |
|
|
|
77 |
|
78 |
def _make_grid(self, nx=20, ny=20, i=0):
|
79 |
d = self.anchors[i].device
|
80 |
+
t = self.anchors[i].dtype
|
81 |
shape = 1, self.na, ny, nx, 2 # grid shape
|
82 |
+
y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
|
83 |
if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
|
84 |
+
yv, xv = torch.meshgrid(y, x, indexing='ij')
|
85 |
else:
|
86 |
+
yv, xv = torch.meshgrid(y, x)
|
87 |
+
grid = torch.stack((xv, yv), 2).expand(shape) - 0.5 # add grid offset, i.e. y = 2.0 * x - 0.5
|
88 |
+
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
|
89 |
return grid, anchor_grid
|
90 |
|
91 |
|