Google recientemente proporcionó acceso gratuito a su unidad de procesamiento de
tensor (TPU) en la
plataforma de aprendizaje automático basada en la nube de
Colaboratory . El procesador tensorial es un circuito integrado especializado (ASIC) desarrollado por Google para tareas de aprendizaje automático utilizando la biblioteca TensorFlow. Decidí intentar aprender una red convolucional de TPU en Keras que reconoce objetos en imágenes CIFAR-10. El código completo de la solución se puede ver y ejecutar en la
computadora portátil .
Foto cloud.google.comProcesadores tensoriales
En Habré ya escribió cómo se organizan los TPU (
aquí ,
aquí y
aquí ), y también
por qué los TPU son adecuados para entrenar redes neuronales . Por lo tanto, no profundizaré en los detalles de la arquitectura de TPU, sino que solo consideraré las características que deben tenerse en cuenta al entrenar redes neuronales.
Ahora hay tres generaciones de procesadores tensoriales, el rendimiento de TPU de última y tercera generación es de 420 TFlops (billones de operaciones de punto flotante por segundo), contiene 128 GB de memoria de alto ancho de banda. Sin embargo, solo los TPU de segunda generación están disponibles en el Colaboratorio, que tienen 180 TFlops de rendimiento y 64 GB de memoria. En el futuro, consideraré estos TPU.
El procesador tensor consiste en cuatro chips, cada uno de los cuales contiene dos núcleos, un total de ocho núcleos en TPU. El entrenamiento de TPU se lleva a cabo en paralelo en todos los núcleos utilizando la replicación: cada núcleo ejecuta una copia del gráfico TensorFlow con un octavo del volumen de datos.
La base del procesador tensor es una unidad de matriz (MXU). Utiliza la astuta estructura de datos de una
matriz sistólica de 128x128 para la implementación eficiente de operaciones matriciales. Por lo tanto, para maximizar el uso de los recursos del equipo de TPU, la dimensión de la minimuestra o las características debe ser un múltiplo de 128 (
fuente ). Además, debido a la naturaleza del sistema de memoria TPU, es deseable que la dimensión de la mini-muestra y las características sean múltiplos de 8.
Plataforma colaborativa
Colaboratory es la plataforma en la nube de Google para avanzar en la tecnología de aprendizaje automático. Puede obtener una máquina virtual con las bibliotecas populares instaladas TensorFlow, Keras, sklearn, pandas, etc. de forma gratuita. Lo más conveniente es que puede ejecutar computadoras portátiles similares a Jupyter en el Colaboratorio. Las computadoras portátiles se almacenan en Google Drive, puede distribuirlas e incluso organizar la colaboración. Así es como se ve la computadora portátil en el Colaboratorio (
se puede hacer clic en la
imagen ):
Escribe el código en un navegador en una computadora portátil, se ejecuta en una máquina virtual en Google Cloud. El auto se le entrega por 12 horas, luego de lo cual se detiene. Sin embargo, nada le impide iniciar otra máquina virtual y trabajar otras 12 horas. Solo tenga en cuenta que después de que la máquina virtual se detiene, se eliminan todos los datos que contiene. Por lo tanto, no olvide guardar los datos necesarios en su computadora o Google Drive, y después de reiniciar la máquina virtual, vuelva a descargarlos.
Las instrucciones detalladas para trabajar en la plataforma de colaboración están
aquí ,
aquí y
aquí .
Conecte el procesador tensor al Colaboratorio
De forma predeterminada, Colaboratory no utiliza aceleradores de cálculo de GPU o TPU. Puede conectarlos en el menú Tiempo de ejecución -> Cambiar tipo de tiempo de ejecución -> Acelerador de hardware. En la lista que aparece, seleccione "TPU":
Después de elegir el tipo de acelerador, la máquina virtual a la que está conectada la computadora portátil Colaboratory se reiniciará y TPU estará disponible.
Si descargó algún dato a la máquina virtual, durante el proceso de reinicio se eliminará. Tienes que descargar los datos nuevamente.
Red neuronal Keras para el reconocimiento CIFAR-10
Como ejemplo, intentemos entrenar una red neuronal Keras en TPU que reconozca imágenes del
conjunto de datos CIFAR-10 . Este es un conjunto de datos popular que contiene pequeñas imágenes de objetos de 10 clases: avión, automóvil, pájaro, gato, venado, perro, rana, caballo, barco y camión. Las clases no se cruzan, cada objeto en la imagen pertenece a una sola clase.
Descargue el conjunto de datos CIFAR-10 usando Keras:
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
Para crear una red neuronal, obtuve una función separada. Crearemos el mismo modelo dos veces: la primera versión del modelo para TPU, en la que entrenaremos, y la segunda para la CPU, donde reconoceremos 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
Hasta ahora, los optimizadores Keras no se pueden usar en TPU, por lo tanto, al compilar un modelo, se especifica el optimizador de TensorFlow.
Creamos un modelo Keras para la CPU, que en el siguiente paso convertiremos a un modelo para TPU:
cpu_model = create_model()
Convierta la red neuronal Keras al modelo TPU
Los modelos en Keras y TensorFlow se pueden entrenar en la GPU sin ningún cambio. Hasta ahora no puede hacer esto en TPU, por lo que debe convertir el modelo que creamos en un modelo para TPU.
Primero debe averiguar dónde se encuentra el TPU disponible para nosotros. En la plataforma de colaboración, esto se puede hacer con el siguiente comando:
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
En mi caso, la dirección de TPU resultó ser así:
grpc://10.102.233.146:8470 . Las direcciones fueron diferentes para diferentes lanzamientos.
Ahora puede obtener el modelo para TPU utilizando la función
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)))
La primera línea incluye el registro en el nivel de información. Aquí está el registro de conversión del 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.Puede ver que se encontró TPU en la dirección que especificamos anteriormente, tiene 8 núcleos. También vemos una advertencia de que
tpu_model es experimental y se puede cambiar o eliminar en cualquier momento. Espero que con el tiempo sea posible entrenar modelos Keras directamente en TPU sin ninguna conversión.
Entrenamos modelo en TPU
El modelo para TPU se puede entrenar de la manera habitual para Keras llamando al método de
fit :
history = tpu_model.fit(x_train, y_train, batch_size=128*8, epochs=50, verbose=2)
¿Cuáles son las características aquí? Recordamos que para usar TPU de manera efectiva, el tamaño de la mini-muestra debe ser un múltiplo de 128. Además, se realiza capacitación en cada núcleo de TPU utilizando un octavo de todos los datos de la mini-muestra. Por lo tanto, establecemos el tamaño de la mini muestra durante el entrenamiento en 128 * 8, obtenemos 128 imágenes para cada núcleo de TPU. Puede usar un tamaño más grande, por ejemplo, 256 o 512, luego el rendimiento será mayor.
En mi caso, el entrenamiento de una era requiere un promedio de 6 s.
La calidad de la educación en la era 50:
Epoch 50/50
- 6s - loss: 0.2727 - sparse_categorical_accuracy: 0.9006La proporción de respuestas correctas a los datos para la capacitación fue del 90.06%. Verificamos la calidad de los datos de prueba usando TPU:
scores = tpu_model.evaluate(x_test, y_test, verbose=0, batch_size=batch_size * 8) print(" : %.2f%%" % (scores[1]*100))
: 80.79%Ahora guarde los pesos del modelo entrenado:
tpu_model.save_weights("cifar10_model.h5")
TensorFlow nos dará un mensaje de que los pesos se transfieren desde la TPU a la CPU:
INFO:tensorflow:Copying TPU weights to the CPUCabe señalar que los pesos de la red capacitada se guardaron en el disco de la máquina virtual colaborativa. Cuando la máquina virtual se detiene, todos los datos de ella se borrarán. Si no desea perder pesos entrenados, guárdelos en su computadora:
from google.colab import files files.download("cifar10_model.h5")
Reconociendo objetos en la CPU
Ahora intentemos usar un modelo entrenado en TPU para reconocer objetos en imágenes usando la CPU. Para hacer esto, cree el modelo nuevamente y cargue los pesos entrenados en TPU en él:
model = create_model() model.load_weights("cifar10_model.h5")
El modelo está listo para usar en el procesador central. Tratemos de reconocer con su ayuda una de las imágenes del conjunto de pruebas CIFAR-10:
index=111 plt.imshow(toimage(x_test[index])) plt.show()
La imagen es pequeña, pero puedes entender que este es un avión. Comenzamos el reconocimiento:
Obtenemos una lista de valores de salida de las neuronas, casi todas están cerca de cero, excepto el primer valor, que corresponde al 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]]
El reconocimiento fue exitoso!
Resumen
Fue posible demostrar la operatividad de TPU en la plataforma de colaboración, se puede utilizar para entrenar redes neuronales en Keras. Sin embargo, el conjunto de datos CIFAR-10 es demasiado pequeño; no es suficiente para cargar completamente los recursos de TPU. La aceleración en comparación con la GPU resultó ser pequeña (puede comprobarlo usted mismo eligiendo la GPU como acelerador en lugar de la TPU y volviendo a entrenar el modelo nuevamente).
En Habré hay un artículo en el que se
mide el rendimiento de TPU y GPU V100 en el entrenamiento de la red ResNet-50 . En esta tarea, la TPU mostró el mismo rendimiento que las cuatro GPU V100. ¡Es bueno que Google proporcione un acelerador de aprendizaje de redes neuronales tan poderoso de forma gratuita!
Video que muestra el entrenamiento de la red neuronal Keras en TPU.
Enlaces utiles
- Computadora portátil colaborativa con código de aprendizaje modelo Keras TPU completo .
- Cuaderno colaborativo con ejemplo de capacitación Keras TPU para reconocer ropa y zapatos de Fashion MNIST .
- Procesadores tensoriales en Google Cloud .
- Características de la arquitectura y uso de procesadores tensoriales .