O Google recentemente forneceu acesso gratuito à sua unidade de processamento
tensor (TPU) na
plataforma de aprendizado de máquina baseada em nuvem do
Colaboratory . O processador tensor é um circuito integrado especializado (ASIC) desenvolvido pelo Google para tarefas de aprendizado de máquina usando a biblioteca TensorFlow. Decidi tentar aprender a rede convolucional TPU em Keras, que reconhece objetos nas imagens CIFAR-10. O código completo da solução pode ser visualizado e executado no
laptop .
Foto cloud.google.comProcessadores tensores
Em Habré já escreveu como as TPUs são organizadas (
aqui ,
aqui e
aqui ) e também
por que as TPUs são adequadas para o treinamento de redes neurais . Portanto, não vou me aprofundar nos detalhes da arquitetura TPU, mas apenas considerarei os recursos que precisam ser considerados ao treinar redes neurais.
Agora, existem três gerações de processadores tensoriais, o desempenho das TPUs da terceira e última geração é de 420 TFlops (trilhões de operações de ponto flutuante por segundo) e contém 128 GB de memória de alta largura de banda. No entanto, apenas as TPUs de segunda geração estão disponíveis no Colaboratory, com 180 TFlops de desempenho e 64 GB de memória. No futuro, vou considerar esses TPUs.
O processador tensor consiste em quatro chips, cada um dos quais contém dois núcleos, um total de oito núcleos em TPU. O treinamento em TPU é realizado em paralelo em todos os núcleos usando replicação: cada núcleo executa uma cópia do gráfico TensorFlow com um oitavo do volume de dados.
A base do processador tensor é uma unidade matricial (MXU). Ele usa a estrutura de dados esperta de uma
matriz sistólica de 128x128 para implementação eficiente de operações da matriz. Portanto, para maximizar o uso dos recursos do equipamento TPU, a dimensão da miniamostra ou dos recursos deve ser um múltiplo de 128 (
fonte ). Além disso, devido à natureza do sistema de memória TPU, é desejável que a dimensão da mini-amostra e recursos seja um múltiplo de 8.
Plataforma Colaborativa
Colaboratory é a plataforma em nuvem do Google para o avanço da tecnologia de aprendizado de máquina. Você pode obter uma máquina virtual gratuitamente com as bibliotecas populares instaladas TensorFlow, Keras, sklearn, pandas etc. O mais conveniente é que você possa executar laptops semelhantes ao Jupyter no Colaboratory. Os laptops são armazenados no Google Drive, você pode distribuí-los e até organizar a colaboração. É assim que o laptop se parece no Colaboratory (a
imagem é clicável ):
Você escreve o código em um navegador em um laptop, ele é executado em uma máquina virtual no Google Cloud. O carro é emitido para você por 12 horas, após as quais para. No entanto, nada impede que você inicie outra máquina virtual e trabalhe outras 12 horas. Lembre-se de que, após a máquina virtual parar, todos os dados dela serão excluídos. Portanto, não esqueça de salvar os dados necessários no seu computador ou no Google Drive e, após reiniciar a máquina virtual, faça o download novamente.
Instruções detalhadas para trabalhar na plataforma Colaboratory estão
aqui ,
aqui e
aqui .
Conecte o processador tensor ao Colaboratory
Por padrão, o Colaboratory não usa aceleradores de cálculo de GPU ou TPU. Você pode conectá-los no menu Tempo de execução -> Alterar tipo de tempo de execução -> Acelerador de hardware. Na lista exibida, selecione "TPU":
Após escolher o tipo de acelerador, a máquina virtual à qual o laptop Colaboratory está conectado será reiniciada e o TPU estará disponível.
Se você baixou quaisquer dados para a máquina virtual, durante o processo de reinicialização, eles serão excluídos. Você precisa baixar os dados novamente.
Rede Neural Keras para reconhecimento CIFAR-10
Como exemplo, vamos tentar treinar uma rede neural Keras em TPU que reconheça imagens do
conjunto de dados CIFAR-10 . Este é um conjunto de dados popular que contém pequenas imagens de objetos de 10 classes: avião, carro, pássaro, gato, veado, cachorro, sapo, cavalo, navio e caminhão. As classes não se cruzam, cada objeto na imagem pertence a apenas uma classe.
Faça o download do conjunto de dados CIFAR-10 usando Keras:
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
Para criar uma rede neural, recebi uma função separada. Criaremos o mesmo modelo duas vezes: a primeira versão do modelo para TPU, na qual iremos treinar, e a segunda para a CPU, onde reconheceremos objetos.
def create_model(): input_layer = Input(shape=(32, 32, 3), dtype=tf.float32, name='Input') x = BatchNormalization()(input_layer) x = Conv2D(32, (3, 3), padding='same', activation='relu')(x) x = Conv2D(32, (3, 3), activation='relu', padding='same')(x) x = MaxPooling2D(pool_size=(2, 2))(x) x = Dropout(0.25)(x) x = BatchNormalization()(x) x = Conv2D(64, (3, 3), padding='same', activation='relu')(x) x = Conv2D(64, (3, 3), activation='relu')(x) x = MaxPooling2D(pool_size=(2, 2))(x) x = Dropout(0.25)(x) x = Flatten()(x) x = Dense(512, activation='relu')(x) x = Dropout(0.5)(x) output_layer = Dense(10, activation='softmax')(x) model = Model(inputs=[input_layer], outputs=[output_layer]) model.compile( optimizer=tf.train.AdamOptimizer(0.001), loss=tf.keras.losses.sparse_categorical_crossentropy, metrics=['sparse_categorical_accuracy']) return model
Até o momento, os otimizadores Keras não podem ser usados em TPUs; portanto, ao compilar um modelo, o otimizador do TensorFlow é especificado.
Criamos um modelo Keras para a CPU, que na próxima etapa iremos converter para um modelo para TPU:
cpu_model = create_model()
Converter rede neural de Keras em modelo TPU
Os modelos no Keras e no TensorFlow podem ser treinados na GPU sem nenhuma alteração. Você ainda não pode fazer isso em TPU, portanto, é necessário converter o modelo que criamos em um modelo para TPU.
Primeiro, você precisa descobrir onde o TPU disponível para nós está localizado. Na plataforma Colaboratory, isso pode ser feito com o seguinte comando:
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
No meu caso, o endereço TPU ficou assim -
grpc://10.102.233.146:8470 . Os endereços foram diferentes para diferentes lançamentos.
Agora você pode obter o modelo para TPU usando a função
keras_to_tpu_model :
tf.logging.set_verbosity(tf.logging.INFO) tpu_model = tf.contrib.tpu.keras_to_tpu_model( cpu_model, strategy=tf.contrib.tpu.TPUDistributionStrategy( tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
A primeira linha inclui o registro no nível de informações. Aqui está o log de conversão do modelo:
INFO:tensorflow:Querying Tensorflow master (b'grpc://10.102.233.146:8470') for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
...
WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.Você pode ver que o TPU foi encontrado no endereço especificado anteriormente, possui 8 núcleos. Também vemos um aviso de que
tpu_model é experimental e pode ser alterado ou excluído a qualquer momento. Espero que, com o tempo, seja possível treinar os modelos Keras diretamente em TPU sem nenhuma conversão.
Treinamos modelo em TPU
O modelo para TPU pode ser treinado da maneira usual para Keras chamando o método de
fit :
history = tpu_model.fit(x_train, y_train, batch_size=128*8, epochs=50, verbose=2)
Quais são os recursos aqui. Lembramos que, para usar de forma eficiente as TPUs, o tamanho da mini amostra deve ser múltiplo de 128. Além disso, o treinamento é realizado em cada núcleo da TPU, usando um oitavo de todos os dados da mini amostra. Portanto, definimos o tamanho da mini-amostra durante o treinamento para 128 * 8, obtemos 128 fotos para cada núcleo de TPU. Você pode usar um tamanho maior, por exemplo, 256 ou 512; o desempenho será maior.
No meu caso, o treinamento de uma época requer uma média de 6 s.
A qualidade da educação na era 50:
Epoch 50/50
- 6s - loss: 0.2727 - sparse_categorical_accuracy: 0.9006A parcela de respostas corretas nos dados para treinamento foi de 90,06%. Verificamos a qualidade dos dados de teste usando TPU:
scores = tpu_model.evaluate(x_test, y_test, verbose=0, batch_size=batch_size * 8) print(" : %.2f%%" % (scores[1]*100))
: 80.79%Agora salve os pesos do modelo treinado:
tpu_model.save_weights("cifar10_model.h5")
O TensorFlow nos dará uma mensagem de que os pesos são transferidos do TPU para a CPU:
INFO:tensorflow:Copying TPU weights to the CPUNote-se que os pesos da rede treinada foram salvos no disco da máquina virtual Colaboratory. Quando a máquina virtual é parada, todos os dados serão apagados. Se você não deseja perder pesos treinados, salve-os no seu computador:
from google.colab import files files.download("cifar10_model.h5")
Reconhecendo objetos na CPU
Agora vamos tentar usar um modelo treinado em TPU para reconhecer objetos em imagens usando a CPU. Para fazer isso, crie o modelo novamente e carregue os pesos treinados em TPU nele:
model = create_model() model.load_weights("cifar10_model.h5")
O modelo está pronto para uso no processador central. Vamos tentar reconhecer, com sua ajuda, uma das imagens da suíte de testes CIFAR-10:
index=111 plt.imshow(toimage(x_test[index])) plt.show()
A imagem é pequena, mas você pode entender que este é um avião. Começamos o reconhecimento:
Obtemos uma lista dos valores de saída dos neurônios, quase todos eles são próximos de zero, exceto o primeiro valor, que corresponde ao plano.
[[9.81738389e-01 2.91262069e-07 1.82225723e-02 9.78524668e-07
5.89265142e-07 6.76223244e-10 1.03252004e-10 9.23009047e-09
3.71878523e-05 3.16599618e-08]]
O reconhecimento foi bem sucedido!
Sumário
Foi possível demonstrar a operacionalidade do TPU na plataforma Colaboratory, que pode ser usada para treinar redes neurais no Keras. No entanto, o conjunto de dados CIFAR-10 é muito pequeno; não basta carregar totalmente os recursos de TPU. A aceleração comparada à GPU acabou sendo pequena (você pode verificar a si mesmo escolhendo a GPU como um acelerador em vez da TPU e treinando novamente o modelo).
Em Habré, há um artigo no qual
medimos o desempenho de TPUs e GPUs V100 no treinamento da rede ResNet-50 . Nesta tarefa, o TPU mostrou o mesmo desempenho que as quatro GPUs V100. É bom que o Google forneça um acelerador de aprendizado de rede neural tão poderoso de graça!
Vídeo demonstrando o treinamento da rede neural Keras em TPU.
Links úteis
- Laptop colaborativo com código completo de aprendizado do modelo Keras TPU .
- Caderno colaborativo com exemplo de treinamento em Keras TPU para reconhecimento de roupas e sapatos do Fashion MNIST .
- Processadores tensores no Google Cloud .
- Características da arquitetura e uso de processadores tensores .