This commit is contained in:
Glenn Jocher 2020-02-19 14:57:58 -08:00
parent f4a9e5cd58
commit a9cbc28214
1 changed files with 20 additions and 25 deletions

View File

@ -264,9 +264,9 @@ class Darknet(nn.Module):
if mtype in ['convolutional', 'upsample', 'maxpool']: if mtype in ['convolutional', 'upsample', 'maxpool']:
x = module(x) x = module(x)
elif mtype == 'shortcut': # sum elif mtype == 'shortcut': # sum
x = module(x, layer_outputs) # weightedFeatureFusion()
if verbose: if verbose:
print('shortcut/add %s' % ([layer_outputs[i].shape for i in module.layers])) print('shortcut/add %s + %s' % (x.shape, [layer_outputs[i].shape for i in module.layers]))
x = module(x, layer_outputs) # weightedFeatureFusion()
elif mtype == 'route': # concat elif mtype == 'route': # concat
layers = [int(x) for x in mdef['layers'].split(',')] layers = [int(x) for x in mdef['layers'].split(',')]
if verbose: if verbose:
@ -354,38 +354,33 @@ def load_darknet_weights(self, weights, cutoff=-1):
ptr = 0 ptr = 0
for i, (mdef, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): for i, (mdef, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])):
if mdef['type'] == 'convolutional': if mdef['type'] == 'convolutional':
conv_layer = module[0] conv = module[0]
if mdef['batch_normalize']: if mdef['batch_normalize']:
# Load BN bias, weights, running mean and running variance # Load BN bias, weights, running mean and running variance
bn_layer = module[1] bn = module[1]
num_b = bn_layer.bias.numel() # Number of biases nb = bn.bias.numel() # number of biases
# Bias # Bias
bn_b = torch.from_numpy(weights[ptr:ptr + num_b]).view_as(bn_layer.bias) bn.bias.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.bias))
bn_layer.bias.data.copy_(bn_b) ptr += nb
ptr += num_b
# Weight # Weight
bn_w = torch.from_numpy(weights[ptr:ptr + num_b]).view_as(bn_layer.weight) bn.weight.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.weight))
bn_layer.weight.data.copy_(bn_w) ptr += nb
ptr += num_b
# Running Mean # Running Mean
bn_rm = torch.from_numpy(weights[ptr:ptr + num_b]).view_as(bn_layer.running_mean) bn.running_mean.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.running_mean))
bn_layer.running_mean.data.copy_(bn_rm) ptr += nb
ptr += num_b
# Running Var # Running Var
bn_rv = torch.from_numpy(weights[ptr:ptr + num_b]).view_as(bn_layer.running_var) bn.running_var.data.copy_(torch.from_numpy(weights[ptr:ptr + nb]).view_as(bn.running_var))
bn_layer.running_var.data.copy_(bn_rv) ptr += nb
ptr += num_b
else: else:
# Load conv. bias # Load conv. bias
num_b = conv_layer.bias.numel() nb = conv.bias.numel()
conv_b = torch.from_numpy(weights[ptr:ptr + num_b]).view_as(conv_layer.bias) conv_b = torch.from_numpy(weights[ptr:ptr + nb]).view_as(conv.bias)
conv_layer.bias.data.copy_(conv_b) conv.bias.data.copy_(conv_b)
ptr += num_b ptr += nb
# Load conv. weights # Load conv. weights
num_w = conv_layer.weight.numel() nw = conv.weight.numel() # number of weights
conv_w = torch.from_numpy(weights[ptr:ptr + num_w]).view_as(conv_layer.weight) conv.weight.data.copy_(torch.from_numpy(weights[ptr:ptr + nw]).view_as(conv.weight))
conv_layer.weight.data.copy_(conv_w) ptr += nw
ptr += num_w
def save_weights(self, path='model.weights', cutoff=-1): def save_weights(self, path='model.weights', cutoff=-1):