model forward() zip() removal
This commit is contained in:
parent
a657345b45
commit
bb59ffe68f
12
models.py
12
models.py
|
@ -230,25 +230,25 @@ class Darknet(nn.Module):
|
||||||
img_size = x.shape[-2:]
|
img_size = x.shape[-2:]
|
||||||
yolo_out, out = [], []
|
yolo_out, out = [], []
|
||||||
if verbose:
|
if verbose:
|
||||||
str = ''
|
|
||||||
print('0', x.shape)
|
print('0', x.shape)
|
||||||
|
str = ''
|
||||||
|
|
||||||
for i, (mdef, module) in enumerate(zip(self.module_defs, self.module_list)):
|
for i, module in enumerate(self.module_list):
|
||||||
mtype = mdef['type']
|
name = module.__class__.__name__
|
||||||
if mtype in ['shortcut', 'route']: # sum, concat
|
if name in ['WeightedFeatureFusion', 'FeatureConcat']: # sum, concat
|
||||||
if verbose:
|
if verbose:
|
||||||
l = [i - 1] + module.layers # layers
|
l = [i - 1] + module.layers # layers
|
||||||
s = [list(x.shape)] + [list(out[i].shape) for i in module.layers] # shapes
|
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)])
|
str = ' >> ' + ' + '.join(['layer %g %s' % x for x in zip(l, s)])
|
||||||
x = module(x, out) # WeightedFeatureFusion(), FeatureConcat()
|
x = module(x, out) # WeightedFeatureFusion(), FeatureConcat()
|
||||||
elif mtype == 'yolo':
|
elif name == 'YOLOLayer':
|
||||||
yolo_out.append(module(x, img_size, out))
|
yolo_out.append(module(x, img_size, out))
|
||||||
else: # run module directly, i.e. mtype = 'convolutional', 'upsample', 'maxpool', 'batchnorm2d' etc.
|
else: # run module directly, i.e. mtype = 'convolutional', 'upsample', 'maxpool', 'batchnorm2d' etc.
|
||||||
x = module(x)
|
x = module(x)
|
||||||
|
|
||||||
out.append(x if self.routs[i] else [])
|
out.append(x if self.routs[i] else [])
|
||||||
if verbose:
|
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 = ''
|
str = ''
|
||||||
|
|
||||||
if self.training: # train
|
if self.training: # train
|
||||||
|
|
Loading…
Reference in New Issue