Modelos salvos do TF Hub no TensorFlow 2

O formato SavedModel do TensorFlow 2 é a maneira recomendada de compartilhar modelos pré-treinados e peças de modelo no TensorFlow Hub. Ele substitui o antigo formato TF1 Hub e vem com um novo conjunto de APIs.

Esta página explica como reutilizar TF2 SavedModels em um programa TensorFlow 2 com a API hub.load() de baixo nível e seu wrapper hub.KerasLayer . (Normalmente, hub.KerasLayer é combinado com outros tf.keras.layers para construir um modelo Keras ou o model_fn de um TF2 Estimator.) Essas APIs também podem carregar os modelos legados no formato TF1 Hub, dentro dos limites, consulte o guia de compatibilidade .

Os usuários do TensorFlow 1 podem atualizar para o TF 1.15 e usar as mesmas APIs. Versões mais antigas do TF1 não funcionam.

Usando SavedModels do TF Hub

Usando um SavedModel em Keras

Keras é a API de alto nível do TensorFlow para construir modelos de aprendizado profundo compondo objetos da camada Keras. A biblioteca tensorflow_hub fornece a classe hub.KerasLayer que é inicializada com a URL (ou caminho do sistema de arquivos) de um SavedModel e, em seguida, fornece o cálculo do SavedModel, incluindo seus pesos pré-treinados.

Aqui está um exemplo de uso de incorporação de texto pré-treinado:

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://tfhub.dev/google/nnlm-en-dim128/2"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

A partir disso, um classificador de texto pode ser construído da maneira usual de Keras:

model = tf.keras.Sequential([
    embed,
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(1, activation="sigmoid"),
])

A colaboração de classificação de texto é um exemplo completo de como treinar e avaliar tal classificador.

Os pesos do modelo em um hub.KerasLayer são definidos como não treináveis ​​por padrão. Consulte a seção sobre ajuste fino abaixo para saber como mudar isso. Os pesos são compartilhados entre todas as aplicações do mesmo objeto de camada, como de costume no Keras.

Usando um SavedModel em um estimador

Os usuários da API Estimator do TensorFlow para treinamento distribuído podem usar SavedModels do TF Hub escrevendo seu model_fn em termos de hub.KerasLayer entre outros tf.keras.layers .

Nos bastidores: download e armazenamento em cache do SavedModel

Usar um SavedModel do TensorFlow Hub (ou outros servidores HTTPS que implementam seu protocolo de hospedagem ) faz download e descompacta-o no sistema de arquivos local, se ainda não estiver presente. A variável de ambiente TFHUB_CACHE_DIR pode ser definida para substituir o local temporário padrão para armazenar em cache os SavedModels baixados e descompactados. Para obter detalhes, consulte Cache .

Usando um SavedModel no TensorFlow de baixo nível

Alças de modelo

SavedModels podem ser carregados a partir de um handle especificado, onde o handle é um caminho do sistema de arquivos, URL do modelo TFhub.dev válido (por exemplo, "https://tfhub.dev/..."). Os URLs dos modelos Kaggle espelham os tratamentos TFhub.dev de acordo com nossos Termos e a licença associada aos ativos do modelo, por exemplo, "https://www.kaggle.com/...". Os identificadores dos modelos Kaggle são equivalentes ao identificador TFhub.dev correspondente.

A função hub.load(handle) baixa e descompacta um SavedModel (a menos que handle já seja um caminho do sistema de arquivos) e, em seguida, retorna o resultado do carregamento com a função integrada do TensorFlow tf.saved_model.load() . Portanto, hub.load() pode lidar com qualquer SavedModel válido (ao contrário de seu antecessor hub.Module para TF1).

Tópico avançado: o que esperar do SavedModel após o carregamento

