addons

addons copied to clipboard

In layers/crf CRF Layer, the output is [decoded_sequence, potentials, sequence_length, self.chain_kernel], however, when I want to get the first element of the output like the code below :

from tensorflow.keras.layers import Input, Embedding, Bidirectional, GRU, Dense
from tensorflow.keras.models import Model
import tensorflow as tf
from tensorflow_addons.layers.crf import CRF
from tensorflow_addons.text.crf import crf_log_likelihood

def unpack_data(data):
    if len(data) == 2:
        return data[0], data[1], None
    elif len(data) == 3:
        return data
    else:
        raise TypeError("Expected data to be a tuple of size 2 or 3.")


class ModelWithCRFLoss(tf.keras.Model):
    """Wrapper around the base model for custom training logic."""

    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model

    def call(self, inputs):
        return self.base_model(inputs)

    def compute_loss(self, x, y, sample_weight, training=False):
        y_pred = self(x, training=training)
        _, potentials, sequence_length, chain_kernel = y_pred

        # we now add the CRF loss:
        crf_loss = -crf_log_likelihood(potentials, y, sequence_length, chain_kernel)[0]

        if sample_weight is not None:
            crf_loss = crf_loss * sample_weight

        return tf.reduce_mean(crf_loss), sum(self.losses)

    def train_step(self, data):
        x, y, sample_weight = unpack_data(data)

        with tf.GradientTape() as tape:
            crf_loss, internal_losses = self.compute_loss(
                x, y, sample_weight, training=True
            )
            total_loss = crf_loss + internal_losses

        gradients = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        return {"crf_loss": crf_loss, "internal_losses": internal_losses}

    def test_step(self, data):
        x, y, sample_weight = unpack_data(data)
        crf_loss, internal_losses = self.compute_loss(x, y, sample_weight)
        return {"crf_loss_val": crf_loss, "internal_losses_val": internal_losses}


def test():
    inputs = Input(shape=(None,), dtype='int32')
    output = Embedding(100, 40, trainable=True, mask_zero=True)(inputs)
    output = Bidirectional(GRU(64, return_sequences=True))(output)
    # output = Dense(9, activation=None)(output)
    crf = CRF(9)
    output = crf(output)[0]
    base_model = Model(inputs, output)
    model = ModelWithCRFLoss(base_model)
    model.build(input_shape=(None, None))
    model.compile('adam')

    x = [[5, 2, 3] * 3] * 10
    y = [[1, 2, 3] * 3] * 10

    model.fit(x=x, y=y, epochs=2, batch_size=2)
    model.save('model')


if __name__ == '__main__':
    test()

The error raised: OperatorNotAllowedInGraphError: iterating over tf.Tensor is not allowed: AutoGraph did not convert this function. Try decorating it directly with @tf.function.

Because I want to use the model on tf-serving, it asked the output to has the same batch dim with the input. So I want to get only the first dimension of the output of the CRF Layer. Is these any way to do it?

The root cause is that autograph doesn't yet support list comprehension. You can use tf.map_fn for the comprehension: return tf.map_fn(lambda i: i ** 2 if i > 0 else i, x)

You can also just do : index = 0 for image in inputs: .... My code ... index += 1

Hi @xuxingya , your model building is incorrect. replace output = crf(output)[0] with output = crf(output) will fix your problem. You can check double it at https://colab.research.google.com/drive/1T5VMO3Q2GCN27tqRpSkrO1a8-SyhBnmG?usp=sharing.

@bhack This issue is CRF related, so please label it with crf.

@xuxingya Also, you may find this PR #2552 helps. It uses Custom Training Loop programming style instead of Custom Training Step. Both methods can work.

TensorFlow Addons is transitioning to a minimal maintenance and release mode. New features will not be added to this repository. For more information, please see our public messaging on this decision: TensorFlow Addons Wind Down

Please consider sending feature requests / contributions to other repositories in the TF community with a similar charters to TFA: Keras Keras-CV Keras-NLP