Armazenando downloads de modelos em cache do TF Hub

Visão geral

A biblioteca tensorflow_hub atualmente oferece suporte a dois modos de download de 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 no TensorFlow. De qualquer forma, as chamadas para funções tensorflow_hub no código Python real podem e devem continuar a usar os URLs canônicos tfhub.dev dos modelos, que são portáveis ​​entre sistemas e navegáveis ​​para documentação. No caso raro em que o código do usuário precise da localização real do sistema de arquivos (após o download e a descompactação, 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) .

Cache de downloads compactados

A biblioteca tensorflow_hub por padrão, armazena em cache os modelos no sistema de arquivos quando eles são baixados de tfhub.dev (ou outros sites de hospedagem ) e descompactados. Este 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 é padronizado como um diretório temporário local, mas pode ser personalizado definindo a variável de ambiente TFHUB_CACHE_DIR (recomendado) ou passando o sinalizador de linha de comando --tfhub_cache_dir . O local de cache padrão /tmp/tfhub_modules (ou qualquer que seja os.path.join(tempfile.gettempdir(), "tfhub_modules") avaliado) deve funcionar na maioria dos casos.

Os usuários que preferem o cache persistente durante as reinicializações do sistema podem, em vez disso, definir TFHUB_CACHE_DIR para um local em seu diretório inicial. Por exemplo, um usuário do shell bash em um sistema Linux pode adicionar uma linha como a seguinte em ~/.bashrc

export TFHUB_CACHE_DIR=$HOME/.cache/tfhub_modules

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

Lendo do armazenamento remoto

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

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

ou definindo o sinalizador 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 que oferecem pouco espaço em disco, mas uma conexão rápida à Internet.

Executando em TPU em notebooks Colab

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

1) Use um bucket do GCS que o trabalhador da TPU possa acessar

A solução mais fácil é instruir a biblioteca tensorflow_hub para ler os modelos do bucket GCS do TF Hub conforme explicado acima. Os usuários com seu próprio bucket do GCS podem especificar um diretório em seu bucket como o local do cache com 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 por meio do host Colab

Outra solução alternativa é redirecionar todas as leituras (mesmo de variáveis ​​grandes) por meio do host 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 .