Dependendo do conteúdo do SavedModel, o resultado de obj = hub.load(...) pode ser invocado de várias maneiras (conforme explicado com muito mais detalhes no SavedModel Guide do TensorFlow:

  • As assinaturas de serviço do SavedModel (se houver) são representadas como um dicionário de funções concretas e podem ser chamadas como tensors_out = obj.signatures["serving_default"](**tensors_in) , com dicionários de tensores digitados pela respectiva entrada e saída nomes e sujeitos às restrições de formato e tipo da assinatura.

  • Os métodos @tf.function decorados do objeto salvo (se houver) são restaurados como objetos tf.function que podem ser chamados por todas as combinações de argumentos Tensor e não Tensor para os quais o tf.function foi rastreado antes de salvar. Em particular, se houver um método obj.__call__ com rastreamentos adequados, o próprio obj poderá ser chamado como uma função Python. Um exemplo simples poderia ser output_tensor = obj(input_tensor, training=False) .

Isso deixa enorme liberdade nas interfaces que SavedModels pode implementar. A interface Reusable SavedModels para obj estabelece convenções para que o código do cliente, incluindo adaptadores como hub.KerasLayer , saiba como usar o SavedModel.

Alguns SavedModels podem não seguir essa convenção, especialmente modelos inteiros que não devem ser reutilizados em modelos maiores e apenas fornecem assinaturas de serviço.

As variáveis ​​treináveis ​​em um SavedModel são recarregadas como treináveis ​​e tf.GradientTape as observará por padrão. Consulte a seção sobre ajuste fino abaixo para algumas advertências e considere evitar isso para começar. Mesmo se você quiser fazer um ajuste fino, você pode querer ver se obj.trainable_variables aconselha treinar novamente apenas um subconjunto das variáveis ​​​​originalmente treináveis.

Criando SavedModels para TF Hub

Visão geral

SavedModel é o formato de serialização padrão do TensorFlow para modelos treinados ou peças de modelo. Ele armazena os pesos treinados do modelo junto com as operações exatas do TensorFlow para realizar seu cálculo. Pode ser usado independentemente do código que o criou. Em particular, ele pode ser reutilizado em diferentes APIs de construção de modelos de alto nível, como Keras, porque as operações do TensorFlow são sua linguagem básica comum.

Salvando de Keras

A partir do TensorFlow 2, tf.keras.Model.save() e tf.keras.models.save_model() assumem como padrão o formato SavedModel (não HDF5). Os SavedModels resultantes que podem ser usados ​​com hub.load() , hub.KerasLayer e adaptadores semelhantes para outras APIs de alto nível à medida que ficam disponíveis.

Para compartilhar um modelo Keras completo, basta salvá-lo com include_optimizer=False .

Para compartilhar uma peça de um modelo Keras, transforme a peça em um modelo e salve-o. Você pode definir o código assim desde o início....

piece_to_share = tf.keras.Model(...)
full_model = tf.keras.Sequential([piece_to_share, ...])
full_model.fit(...)
piece_to_share.save(...)

...ou recorte a peça para compartilhar após o fato (se estiver alinhada com as camadas do seu modelo completo):

full_model = tf.keras.Model(...)
sharing_input = full_model.get_layer(...).get_output_at(0)
sharing_output = full_model.get_layer(...).get_output_at(0)
piece_to_share = tf.keras.Model(sharing_input, sharing_output)
piece_to_share.save(..., include_optimizer=False)

Os modelos TensorFlow no GitHub usam a abordagem anterior para BERT (consulte nlp/tools/export_tfhub_lib.py , observe a divisão entre core_model para exportação e o pretrainer para restaurar o ponto de verificação) e a última abordagem para ResNet (consulte legado/image_classification/tfhub_export.py ).

Salvando no TensorFlow de baixo nível

Isso requer boa familiaridade com o SavedModel Guide do TensorFlow.

Se quiser fornecer mais do que apenas uma assinatura de serviço, você deve implementar a interface Reusable SavedModel . Conceitualmente, isso parece

class MyMulModel(tf.train.Checkpoint):
  def __init__(self, v_init):
    super().__init__()
    self.v = tf.Variable(v_init)
    self.variables = [self.v]
    self.trainable_variables = [self.v]
    self.regularization_losses = [
        tf.function(input_signature=[])(lambda: 0.001 * self.v**2),
    ]

  @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
  def __call__(self, inputs):
    return tf.multiply(inputs, self.v)

tf.saved_model.save(MyMulModel(2.0), "/tmp/my_mul")

layer = hub.KerasLayer("/tmp/my_mul")
print(layer([10., 20.]))  # [20., 40.]
layer.trainable = True
print(layer.trainable_weights)  # [2.]
print(layer.losses)  # 0.004

Afinação

Treinar as variáveis ​​já treinadas de um SavedModel importado junto com aquelas do modelo ao seu redor é chamado de ajuste fino do SavedModel. Isso pode resultar em melhor qualidade, mas muitas vezes torna o treinamento mais exigente (pode levar mais tempo, depender mais do otimizador e de seus hiperparâmetros, aumentar o risco de overfitting e exigir aumento do conjunto de dados, especialmente para CNNs). Aconselhamos os consumidores do SavedModel a fazer o ajuste fino somente depois de terem estabelecido um bom regime de treinamento e somente se o editor do SavedModel o recomendar.

O ajuste fino altera os parâmetros do modelo "contínuo" que são treinados. Ele não altera as transformações codificadas, como a tokenização da entrada de texto e o mapeamento de tokens para suas entradas correspondentes em uma matriz de incorporação.

Para consumidores SavedModel

Criando um hub.KerasLayer como

layer = hub.KerasLayer(..., trainable=True)

permite o ajuste fino do SavedModel carregado pela camada. Ele adiciona os pesos treináveis ​​​​e os regularizadores de peso declarados no SavedModel ao modelo Keras e executa o cálculo do SavedModel no modo de treinamento (pense em abandono, etc.).

O colab de classificação de imagens contém um exemplo completo com ajuste fino opcional.

Reexportando o resultado do ajuste fino

Usuários avançados podem querer salvar os resultados do ajuste fino em um SavedModel que pode ser usado em vez do carregado originalmente. Isso pode ser feito com código como

loaded_obj = hub.load("https://tfhub.dev/...")
hub_layer = hub.KerasLayer(loaded_obj, trainable=True, ...)

model = keras.Sequential([..., hub_layer, ...])
model.compile(...)
model.fit(...)

export_module_dir = os.path.join(os.getcwd(), "finetuned_model_export")
tf.saved_model.save(loaded_obj, export_module_dir)

Para criadores de SavedModel

Ao criar um SavedModel para compartilhamento no TensorFlow Hub, pense antecipadamente se e como seus consumidores devem ajustá-lo e forneça orientação na documentação.

Salvar a partir de um modelo Keras deve fazer com que toda a mecânica de ajuste fino funcione (economizando perdas de regularização de peso, declarando variáveis ​​​​treináveis, rastreando __call__ para ambos training=True e training=False , etc.)

Escolha uma interface de modelo que funcione bem com fluxo gradiente, por exemplo, logits de saída em vez de probabilidades softmax ou previsões top-k.

Se o modelo usar eliminação, normalização de lote ou técnicas de treinamento semelhantes que envolvam hiperparâmetros, defina-os com valores que façam sentido em muitos problemas-alvo e tamanhos de lote esperados. (No momento em que este livro foi escrito, economizar no Keras não tornava mais fácil permitir que os consumidores os ajustassem.)

Os regularizadores de peso em camadas individuais são salvos (com seus coeficientes de força de regularização), mas a regularização de peso de dentro do otimizador (como tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) é perdida. Aconselhe os consumidores sobre seu SavedModel adequadamente.