From 5c470d235e1cd8ae672035b4bddc56fdd745aaa7 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 5 Jun 2020 20:36:29 -0700 Subject: [PATCH] multi-gpu test bug fix #15 --- models/yolo.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/models/yolo.py b/models/yolo.py index 871312e..38fa3c0 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -20,7 +20,7 @@ class Detect(nn.Module): self.export = False # onnx export def forward(self, x): - x = x.copy() # for profiling + x = x.copy() # for profiling z = [] # inference output self.training |= self.export for i in range(self.nl): @@ -67,6 +67,25 @@ class Model(nn.Module): print('') def forward(self, x, augment=False, profile=False): + if augment: + img_size = x.shape[-2:] # height, width + s = [0.83, 0.67] # scales + y = [] + for i, xi in enumerate((x, + torch_utils.scale_img(x.flip(3), s[0], same_shape=False), # flip-lr and scale + torch_utils.scale_img(x, s[1], same_shape=False), # scale + )): + # cv2.imwrite('img%g.jpg' % i, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) + y.append(self.forward_once(xi)[0]) + + y[1][..., :4] /= s[0] # scale + y[1][..., 0] = img_size[1] - y[1][..., 0] # flip lr + y[2][..., :4] /= s[1] # scale + return torch.cat(y, 1), None # augmented inference, train + else: + return self.forward_once(x, profile) # single-scale inference, train + + def forward_once(self, x, profile=False): y, dt = [], [] # outputs for m in self.model: if m.f != -1: # if not from previous layer