与我常在

未经审视的人生不值一过

0%

Keras CNN模型转换为PyTorch模型

发表于 更新于 分类于 阅读次数： 评论数：
本文字数： 2.9k 阅读时长 ≈ 3 分钟

最近想用两种方法跑同一个CNN模型，但一个用Keras写的，一个用PyTorch写的，需要把原先的Keras模型转换为等价的PyTorch模型。目前没有找到合适的自动转换工具，因此只能手动在PyTorch中创建一个一模一样的模型，然后再将权重复制进去：

模型结构复制

以Keras里的LeNet4模型为例：

1
2
3
4
5
6
7
8
9
10
11
12
# block1
x = Convolution2D(6, kernel_size, activation='relu', padding='same', name='block1_conv1')(input_tensor)
x = MaxPooling2D(pool_size=(2, 2), name='block1_pool1')(x)

# block2
x = Convolution2D(16, kernel_size, activation='relu', padding='same', name='block2_conv1')(x)
x = MaxPooling2D(pool_size=(2, 2), name='block2_pool1')(x)

x = Flatten(name='flatten')(x)
x = Dense(84, activation='relu', name='fc1')(x)
x = Dense(nb_classes, name='before_softmax')(x)
x = Activation('softmax', name='predictions')(x)

model.summary()查看参数情况：

keras

在PyTorch中复现相同的结构（忽略最后的softmax层）：

1
2
3
4
5
6
7
8
9
10
11
self.conv2 = nn.Sequential(
    nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, padding=2),
    nn.ReLU(),
    nn.MaxPool2d(2)
)

self.fc1 = nn.Sequential(
    nn.Linear(16 * 7 * 7, 84),
    nn.ReLU()
)
self.before_softmax = nn.Linear(84, 10)

pytorch

权重参数转换

用model.get_weights可以获得Keras模型的权重weights，它是一个list，list中每一项是一个numpy矩阵，依次表示每层的weight、bias。由于各层的bias只有一维，因此两个模型对应位置的bias可以直接互相赋值,如：

1
self.conv1[0].bias.data = from_numpy(weights[1])

但weight是多维矩阵，由于keras和PyTorch的通道维度存储方式不同，需要进行转换，方法如下：

卷积层权重

在Keras的卷积层中kernel weights存储格式是[kernel_height, kernel_width, kernel_channel, kernel_number]，在Pytorch的卷积层中，kernel weights存储格式是[kernel_number, kernel_channel, kernel_height, kernel_width]，因此将Keras模型权重复制到Pytorch模型之前需要进行转换，使用numpy.transpose方法可以方便地对矩阵维度进行位置调换：

1
self.conv1[0].weight.data = from_numpy(np.transpose(weights[0], (3, 2, 0, 1)))

flatten层/第一个全连接层

在将卷积层的输出拉平输入到全连接层时，注意由于Keras和PyTorch通道顺序不同，需要对第一个全连接层的权重进行特殊处理，比如上面的例子中Keras模型第二个池化层的输出形状为（7,7,16），而PyTorch模型的输出形状为（16,7,7），如果不经处理flatten后神经元的顺序将不一样，因此需要先将第一个全连接层的第一个维度复原为Keras中的形状，然后进行转置操作，最后再重新flatten并赋值给PyTorch模型的权重：

1
2
3
w4 = weights[4].reshape((7,7,16,84))    # 复原形状(7,7,16)
w4 = np.transpose(w4, (2,0,1,3)).reshape(784, 84)    # 转置为(16,7,7)
self.fc1[0].weight.data = from_numpy(np.transpose(w4))

全连接层

全连接层权重矩阵的维度为2，直接将矩阵转置即可：

1
self.before_softmax.weight.data = from_numpy(np.transpose(weights[6]))

完整权重转换代码：

1
2
3
4
5
6
7
8
9
10
self.conv1[0].weight.data = from_numpy(np.transpose(weights[0], (3, 2, 0, 1)))
self.conv1[0].bias.data = from_numpy(weights[1])
self.conv2[0].weight.data = from_numpy(np.transpose(weights[2], (3, 2, 0, 1)))
self.conv2[0].bias.data = from_numpy(weights[3])
w4 = weights[4].reshape((7,7,16,84))
w4 = np.transpose(w4, (2,0,1,3)).reshape(784, 84)
self.fc1[0].weight.data = from_numpy(np.transpose(w4))
self.fc1[0].bias.data = from_numpy(weights[5])
self.before_softmax.weight.data = from_numpy(np.transpose(weights[6]))
self.before_softmax.bias.data = from_numpy(weights[7])

参考：

[1] HOW TO TRANSFER A SIMPLE KERAS MODEL TO PYTORCH – THE HARD WAY https://gereshes.com/2019/06/24/how-to-transfer-a-simple-keras-model-to-pytorch-the-hard-way/

[2] Pytorch与Tensorflow权重互转 https://blog.csdn.net/qq_37541097/article/details/113862998