def save_weights(self, outfile, cutoff=0):
if cutoff <= 0:
cutoff = len(self.blocks) - 1
fp = open(outfile, 'wb')
self.header[3] = self.seen
header = self.header
header.numpy().tofile(fp)
ind = -1
for blockId in range(1, cutoff + 1):
ind = ind + 1
block = self.blocks[blockId]
if block['type'] == 'convolutional':
model = self.models[ind]
batch_normalize = int(block['batch_normalize'])
if batch_normalize:
cfg.save_conv_bn(fp, model[0], model[1])
else:
cfg.save_conv(fp, model[0])
elif block['type'] == 'connected':
model = self.models[ind]
if block['activation'] != 'linear':
cfg.save_fc(fc, model)
else:
cfg.save_fc(fc, model[0])
elif block['type'] == 'maxpool':
pass
elif block['type'] == 'reorg':
pass
elif block['type'] == 'upsample':
pass
elif block['type'] == 'route':
pass
elif block['type'] == 'shortcut':
pass
elif block['type'] == 'region':
pass
elif block['type'] == 'yolo':
pass
elif block['type'] == 'avgpool':
pass
elif block['type'] == 'softmax':
pass
elif block['type'] == 'cost':
pass
else:
print('unknown type %s' % (block['type']))
fp.close()
def save_resnet_weights(model, filename):
fp = open(filename, 'wb')
header = torch.IntTensor([0, 0, 0, 0])
header.numpy().tofile(fp)
save_conv_bn(fp, model.conv1, model.bn1)
for i in range(len(model.layer1._modules)):
save_bottlenet_weights(model.layer1[i], fp)
for i in range(len(model.layer2._modules)):
save_bottlenet_weights(model.layer2[i], fp)
for i in range(len(model.layer3._modules)):
save_bottlenet_weights(model.layer3[i], fp)
for i in range(len(model.layer4._modules)):
save_bottlenet_weights(model.layer4[i], fp)
save_fc(fp, model.fc)
fp.close()
def save_bottlenet_weights(model, fp):
save_conv_bn(fp, model.conv1, model.bn1)
save_conv_bn(fp, model.conv2, model.bn2)
save_conv_bn(fp, model.conv3, model.bn3)
if model.downsample:
save_conv_bn(fp, model.downsample[0], model.downsample[1])