From bb59ffe68f3c1b7cfbe384a12cea626efb64a9cf Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 5 Apr 2020 15:22:32 -0700 Subject: [PATCH] model forward() zip() removal --- models.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/models.py b/models.py index f84da507..557fe187 100755 --- a/models.py +++ b/models.py @@ -230,25 +230,25 @@ class Darknet(nn.Module): img_size = x.shape[-2:] yolo_out, out = [], [] if verbose: - str = '' print('0', x.shape) + str = '' - for i, (mdef, module) in enumerate(zip(self.module_defs, self.module_list)): - mtype = mdef['type'] - if mtype in ['shortcut', 'route']: # sum, concat + for i, module in enumerate(self.module_list): + name = module.__class__.__name__ + if name in ['WeightedFeatureFusion', 'FeatureConcat']: # sum, concat if verbose: l = [i - 1] + module.layers # layers s = [list(x.shape)] + [list(out[i].shape) for i in module.layers] # shapes str = ' >> ' + ' + '.join(['layer %g %s' % x for x in zip(l, s)]) x = module(x, out) # WeightedFeatureFusion(), FeatureConcat() - elif mtype == 'yolo': + elif name == 'YOLOLayer': yolo_out.append(module(x, img_size, out)) else: # run module directly, i.e. mtype = 'convolutional', 'upsample', 'maxpool', 'batchnorm2d' etc. x = module(x) out.append(x if self.routs[i] else []) if verbose: - print('%g/%g %s -' % (i, len(self.module_list), mtype), list(x.shape), str) + print('%g/%g %s -' % (i, len(self.module_list), name), list(x.shape), str) str = '' if self.training: # train