Participe do Simpósio Women in ML em 7 de dezembro Inscreva-se agora

SavedModels do TF Hub no TensorFlow 2

Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

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 formato antigo do TF1 Hub e vem com um novo conjunto de APIs.

Esta página explica como reutilizar SavedModels do TF2 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 no Keras

Keras é a API de alto nível do TensorFlow para criar modelos de aprendizado profundo compondo objetos Keras Layer. 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 como usar uma incorporação de texto pré-treinada:

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 do Keras:

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

O colab 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. Veja a seção sobre ajuste fino abaixo para saber como mudar isso. Os pesos são compartilhados entre todos os aplicativos 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

O uso de um SavedModel do TensorFlow Hub (ou outros servidores HTTPS que implementam seu protocolo de hospedagem ) faz o download e o descompacta 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

A função hub.load(handle) baixa e descompacta um SavedModel (a menos que o handle já seja um caminho do sistema de arquivos) e, em seguida, retorna o resultado de carregá-lo com a função tf.saved_model.load() do TensorFlow . Portanto, hub.load() pode manipular qualquer SavedModel válido (ao contrário de seu predecessor 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 Guia SavedModel 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 codificados pela respectiva entrada e saída nomes e sujeito às restrições de forma e dtype da assinatura.

  • Os métodos @tf.function -decorated 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 pode ser chamado como uma função Python. Um exemplo simples pode ser output_tensor = obj(input_tensor, training=False) .

Isso deixa uma 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. Veja a seção sobre ajuste fino abaixo para algumas ressalvas e considere evitar isso para começar. Mesmo se você quiser ajustar, você pode querer ver se obj.trainable_variables aconselha a retreinar 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. Ele 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.

Economizando de Keras

A partir do TensorFlow 2, tf.keras.Model.save() e tf.keras.models.save_model() 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 se tornam disponíveis.

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

Para compartilhar uma peça de um modelo Keras, faça da peça um modelo em si e salve-o. Você pode colocar 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 do 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 pré- pretrainer para restaurar o ponto de verificação) e a última abordagem para ResNet (consulte legacy/image_classification/tfhub_export.py ).

Salvando do TensorFlow de baixo nível

Isso requer uma boa familiaridade com o SavedModel Guide do TensorFlow.

Se você deseja fornecer mais do que apenas uma assinatura de serviço, 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 em torno dele é 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 seus hiperparâmetros, aumentar o risco de overfitting e exigir aumento do conjunto de dados, especialmente para CNNs). Aconselhamos os consumidores do SavedModel a procurar o ajuste fino somente depois de estabelecer um bom regime de treinamento e somente se o editor do SavedModel o recomendar.

O ajuste fino altera os parâmetros de modelo "contínuos" que são treinados. Ele não altera as transformações codificadas, como tokenização de entrada de texto e tokens de mapeamento 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 a computação do SavedModel no modo de treinamento (pense em dropout etc.).

A colab de classificação de imagem contém um exemplo de ponta a ponta com ajuste fino opcional.

Reexportando o resultado do ajuste fino

Usuários avançados podem querer salvar os resultados do ajuste fino de volta em um SavedModel que pode ser usado em vez do originalmente carregado. 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 no futuro se e como seus consumidores devem ajustá-lo e forneça orientações na documentação.

Salvar de um modelo Keras deve fazer toda a mecânica de ajuste fino funcionar (economizar perdas de regularização de peso, declarar variáveis ​​treináveis, rastrear __call__ tanto para training=True quanto training=False , etc.)

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

Se o modelo usar dropout, normalização de lote ou técnicas de treinamento semelhantes que envolvem hiperparâmetros, defina-os para valores que façam sentido em muitos problemas de destino esperados e tamanhos de lote. (No momento em que este artigo foi escrito, salvar do Keras não facilita permitir que os consumidores os ajustem.)

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