def save_vgg16_weights(model, filename): fp = open(filename, 'wb') header = torch.IntTensor([0, 0, 0, 0]) header.numpy().tofile(fp) for layer in model.features: if type(layer) == torch.nn.Conv2d: print(layer) save_conv(fp, layer) for layer in model.classifier: if type(layer) == torch.nn.Linear: print(layer) save_fc(fp, layer)
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_alexnet_weights(model, filename): fp = open(filename, 'wb') header = torch.IntTensor([0, 0, 0, 0]) header.numpy().tofile(fp) for layer in model.model: print layer if type(layer) == torch.nn.Conv2d: print "1" print(layer) save_conv(fp, layer) #for layer in model.fc: if type(model.fc) == torch.nn.Linear: print(layer) save_fc(fp, model.fc)
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()