convert vgg-face model weight from caffe to pytorch

#coding=UTF-8
import torch 
import torch.nn as nn
import math
import numpy as np
from PIL import Image,ImageDraw
import matplotlib.pyplot as plt
import collections
import matplotlib.cm as cm

from torch.autograd import Variable
from torchvision import models
import caffe


def vis_square(data):
    """Take an array of shape (n, height, width) or (n, height, width, 3)
       and visualize each (height, width) thing in a grid of size approx. sqrt(n) by sqrt(n)"""

    # normalize data for display
    data = (data - data.min()) / (data.max() - data.min())

    # force the number of filters to be square
    n = int(np.ceil(np.sqrt(data.shape[0])))
    padding = (((0, n ** 2 - data.shape[0]),
               (0, 1), (0, 1))                 # add some space between filters
               + ((0, 0),) * (data.ndim - 3))  # don't pad the last dimension (if there is one)
    data = np.pad(data, padding, mode='constant', constant_values=1)  # pad with ones (white)

    # tile the filters into an image
    data = data.reshape((n, n) + data.shape[1:]).transpose((0, 2, 1, 3) + tuple(range(4, data.ndim + 1)))
    data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])

    plt.imshow(data); plt.axis('off')

def vis_activation(activations, img_arr):
    assert(len(activations.shape) == 3),'make sure : nchannels * height * width '
    n_nodes = activations.shape[0] # number of nodels

    print('n_nodes: ',n_nodes)
    vis_size = activations.shape[1:] #visualization shape
    print('visual_size: ',vis_size)
    vis = np.zeros((vis_size[0], vis_size[1]), dtype=np.float32)

    #generating saliency image
    for i in range(n_nodes):
        activation = activations[i, :, :]
        weight = 1
        weighted_activation = activation*weight
        vis += weighted_activation
    vis = np.maximum(vis, 0) # relu
    vis_img = Image.fromarray(vis, None)
    vis_img = vis_img.resize((224,224),Image.BICUBIC)
    vis_img = vis_img / np.max(vis_img)


    vis_img = Image.fromarray(np.uint8(cm.jet(vis_img) * 255)) # winter ,jet
    vis_img = vis_img.convert('RGB') # dropping alpha channel


    input_image = Image.fromarray(img_arr)
    input_image = input_image.resize((224,224))
    input_image = input_image.convert('RGB')
    plt.imshow(vis_img)
    #plt.show()
    #print vis_img.size, input_i
    heat_map = Image.blend(input_image, vis_img, 0.5)
    plt.imshow(heat_map)
    plt.axis('off')

