`InaccessibleTensorError` с лямбда-функцией keras с использованием обнимающего лица BertModel

Question

Я загружаю предварительно обученную модель Берта, чтобы выполнить классификацию текста. По умолчанию transformers.TFBertModel возвращает hidden_states только последнего слоя, где я могу извлечь первый токен и использовать его для классификации текста.

Однако я хочу извлечь векторные представления первого токена не из последнего ОДНОГО слоя, а из последних ЧЕТЫРЕХ слоев. Поэтому я пишу лямбда-слой, чтобы объединить четыре вектора первого токена (каждый из одного из четырех последних слоев) в один более длинный вектор. Вот где у меня проблема. Вот код для импорта модулей и определения функции.

import tensorflow as tf
from keras.layers import Input, Lambda
from transformers import BertConfig, TFBertModel

config = BertConfig(output_hidden_states=True)
bert_layer = TFBertModel.from_pretrained('bert-base-uncased', config=config)

@tf.function
def concat_layers(layers, nlayer=4):
    # input layers shape: a tuple of 13 layers, where each is of [# batches, # tokens, # features]
    # output ans shape: [# batches, # tokens, # features*nlayer]

    # stack the last 12 layers into dim of [# layers, # batches, # tokens, # features]
    layers = tf.stack(layers[1:], axis=0)

    # permute dimensions into [# batches, # tokens, # layers, # features]
    layers = tf.transpose(layers, (1, 2, 0, 3))

    # for each token, loop to slice last n-layers and reshape them into a vector
    ans = []
    for batch in layers:
        tans = []
        for token in batch:
            vec = tf.squeeze(tf.reshape(token[-nlayer:], (1, -1)))
            tans.append(vec)
        tans = tf.stack(tans, axis=0)
        ans.append(tans)
    ans = tf.stack(ans, axis=0)

    return ans

Хорошо работает с

input1 = Input(shape=(500,), dtype=tf.int32, name='inputs')
input2 = Input(shape=(500,), dtype=tf.int32, name='masks')
input3 = Input(shape=(500,), dtype=tf.int32, name='segments')
last_layer, first_token, layers = bert_layer([input1, input2, input3])

Но когда я делаю

layers = Lambda(function=concat_layers)(layers)

Он сообщает об ошибке :

InaccessibleTensorError: in converted code:

    <ipython-input-125-09697873bb2c>:17 concat_layers  *
        tans = tf.stack(tans, axis=0)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/util/dispatch.py:180 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/array_ops.py:1255 stack
        return gen_array_ops.pack(values, axis=axis, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/gen_array_ops.py:5704 pack
        "Pack", values=values, axis=axis, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/op_def_library.py:742 _apply_op_helper
        attrs=attr_protos, op_def=op_def)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py:591 _create_op_internal
        inp = self.capture(inp)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py:641 capture
        % (tensor, tensor.graph, self))

    InaccessibleTensorError: The tensor 'Tensor("Squeeze:0", shape=(3072,), dtype=float32)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=while_body_687352, id=139971488261848); accessed from: FuncGraph(name=while_body_687327, id=139971488296744).

UPDATE NOTE

Я сузил проблему. Но не знаю, как это решить. Проблема связана с циклом. Именно это конкретное утверждение и вызывает проблему.

vec = tf.squeeze(tf.reshape(token[-nlayer:], (1, -1)))

Если я заменим решение с зацикливанием тока на векторизованное решение, как показано ниже, то оно будет работать.

ans = tf.reshape(tmat, (bs, nt, -1, nlayer, 768))
ans = ans[:, :, -1, :, :]
ans = tf.squeeze(tf.reshape(ans, (bs, nt, 1, -1)), 2)

, где bs - размер партии, nt - количество токенов (ограничено 512 для Берта).

Возможно ли решение для зацикливания, которое является более гибким, чем мое векторизованное решение, которое может только 1, 2, 3, 4, 6, 12 слоев, но не 5, 7, 8, 9, 10, 11 (т.е. 12% nlayer == 0).