glenn-jocher commited on
Commit
4f839b7
1 Parent(s): ffcbd8c

Refactor out-of-place `Detect()` for reduced ops (#7257)

Browse files
Files changed (1) hide show
  1. models/yolo.py +2 -2
models/yolo.py CHANGED
@@ -66,9 +66,9 @@ class Detect(nn.Module):
66
  y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
67
  y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
68
  else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
69
- xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
70
  wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
71
- y = torch.cat((xy, wh, y[..., 4:]), -1)
72
  z.append(y.view(bs, -1, self.no))
73
 
74
  return x if self.training else (torch.cat(z, 1), x)
 
66
  y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
67
  y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
68
  else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
69
+ xy = (y[..., 0:2] * 2 + (self.grid[i] - 0.5)) * self.stride[i] # xy
70
  wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
71
+ y = torch.cat((xy, wh, y[..., 4:]), 4)
72
  z.append(y.view(bs, -1, self.no))
73
 
74
  return x if self.training else (torch.cat(z, 1), x)