updates
This commit is contained in:
parent
ee4f7a324d
commit
07c8a03aa0
10
models.py
10
models.py
|
@ -224,6 +224,9 @@ class Darknet(nn.Module):
|
||||||
img_size = x.shape[-2:]
|
img_size = x.shape[-2:]
|
||||||
layer_outputs = []
|
layer_outputs = []
|
||||||
output = []
|
output = []
|
||||||
|
verbose = False
|
||||||
|
if verbose:
|
||||||
|
print('0', x.shape)
|
||||||
|
|
||||||
for i, (mdef, module) in enumerate(zip(self.module_defs, self.module_list)):
|
for i, (mdef, module) in enumerate(zip(self.module_defs, self.module_list)):
|
||||||
mtype = mdef['type']
|
mtype = mdef['type']
|
||||||
|
@ -241,10 +244,15 @@ class Darknet(nn.Module):
|
||||||
x = torch.cat([layer_outputs[i] for i in layers], 1)
|
x = torch.cat([layer_outputs[i] for i in layers], 1)
|
||||||
# print(''), [print(layer_outputs[i].shape) for i in layers], print(x.shape)
|
# print(''), [print(layer_outputs[i].shape) for i in layers], print(x.shape)
|
||||||
elif mtype == 'shortcut':
|
elif mtype == 'shortcut':
|
||||||
x = x + layer_outputs[int(mdef['from'])]
|
j = int(mdef['from'])
|
||||||
|
if verbose:
|
||||||
|
print('shortcut adding layer %g-%s to %g-%s' % (j, layer_outputs[j].shape, i - 1, x.shape))
|
||||||
|
x = x + layer_outputs[j]
|
||||||
elif mtype == 'yolo':
|
elif mtype == 'yolo':
|
||||||
output.append(module(x, img_size))
|
output.append(module(x, img_size))
|
||||||
layer_outputs.append(x if i in self.routs else [])
|
layer_outputs.append(x if i in self.routs else [])
|
||||||
|
if verbose:
|
||||||
|
print(i, x.shape)
|
||||||
|
|
||||||
if self.training:
|
if self.training:
|
||||||
return output
|
return output
|
||||||
|
|
Loading…
Reference in New Issue