Treinando modelos de detecção de objetos mais rápido com TPUs
Modelos estado da arte em minutos, de graça!
A detecção de objetos é uma técnica de deep learning bastante utilizada, essencialmente quando o objetivo da tarefa envolve a localização dos objetos na imagem ou quando mais de um objeto de interesse aparece simultaneamente.
Porém, uma das desvantagens desse método é o custo computacional elevado, requisitando o uso de GPUs modernas e, ainda assim, podendo levar horas ou até dias para se completar o treinamento. Pensando nisso, o objetivo desse artigo é mostrar como é possível utilizar um hardware específico para deep learning, as TPUs, e treinar modelos dezenas de vezes mais rápido utilizando a API de detecção de objetos do TensorFlow na plataforma gratuita do Google Colab.
O que são TPUs e por que elas funcionam?
Cada componente de um computador possui uma função específica e sua arquitetura física foi pensada para melhor servir ao seu objetivo. As CPUs, por exemplo, conseguem lidar com diversas tarefas simultaneamente, além de realizar operações de leitura e escrita utilizadas o tempo todo em nossa interação com o sistema operacional. De igual modo, as GPUs são unidades especializadas em realizar operações matemáticas de maneira rápida possível, sendo capazes de lidar com tarefas complexas como renderizações 3D e simulações físicas.
Dito isso, as TPUs (Tensor Processing Units) foram criadas pelo Google e anunciadas em 2016 com o objetivo de ser um hardware especializado em tarefas de deep learning, possuindo uma arquitetura que consegue lidar com os tensores multidimensionais de maneira eficiente, conforme explicado em detalhes aqui.
Comumente, as TPUs são utilizadas em clusters com dezenas ou centenas de nós, acessíveis através de serviços de cloud. Felizmente, o Google Colab permite o uso gratuito de um cluster com alguns nós de TPU, sendo capaz de treinar modelos mais rapidamente do que qualquer outra GPU disponível pelo serviço.
Utilizando TPUs no Google Colab
Para exemplificar o uso das TPUs no Colab, vamos utilizar um dataset do Kaggle para detectar gatos e cachorros. As anotações foram produzidas pelo software LabelImg e todo o código fonte se encontra no Github. Para habilitar o cluster de TPU do Colab basta acessar o menu de ambiente de execução e clicar em “alterar tipo do ambiente”, conforme ilustrado abaixo:
Feito isso, é possível selecionar o uso da TPU. Tenha em mente que a quota de uso desse recurso é limitado, devido ao fato do serviço ser gratuito (embora haja uma versão pro). O Google não especifica quanto tempo o usuário consegue ter acesso ao hardware, mas minha recomendação é não selecionar o uso de TPUs a menos que de fato estiver utilizando uma!
Criando os arquivos TFRecord
O primeiro passo necessário em qualquer pipeline de detecção de objetos com o TensorFlow é transformar os arquivos XML contendo as labels das imagens em arquivos binários do tipo TFRecord. Esses arquivos são otimizados para funcionar em diversas APIs do TensorFlow, entre elas o TF Data, utilizado para ingestão de dados de forma paralela e com alto desempenho.
O TF Data é o padrão de ingestão de dados na API de detecção de objetos do TensorFlow, visando aumentar a performance do treinamento dos modelos reduzindo o gargalo na leitura e modificação dos dados.
Saber da existência e funcionamento da ingestão de dados é essencial para se utilizar as TPUs de maneira correta, visto que o processamento de dados nesse tipo de dispositivo ocorre de maneira muito acelerada, gerando um gargalo no processador para ler e transformar arquivos em massa. Explicando de maneira simples (detalhes podem ser vistos aqui) o treinamento dos modelos de deep learning geralmente passam por duas etapas. Na primeira etapa, a CPU faz a leitura de um batch de arquivos (24 imagens por exemplo) e aplica transformações a essas imagens (reduz a dimensão e aplica técnicas de aumento de dados). Feito isso, as imagens pré-processadas são transmitidas a um hardware de alto desempenho (GPUs ou TPUs) e estes realizam o treinamento do modelo.
Entretanto, quando o hardware processa as imagens a uma velocidade muito alta, é possível que a CPU não consiga ler e pré-processar o batch de dados na mesma velocidade que o modelo requer para ser treinado.
Isso gera um gargalo que limita a velocidade do treinamento do modelo ao número de imagens por segundo que o processador consegue ler e transformar.
Esse problema é resolvido através da API do TF data, com a leitura eficiente e paralelizada dos dados utilizados em treinamento. Porém, um detalhe que muita das vezes é ignorado é a forma em que os arquivos TFRecord devem ser criados. A documentação do TensorFlow recomenda criar um conjunto de arquivos que não sejam nem tão grandes e nem tão pequenos. Dito isso, uma boa prática é, ao invés de possuir um único arquivo TFRecord com 5 GB, ter 10 arquivos com 500 MB cada um, facilitando a paralelização da leitura.
O primeiro passo para atingir isso é mapear as anotações XML e as imagens em um arquivos CSV. Isso pode ser feito através do script “xml_to_csv” disponível no Github com o seguinte comando:
python3 xml_to_csv.py -i caminho-do-dataset -o ./train_labels.csvAs flags -i e -o são utilizadas para especificar o caminho em que o seu dataset se encontra e o caminho em que o resultado do arquivo XML será gerado.
Feito isso, é preciso executar o script “generate_tfrecord” para transformar as anotações efetivamente nos arquivos binários. O seguinte comando pode ser utilizado para esse fim:
python3 generate_tfrecord.py --csv_input tpu-training/train_labels.csv --output_path tpu-training/train.tfrecord --label_map tpu-training/label_map.pbtxt --n_splits 2 --img_path tpu-training/dataset/trainOs parâmetros utilizados no script são explicados em detalhes abaixo:
- csv_input: Caminho do arquivo CSV gerado no passo anterior.
- output_path: Nome base e caminho do arquivo TFRecord que será criado.
- label_map: Caminho do arquivo contendo as labels das classes a serem detectadas.
- n_splits: Número de arquivos TFRecord a serem gerados.
- img_path: Caminho das imagens do dataset.
Após a execução do algoritmo, dois arquivos TFRecord são gerados, conforme especificado nos parâmetros. Nesse exemplo em específico, os arquivos têm pouco mais de 10MB e poderiam ser apenas um único arquivo. O número de arquivos pode variar conforme você quiser, porém, eu geralmente tento manter um máximo de 300MB cada um.
Ambos os scripts foram criados originalmente por esse autor, porém, o arquivo de geração de TFRecord foi modificado por mim de modo a possibilitar a criação de múltiplos arquivos, além de algumas melhorias feitas na geração do CSV!
Google Cloud Service
Ao criar os arquivos TFRecord estamos habilitados a treinar o modelo. Porém, existe um importante detalhe que ainda precisa ser feito para possibilitar o uso do cluster de TPU no treinamento.
Normalmente, quando realizamos o treinamento de um modelo no Google Colab utilizando uma GPU, por exemplo, a execução do código acontece no mesmo ambiente em que o notebook é instanciado, facilitando a ingestão de dados. Por conta disso, podemos facilmente nos conectar ao Google Drive, clonar um repositório do git ou ainda fazer upload diretamente, como forma de acessar os dados (TFRecords) que serão utilizados no treinamento.
Quando utilizamos as TPUs, por outro lado, o treinamento ocorre em um cluster hospedado em uma cloud fora do ambiente jupyter em que estamos trabalhando. Por conta disso, os dados utilizados (TFRecords no nosso caso) precisam estar hospedados em algum serviço de nuvem, além de estarem disponíveis para acesso externo. Os dados presentes na sessão do Colab ou no Google Drive não conseguem ser diretamente acessados ao realizar o treinamento por TPU.
Dentre as inúmeras formas possíveis de se fazer isso, a mais fácil que encontrei é utilizar os serviços do Google Cloud Platform (GCP), mais especificamente o Cloud Storage. O primeiro passo é criar e ativar uma conta no GCP, sendo não só gratuito como dando $300 dólares de créditos para utilizar em testes dos serviços do GCP. Entretando, será necessário cadastrar um cartão de crédito válido. Fique tranquilo, nenhuma cobrança será realizada enquanto você tiver os créditos gratuitos na sua conta.
Feito isso, é preciso acessar o serviço de Cloud Storage, onde iremos armazenar os arquivos TFRecord que serão utilizados no treinamento. É possível acessar esse serviço no menu do console do GCP.
Depois, basta clicar em “criar intervalo” no menu superior esquerdo, para criar um bucket de armazenamento de dados e seguir as seguintes configurações:
- Após nomear o bucket conforme achar melhor, é preciso selecionar a região de armazenamento. Para o propósito de armazenar os dados para treinamento eu recomendo utilizar uma única região, de preferência na Carolina do Sul (us-east1).
- A classe de armazenamento pode ser standard.
- No menu de controle de acessos, remover a opção “aplicar a prevenção de acesso público neste bucket”.
- Na proteção de dados, pode-se manter a opção “nenhum”.
- Basta clicar em criar e seu bucket estará disponível.
Por fim, vamos entrar no bucket e criar uma pasta para treinamento e outra para teste. Em cada pasta, basta fazer o upload dos arquivos TFRecord respectivos arrastando eles para a pasta ou clicando em “fazer upload de arquivos”.
Upload do modelo base
Uma técnica muito comum na detecção de objetos é utilizar modelos base para transferência de aprendizado. Esses modelos podem ser escolhidos no model zoo do TensorFlow baseado no trade-off entre custo computacional e mAP. Ao escolher o modelo base, é preciso baixá-lo e fazer o upload no bucket do GCP, visto que o mesmo também deverá estar acessível pelo cluster de TPUs que será utilizado no treinamento.
O arquivo contendo as labels do modelo, geralmente chamado de label_map.pbtxt também deverá ser disponibilizado no bucket. No total, são 4 objetos que deverão ser colocados no bucket, são eles: pasta de treinamento e pasta de teste com os respectivos TFRecords, modelo base pré-treinado e o label_map.
Configurando a pipeline
Todo modelo de detecção de objetos do TensorFlow possui um arquivo de configuração, geralmente chamado pipeline.config. Esse arquivo possui, como o nome sugere, uma série de hiperparâmetros e configurações base para a execução do modelo, conforme especificado pelo usuário. Um passo necessário antes de realizar o treinamento no Colab é garantir que esse arquivo irá conter os caminhos dos arquivos de treinamento, validação e do modelo pré treinado corretamente configurados, conforme instruido:
- Altere o “fine_tune_checkpoint” pelo caminho do checkpoint do modelo pré-treinado no bucket do GCP. No meu caso, o resultado é esse:
fine_tune_checkpoint: "gs://tpu-object-detection/pretrained_model/efficientdet_d1_coco17_tpu-32/checkpoint/ckpt-0"Onde “tpu-object-detection” é o nome do meu bucket e “pretrained_model” é uma pasta criada por mim para colocar o modelo.
- Altere o label_map_path (em dois lugares) pelo caminho do seu label_map
label_map_path: "gs://tpu-object-detection/label_map.pbtxt"- Altere o “input_path” que fica dentro de “train_input_reader” com os caminhos dos TFRecords de treinamento
- Altere o “input_path” que fica dentro de “eval_input_reader” com os caminhos dos TFRecords de teste.
Treinando o modelo com o Google Colab
Finalmente, é possível criar a estrutura e realizar o treinamento no Colab. Para tal, é necessário primeiro iniciar a conexão com o cluster de TPUs na cloud, através do seguinte código:
Ao executar esse comando, você verá informações relativas ao cluster de TPU a qual conseguiu acesso, como a quantidade de nós disponíveis:
O próximo passo é fazer a definição dos hiperparâmetros necessários para execução do modelo de detecção de objetos, como o número de steps de treinamento e validação e o batch size. É interessante notar que o batch size é multiplicado pelo número de réplicas (nós) disponíveis no cluster, conforme mostrado no código abaixo:
Antes de efetuar o treinamento, é preciso se autenticar com a conta do Google para garantir que você poderá acessar os itens no bucket, executando o seguinte código:
from google.colab import authauth.authenticate_user()
Por fim, basta executar o comando que inicia o treinamento do modelo:
Se tudo estiver funcionado o treinamento deve ser iniciado em alguns instantes, levando um pouco mais de tempo do que o comum devido ao fato dos dados precisarem ser dividos pelo cluster.
Conforme você irá observar, o treinamento com TPUs é muito mais rápido do que utilizando apenas uma GPU, além do fato de conseguir utilizar um batch maior de imagens por iteração, aumentando a velocidade em que uma época é processada.
Conclusão
O treinamento utilizando TPUs é um pouco mais trabalhoso do que em uma simples CPU ou GPU. Entretando, tarefas complexas como detecção de objetos podem ser acelerados dezenas de vezes através do uso desse tipo de hardware, aumentando a velocidade de iteração e testes de hipótese do modelo.
