keras设置为TF格式时加载TH格式的权重

Loading weights in TH format when keras is set to TF format

我将 Keras 的 image_dim_ordering 属性 设置为 'tf',所以我将我的模型定义为:

model = Sequential()
model.add(ZeroPadding2D((1, 1), input_shape=(224, 224, 3)))
model.add(Convolution2D(64, 3, 3, activation='relu'))

但是当我调用 load_weights 方法时,它崩溃了,因为我的模型是使用 "th" 格式保存的:

Exception: Layer weight shape (3, 3, 3, 64) not compatible with provided weight shape (64, 3, 3, 3)

如何加载这些权重并自动转置它们以修复 Tensorflow 的格式?

我就此问题询问了 Francois Chollet(他没有 SO 帐户),他友好地转达了以下回复:


"th" 格式意味着卷积核的形状为 (depth, input_depth, rows, cols)

"tf" 格式意味着卷积核的形状为 (rows, cols, input_depth, depth)

因此,您可以通过 np.transpose(x, (2, 3, 1, 0)) 将前者转换为后者,其中 x 是卷积核的值。

下面是一些进行转换的代码:

from keras import backend as K

K.set_image_dim_ordering('th')

# build model in TH mode, as th_model
th_model = ...
# load weights that were saved in TH mode into th_model
th_model.load_weights(...)

K.set_image_dim_ordering('tf')

# build model in TF mode, as tf_model
tf_model = ...

# transfer weights from th_model to tf_model
for th_layer, tf_layer in zip(th_model.layers, tf_model.layers):
   if th_layer.__class__.__name__ == 'Convolution2D':
      kernel, bias = layer.get_weights()
      kernel = np.transpose(kernel, (2, 3, 1, 0))
      tf_layer.set_weights([kernel, bias])
  else:
      tf_layer.set_weights(tf_layer.get_weights())

如果模型包含 Convolution2D 层下游的 Dense 层,则第一个 Dense 层的权重矩阵也需要打乱。

您可以 Use This Script 自动将 theano/tensorflow 后端训练模型权重直接转换为后端/模糊排序的其他 3 种可能组合。