Browse Source

cleanup

5.0
Glenn Jocher 4 years ago
parent
commit
8f17a96ff0
1 changed files with 0 additions and 28 deletions
  1. +0
    -28
      models/yolo.py

+ 0
- 28
models/yolo.py View File

@@ -38,34 +38,6 @@ class Detect(nn.Module):

return x if self.training else (torch.cat(z, 1), x)

def forward_(self, x):
if hasattr(self, 'nx'):
z = [] # inference output
for (y, gi, agi, si, nyi, nxi) in zip(x, self.grid, self.ag, self.stride, self.ny, self.nx):
m = self.na * nxi * nyi
y = y.view(1, self.na, self.no, nyi, nxi).permute(0, 1, 3, 4, 2).contiguous().view(m, self.no).sigmoid()

xy = (y[..., 0:2] * 2. - 0.5 + gi) * si # xy
wh = (y[..., 2:4] * 2) ** 2 * agi # wh
p_cls = y[:, 4:5] if self.nc == 1 else y[:, 5:self.no] * y[:, 4:5] # conf
z.append([p_cls, xy, wh])

z = [torch.cat(x, 0) for x in zip(*z)]
return z[0], torch.cat(z[1:3], 1) # scores, boxes: 3780x80, 3780x4

else: # dry run
self.nx = [0] * self.nl
self.ny = [0] * self.nl
self.ag = [0] * self.nl
for i in range(self.nl):
bs, _, ny, nx = x[i].shape
m = self.na * nx * ny
self.grid[i] = self._make_grid(nx, ny).repeat(1, self.na, 1, 1, 1).view(m, 2) / torch.tensor([[nx, ny]])
self.ag[i] = self.anchor_grid[i].repeat(1, 1, nx, ny, 1).view(m, 2)
self.nx[i] = nx
self.ny[i] = ny
return None

@staticmethod
def _make_grid(nx=20, ny=20):
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])

Loading…
Cancel
Save