TensorFlow Keras CuDNNGRU到GRU转换

Question

我在TensorFlow 1.14中构建了一个经过训练的模型，该模型使用(现在已被废弃的) tf.keras.layers.CuDNNGRU层(在tf.compat.v1中的TensorFlow 2.0中可用)，并且我正在尝试将旧层的权重移植到一个使用tf.keras.layers.GRU构建的新TensorFlow 2.0模型中，以获得一个等效的模型。

这样做的一个动机是能够在CPU上进行推断( tf.compat.v1.keras.layers.CuDNNGRU层只运行在GPU上)。另一个动机是对模型进行未来验证。

问题

如何将经过训练的tf.contrib.v1.keras.layers.CuDNNGRU层转换为等效的tf.keras.layers.GRU层？

Answer 1

下面的tensorflow.python.keras.saving.hdf5_format中的私有助手函数似乎可以完成这个任务。该函数执行更一般的任务，即在CuDNNGRU/GRU和CuDNNLSTM/LSTM格式之间转换权重，因此它不仅适用于我的用例。该函数似乎起源于独立Keras中的此拉请求。

import numpy as np


def _convert_rnn_weights(layer, weights):
  """Converts weights for RNN layers between native and CuDNN format.

  Input kernels for each gate are transposed and converted between Fortran
  and C layout, recurrent kernels are transposed. For LSTM biases are summed/
  split in half, for GRU biases are reshaped.

  Weights can be converted in both directions between `LSTM` and`CuDNNSLTM`
  and between `CuDNNGRU` and `GRU(reset_after=True)`. Default `GRU` is not
  compatible with `CuDNNGRU`.

  For missing biases in `LSTM`/`GRU` (`use_bias=False`) no conversion is made.

  Arguments:
      layer: Target layer instance.
      weights: List of source weights values (input kernels, recurrent
          kernels, [biases]) (Numpy arrays).

  Returns:
      A list of converted weights values (Numpy arrays).

  Raises:
      ValueError: for incompatible GRU layer/weights or incompatible biases
  """


  def transform_kernels(kernels, func, n_gates):
    """Transforms kernel for each gate separately using given function.

    Arguments:
        kernels: Stacked array of kernels for individual gates.
        func: Function applied to kernel of each gate.
        n_gates: Number of gates (4 for LSTM, 3 for GRU).

    Returns:
        Stacked array of transformed kernels.
    """
    return np.hstack([func(k) for k in np.hsplit(kernels, n_gates)])


  def transpose_input(from_cudnn):
    """Makes a function that transforms input kernels from/to CuDNN format.

    It keeps the shape, but changes between the layout (Fortran/C). Eg.:

    ```
    Keras                 CuDNN
    [[0, 1, 2],  <--->  [[0, 2, 4],
     [3, 4, 5]]          [1, 3, 5]]
    ```

    It can be passed to `transform_kernels()`.

    Arguments:
        from_cudnn: `True` if source weights are in CuDNN format, `False`
            if they're in plain Keras format.

    Returns:
        Function that converts input kernel to the other format.
    """
    order = 'F' if from_cudnn else 'C'


    def transform(kernel):
      return kernel.T.reshape(kernel.shape, order=order)


    return transform


  target_class = layer.__class__.__name__


  # convert the weights between CuDNNLSTM and LSTM
  if target_class in ['LSTM', 'CuDNNLSTM'] and len(weights) == 3:
    # determine if we're loading a CuDNNLSTM layer
    # from the number of bias weights:
    # CuDNNLSTM has (units * 8) weights; while LSTM has (units * 4)
    # if there's no bias weight in the file, skip this conversion
    units = weights[1].shape[0]
    bias_shape = weights[2].shape
    n_gates = 4


    if bias_shape == (2 * units * n_gates,):
      source = 'CuDNNLSTM'
    elif bias_shape == (units * n_gates,):
      source = 'LSTM'
    else:
      raise ValueError('Invalid bias shape: ' + str(bias_shape))


    def convert_lstm_weights(weights, from_cudnn=True):
      """Converts the weights between CuDNNLSTM and LSTM.

      Arguments:
        weights: Original weights.
        from_cudnn: Indicates whether original weights are from CuDNN layer.

      Returns:
        Updated weights compatible with LSTM.
      """


      # Transpose (and reshape) input and recurrent kernels
      kernels = transform_kernels(weights[0], transpose_input(from_cudnn),
                                  n_gates)
      recurrent_kernels = transform_kernels(weights[1], lambda k: k.T, n_gates)
      if from_cudnn:
        # merge input and recurrent biases into a single set
        biases = np.sum(np.split(weights[2], 2, axis=0), axis=0)
      else:
        # Split single set of biases evenly to two sets. The way of
        # splitting doesn't matter as long as the two sets sum is kept.
        biases = np.tile(0.5 * weights[2], 2)
      return [kernels, recurrent_kernels, biases]


    if source != target_class:
      weights = convert_lstm_weights(weights, from_cudnn=source == 'CuDNNLSTM')


  # convert the weights between CuDNNGRU and GRU(reset_after=True)
  if target_class in ['GRU', 'CuDNNGRU'] and len(weights) == 3:
    # We can determine the source of the weights from the shape of the bias.
    # If there is no bias we skip the conversion since
    # CuDNNGRU always has biases.


    units = weights[1].shape[0]
    bias_shape = weights[2].shape
    n_gates = 3


    def convert_gru_weights(weights, from_cudnn=True):
      """Converts the weights between CuDNNGRU and GRU.

      Arguments:
        weights: Original weights.
        from_cudnn: Indicates whether original weights are from CuDNN layer.

      Returns:
        Updated weights compatible with GRU.
      """


      kernels = transform_kernels(weights[0], transpose_input(from_cudnn),
                                  n_gates)
      recurrent_kernels = transform_kernels(weights[1], lambda k: k.T, n_gates)
      biases = np.array(weights[2]).reshape((2, -1) if from_cudnn else -1)
      return [kernels, recurrent_kernels, biases]


    if bias_shape == (2 * units * n_gates,):
      source = 'CuDNNGRU'
    elif bias_shape == (2, units * n_gates):
      source = 'GRU(reset_after=True)'
    elif bias_shape == (units * n_gates,):
      source = 'GRU(reset_after=False)'
    else:
      raise ValueError('Invalid bias shape: ' + str(bias_shape))


    if target_class == 'CuDNNGRU':
      target = 'CuDNNGRU'
    elif layer.reset_after:
      target = 'GRU(reset_after=True)'
    else:
      target = 'GRU(reset_after=False)'


    # only convert between different types
    if source != target:
      types = (source, target)
      if 'GRU(reset_after=False)' in types:
        raise ValueError('%s is not compatible with %s' % types)
      if source == 'CuDNNGRU':
        weights = convert_gru_weights(weights, from_cudnn=True)
      elif source == 'GRU(reset_after=True)':
        weights = convert_gru_weights(weights, from_cudnn=False)


  return weights

