Armazenando downloads de modelos em cache do TF Hub

Visão geral

A biblioteca tensorflow_hub atualmente suporta dois modos para baixar modelos. Por padrão, um modelo é baixado como um arquivo compactado e armazenado em cache no disco. Em segundo lugar, os modelos podem ser lidos diretamente do armazenamento remoto para o TensorFlow. De qualquer forma, as chamadas para as funções tensorflow_hub no código Python podem e devem continuar a usar os URLs canônicos do tfhub.dev dos modelos, que são portáteis entre sistemas e navegáveis ​​para documentação. No raro caso em que o código do usuário precise do local real do sistema de arquivos (após baixar e descompactar, ou após resolver um identificador de modelo em um caminho do sistema de arquivos), ele pode ser obtido pela função hub.resolve(handle) .

Armazenamento em cache de downloads compactados

Por padrão, a biblioteca tensorflow_hub armazena em cache os modelos no sistema de arquivos quando eles são baixados do tfhub.dev (ou de outros sites de hospedagem ) e descompactados. Esse modo é recomendado para a maioria dos ambientes, exceto se o espaço em disco for escasso, mas a largura de banda e a latência da rede forem excelentes.

O local de download é definido por padrão como um diretório temporário local, mas pode ser personalizado definindo a variável de ambiente TFHUB_CACHE_DIR (recomendado) ou passando o parâmetro de linha de comando --tfhub_cache_dir . O local de cache padrão /tmp/tfhub_modules (ou qualquer valor retornado por os.path.join(tempfile.gettempdir(), "tfhub_modules") deve funcionar na maioria dos casos.

Usuários que preferem o cache persistente mesmo após reinicializações do sistema podem definir a variável de ambiente TFHUB_CACHE_DIR para um local em seu diretório pessoal. Por exemplo, um usuário do shell bash em um sistema Linux pode adicionar uma linha como a seguinte ao arquivo ~/.bashrc :

export TFHUB_CACHE_DIR=$HOME/.cache/tfhub_modules

...reinicie o shell e então este local será utilizado. Ao usar um local persistente, esteja ciente de que não há limpeza automática.

Leitura de dados armazenados remotamente

Os usuários podem instruir a biblioteca tensorflow_hub a ler diretamente os modelos do armazenamento remoto (GCS) em vez de baixá-los localmente com:

os.environ["TFHUB_MODEL_LOAD_FORMAT"] = "UNCOMPRESSED"

ou definindo o parâmetro de linha de comando --tfhub_model_load_format como UNCOMPRESSED . Dessa forma, nenhum diretório de cache é necessário, o que é especialmente útil em ambientes com pouco espaço em disco, mas com uma conexão de internet rápida.

Executando em TPU em notebooks do Colab

Em colab.research.google.com , o download de modelos compactados entrará em conflito com o ambiente de execução da TPU, visto que a carga de trabalho computacional é delegada a outra máquina que, por padrão, não tem acesso ao local do cache. Existem duas soluções alternativas para essa situação:

1) Utilize um bucket do GCS que o worker da TPU possa acessar.

A solução mais simples é instruir a biblioteca tensorflow_hub a ler os modelos do bucket GCS do TF Hub, conforme explicado acima. Usuários com seu próprio bucket GCS podem, em vez disso, especificar um diretório em seu bucket como local de cache com um código como:

import os
os.environ["TFHUB_CACHE_DIR"] = "gs://my-bucket/tfhub-modules-cache"

...antes de chamar a biblioteca tensorflow_hub .

2) Redirecione todas as leituras através do host do Colab.

Outra solução alternativa é redirecionar todas as leituras (mesmo de variáveis ​​grandes) através do host do Colab:

load_options =
tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')
reloaded_model = hub.load("https://tfhub.dev/...", options=load_options)

Nota: Veja mais informações sobre identificadores válidos aqui .