zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Clip TTA Augmented Tails (#5028) 

* Clip TTA Augmented Tails

Experimental TTA update.

* Update yolo.py

* Update yolo.py

* Update yolo.py

* Update yolo.py
This commit is contained in:
Glenn Jocher 2021-10-04 15:48:55 -07:00 committed by GitHub
parent 1922ddeac0
commit d133968b1e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
models/yolo.py
View File

@ -134,6 +134,7 @@ class Model(nn.Module):
# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
yi = self._descale_pred(yi, fi, si, img_size) yi = self._descale_pred(yi, fi, si, img_size)
y.append(yi) y.append(yi)
y = self._clip_augmented(y) # clip augmented tails
return torch.cat(y, 1), None # augmented inference, train return torch.cat(y, 1), None # augmented inference, train
def _forward_once(self, x, profile=False, visualize=False): def _forward_once(self, x, profile=False, visualize=False):
@ -166,6 +167,17 @@ class Model(nn.Module):
p = torch.cat((x, y, wh, p[..., 4:]), -1) p = torch.cat((x, y, wh, p[..., 4:]), -1)
return p return p
def _clip_augmented(self, y):
# Clip YOLOv5 augmented inference tails
nl = self.model[-1].nl # number of detection layers (P3-P5)
g = sum(4 ** x for x in range(nl)) # grid points
e = 1 # exclude layer count
i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e)) # indices
y[0] = y[0][:, :-i] # large
i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # indices
y[-1] = y[-1][:, i:] # small
return y
def _profile_one_layer(self, m, x, dt): def _profile_one_layer(self, m, x, dt):
c = isinstance(m, Detect) # is final layer, copy input as inplace fix c = isinstance(m, Detect) # is final layer, copy input as inplace fix
o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs