From fff45c39a83d444a9d5dda5c9f846416a0274bd6 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 19 Mar 2020 16:41:42 -0700 Subject: [PATCH] cleanup/speedup --- models.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/models.py b/models.py index 6b1dddf3..4a815e7d 100755 --- a/models.py +++ b/models.py @@ -105,7 +105,10 @@ def create_modules(module_defs, img_size): module_list.append(modules) output_filters.append(filters) - return module_list, routs + routs_binary = [False] * (i + 1) + for i in routs: + routs_binary[i] = True + return module_list, routs_binary class weightedFeatureFusion(nn.Module): # weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 @@ -292,7 +295,7 @@ class Darknet(nn.Module): # print(''), [print(out[i].shape) for i in layers], print(x.shape) elif mtype == 'yolo': yolo_out.append(module(x, img_size, out)) - out.append(x if i in self.routs else []) + out.append(x if self.routs[i] else []) if verbose: print('%g/%g %s -' % (i, len(self.module_list), mtype), list(x.shape), str) str = '' @@ -342,7 +345,7 @@ def create_grids(self, img_size=416, ng=(13, 13), device='cpu', type=torch.float # build wh gains self.anchor_vec = self.anchors.to(device) / self.stride - self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2).to(device).type(type) + self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2).type(type) self.ng = torch.Tensor(ng).to(device) self.nx = nx self.ny = ny