class vgg16(nn.Module):
    def __init__(self,num_classes=1000):
        super(vgg16,self).__init__() # call  parents' init method'
        inplace = True
        self.features = nn.Sequential(
            nn.Conv2d(3,64,kernel_size=(3,3),stride=(1,1),padding=(1,1)), # input_channel, output_channel
            nn.ReLU(inplace),
            nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1), ceil_mode=False),
            nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1), ceil_mode=False),
            nn. Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1), ceil_mode=False),
            nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1), ceil_mode=False),
            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(inplace),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1), ceil_mode=False)  
            )
        self.classifier = nn.Sequential(
            nn.Linear(in_features=25088, out_features=4096, bias=True),
            nn.ReLU(inplace),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=4096, bias=True),
            nn.ReLU(inplace),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=num_classes, bias=True)
            )
        self._initialize_weights()
    def forward(self,x):
        x = self.features(x)
        x = x.view(x.size(0),-1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

model = vgg16()
# print items 
#print(model)
params = model.state_dict()
print(type(params))
for key,item in params.items():
    print(key,': ', type(item)) # tensor

filepath = '../my_headpose_using_VGG_version2/pretrained_model/vgg_face_caffe/'
pretrained_weight = filepath + 'VGG_FACE.caffemodel'
deploy_pro = filepath + 'VGG_FACE_deploy.prototxt'
net = caffe.Net(deploy_pro,pretrained_weight,caffe.TEST)
print('load caffemodel over ....')

'''
orderdict = collections.OrderedDict()
orderdict['features.0'] = 'conv1_1'
orderdict['features.2'] = 'conv1_2'
orderdict['features.5'] = 'conv2_1'
orderdict['features.7'] = 'conv2_2'
orderdict['features.10'] = 'conv3_1'
orderdict['features.12'] = 'conv3_2'
orderdict['features.14'] = 'conv3_3'
orderdict['features.17'] = 'conv4_1'
orderdict['features.19'] = 'conv4_2'
orderdict['features.21'] = 'conv4_3'
orderdict['features.24'] = 'conv5_1'
orderdict['features.26'] = 'conv5_2'
orderdict['features.28'] = 'conv5_3'
orderdict['classifier.0'] = 'fc6'
orderdict['classifier.3'] = 'fc7'
orderdict['classifier.6'] = 'fc8'
'''

dict_feature = {0:'conv1_1',
                2:'conv1_2',
                5:'conv2_1',
                7:'conv2_2',
                10:'conv3_1',
                12:'conv3_2',
                14:'conv3_3',
                17:'conv4_1',
                19:'conv4_2',
                21:'conv4_3',
                24:'conv5_1',
                26:'conv5_2',
                28:'conv5_3'}

dict_classifier = {0:'fc6',
                   3:'fc7'}
#                   6:'fc8'}


for i in dict_feature:

    caffe_weight = net.params[dict_feature[i]][0].data
    #print(type(model.features[i].weight))
    caffe_weight = torch.from_numpy(caffe_weight).view_as(model.features[i].weight)

    model.features[i].weight.data.copy_(caffe_weight)
    model.features[i].bias.data.copy_(torch.from_numpy(np.array(net.params[dict_feature[i]][1].data)))

for i in dict_classifier:

    model.classifier[i].weight.data.copy_(torch.from_numpy(np.array(net.params[dict_classifier[i]][0].data)))

    model.classifier[i].bias.data.copy_(torch.from_numpy(np.array(net.params[dict_classifier[i]][1].data)))

print('copy weight from caffemodel to pytorch over....')

########## check #####################
imgSize = [224,224]

imgpath = '../copy_caffemodel_2_pytorch/cat.jpg'
img = Image.open(imgpath)
res_img = img.resize((imgSize[0],imgSize[1]))
img = np.double(res_img)
temp_img = np.uint8(res_img) # for vis
img = img[:,:,(2,1,0)] # rgb 2 bgr
img = np.transpose(img, (2,0,1)) # h * w *c==> c*h*w

print(img.shape)

data_arr = np.zeros(shape=(1,3,imgSize[0],imgSize[1]),dtype=np.float32)
data_arr[0,...] = img
input_data = Variable(torch.from_numpy(data_arr).type(torch.FloatTensor))
feat_result  = []
def get_features_hook(self,input,output):
    # number of input:
    print('len(input): ',len(input))
    # number of output:
    print('len(output): ',len(output))
    print('###################################')
    print(input[0].shape) # torch.Size([1, 3, 224, 224])

    print('###################################')
    print(output[0].shape) # 
    feat_result.append(output.data.cpu().numpy())

handle_feat = model.features[0].register_forward_hook(get_features_hook)  # conv1_1  
handle_heat = model.features[30].register_forward_hook(get_features_hook) # pool5
handle_fc7 = model.classifier[3].register_forward_hook(get_features_hook) # fc7

model.eval() # make dropout invalid in test stage
score  = model(input_data)

handle_feat.remove()
handle_heat.remove()

feat1 = feat_result[0]
feat1_heat = feat_result[1]


vis_square(feat1[0,...])
#plt.show()
plt.savefig('feat_visual_pytorch.png')
####################################### pytorch heatmap ####################################
vis_activation(feat1_heat[0,...],temp_img)
print('show heatmap for pytorch...')
#plt.show()
plt.savefig('heatmap_visual_pytorch.png')
############################################  for caffe ############################################
net.blobs['data'].reshape(1,3,imgSize[0],imgSize[1])
net.blobs['data'].data[...] = data_arr
output = net.forward()

feat2 = net.blobs['conv1_1'].data[0,...] 
vis_square(feat2)
#plt.show()
plt.savefig('feat_visual_caffe.png')

############################# caffe heatmap ################################
feat2_heat = net.blobs['pool5'].data[0,...] 
vis_activation(feat2_heat,temp_img)
#plt.show()
plt.savefig('heatmap_visual_caffe.png')

######################### check fc7 layer#######################
fc7_pytorch = feat_result[2]
fc7_caffe = net.blobs['fc7'].data
print(fc7_pytorch.shape)
print(fc7_caffe.shape)

err = np.max(np.abs(fc7_pytorch - fc7_caffe))
print(err) #6.914139e-06

验证结果图：

下图分别是： feat_visual_pytorch.png 和 heatmap_visual_pytorch.png
对于caffe现实的结果相同，不再附图。

参考文献：
1.https://github.com/marvis/pytorch-caffe