对于我的用例(将CuDNNGRU权重放入GRU)，使用此函数的解决方案如下：

# cudnn_gru and gru are built CuDNNGRU and GRU layers, respectively
kernel, recurrent_kernel, bias = _convert_rnn_weights(
    layer=gru,
    weights=[
        cudnn_gru.kernel.numpy(),
        cudnn_gru.recurrent_kernel.numpy(),
        cudnn_gru.bias.numpy(),
    ],
)
gru.cell.kernel.assign(kernel)
gru.cell.recurrent_kernel.assign(recurrent_kernel)
gru.cell.bias.assign(bias)

注意，要使用与cuDNN兼容的tf.keras.layers.GRU实现，必须使用使用特定的参数组合 (特别是use_bias=True)。

Answer 2

我知道这个线程有点旧，但是我想补充一下，如何在Keras/TF 2.6中将CuDNNGRUs/CuDNNLSTM转换为GRUs/LSTM(接受的答案对我无效，因为gru.cell的属性似乎已经改变了)。

背景:我想用GPU(比在GPU上训练标准GRU更快)来训练CuDNNGRU，并将它转换成标准GRU进行CPU推理。

这个解决方案来自一个名为GitHub 萨米克尼克的家伙。

1. 创建包含CuDNNGRU gru_cudnn (或CuDNNLSTM)的模型，对其进行训练，并保存其权重： gru_cudnn = CuDNNGRU(n_units)模型=.用gru_cudnn制作模型..。Model.fit(.)model.save_weights('weights_cudnn.h5')
2. 使用标准GRU gru (或LSTM)而不是CuDNNGRU (或CuDNNLSTM)创建具有相同体系结构的模型，并从1加载保存的CuDNN权重： gru = GRU(n_units，reset_after=True，递归激活=‘sigmoid’)模型=.用格鲁做模特..。model.load_weights('weights_cudnn.h5')

我希望这对将来偶然发现这条线的人有帮助。