updates
This commit is contained in:
parent
f4a9e5cd58
commit
a9cbc28214
45
models.py
45
models.py
|
@ -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):
|
||||||
|
|
Loading…
Reference in New Issue