Keras预训练模型转Pytorch

Weights_Keras_2_Pytorch

最近想在Pytorch项目里使用一下谷歌的NIMA，但是发现没有预训练好的pytorch权重，于是整理了一下将Keras预训练权重转为Pytorch的代码，目前是支持Keras的Conv2D, Dense, DepthwiseConv2D, BatchNormlization的转换。需要注意的是在Pytorch模型内需要给每一层命名为与Keras每一层相同的名字，才能对应转换。

代码地址：

https://github.com/AgCl-LHY/Weights_Keras_2_Pytorch

核心代码：

def keras_to_pyt(km, pm):
    weight_dict = dict()
    for layer in km.layers:
        if type(layer) is keras.layers.convolutional.Conv2D:
            if (len(layer.get_weights()) >= 1):
                weight_dict[layer.get_config()['name'] + '.weight'] = np.transpose(layer.get_weights()[0], (3, 2, 0, 1))
            if (len(layer.get_weights()) >= 2):
                weight_dict[layer.get_config()['name'] + '.bias'] = layer.get_weights()[1]
        elif type(layer) is keras.layers.Dense:
            if (len(layer.get_weights()) >= 1):
                weight_dict[layer.get_config()['name'] + '.weight'] = np.transpose(layer.get_weights()[0], (1, 0))
            if (len(layer.get_weights()) >= 2):
                weight_dict[layer.get_config()['name'] + '.bias'] = layer.get_weights()[1]
        elif type(layer) is keras.layers.DepthwiseConv2D:
            if (len(layer.get_weights()) >= 1):
                weight_dict[layer.get_config()['name'] + '.weight'] = np.transpose(layer.get_weights()[0], (2, 3, 0, 1))
            if (len(layer.get_weights()) >= 2):
                weight_dict[layer.get_config()['name'] + '.bias'] = layer.get_weights()[1]
        elif type(layer) is keras.layers.BatchNormalization:
            if (len(layer.get_weights()) >= 1):
                weight_dict[layer.get_config()['name'] + '.weight'] = layer.get_weights()[0]
            if (len(layer.get_weights()) >= 2):
                weight_dict[layer.get_config()['name'] + '.bias'] = layer.get_weights()[1]
            if (len(layer.get_weights()) >= 3):
                weight_dict[layer.get_config()['name'] + '.running_mean'] = layer.get_weights()[2]
            if (len(layer.get_weights()) >= 4):
                weight_dict[layer.get_config()['name'] + '.running_var'] = layer.get_weights()[3]
        elif type(layer) is keras.layers.ReLU:
            pass
        elif type(layer) is keras.layers.Dropout:
            pass

文件介绍：

weights_keras2pytorch.py 是Keras预训练权重转为Pytorch的代码；

model_keras_NIMA.py 是谷歌NR-IQA NIMA的Keras版模型代码；

model_pytorch_NIMA.py 是NIMA的Pytorch版模型代码；

mobilenet_weights.h5 是用mobilenet实现NIMA的预训练权重；

NIMA_pytorch_model.pth 是用转换代码转换出来的权重；

Requirements：

h5py==3.1.0

Keras

keras2.6.0
Keras-Preprocessing1.1.2

Tensorflow

tensorboard2.7.0
tensorboard-data-server0.6.1
tensorboard-plugin-wit1.8.0
tensorflow-estimator2.7.0
tensorflow-gpu==2.6.0

Pytorch

torch1.10.0
torchaudio0.10.0
torchvision==0.11.1

上述只是我用的环境，并不是一定需要这么高的版本。