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])