Soluciona problemas de TensorFlow: TPU

Esta guía, junto con las Preguntas frecuentes, proporciona ayuda para la solución de problemas a los usuarios que entrenan modelos de TensorFlow en Cloud TPU. Si estás solucionando problemas PyTorch o JAX, puedes consultar los documentos de solución de problemas para esos marcos de trabajo:

Para obtener guías más generales sobre cómo usar Cloud TPU, consulta:

Descripción general

Los problemas habituales que se encuentran con las Cloud TPU se incluyen en las siguientes categorías:

  1. Problemas para conectarte a la TPU

  2. Cómo depurar errores comunes

  3. Cómo reducir el uso de la memoria

  4. Mejora la velocidad de entrenamiento

  5. Depura caídas en la exactitud del modelo

Problemas para establecer la conexión al servidor de TPU

En esta sección, se describe cómo solucionar problemas situaciones en las que TensorFlow deja de responder o imprime un error cuando conectando a la TPU. El paso de compilación de grafos de TPU puede llevar mucho tiempo a los modelos, así que deja que la secuencia de comandos se ejecute por al menos 5 minutos antes de concluir dejó de responder.

El primer paso es verificar si el problema es con el servidor o con la canalización de entrenamiento de TensorFlow. Para ello, ejecuta el instructivo de MNIST con la URL del servidor de TPU y verifica que funcione correctamente. Si todavía hay problemas de conexión con el instructivo de MNIST, esto confirma que el problema está relacionado con el servidor de TPU. En este caso, ocurre lo siguiente:

  1. Ejecuta el siguiente comando para mostrar las TPU disponibles. Reemplaza zone y project-id por tu zona y el ID de tu proyecto.

    (vm)$ gcloud compute tpus tpu-vm list --zone zone --project project-id

    Esto muestra una salida como la siguiente:

    NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
    demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

  2. Verifica que estés pasando el valor correcto a --tpu (demo-tpu en el ejemplo anterior) y que esta TPU aparezca como READY.

  3. Si tu TPU no aparece como READY o si aún tienes problemas para conectarte, reinicia el servidor de forma manual con lo siguiente:

    (vm)$ gcloud compute tpus tpu-vm stop $TPU_SERVER_NAME && gcloud compute tpus tpu-vm start $TPU_SERVER_NAME

    En el ejemplo anterior, $TPU_SERVER_NAME es demo-tpu. Esto puede tardar varios minutos en completarse.

  4. Vuelve a ejecutar el comando ... tpus list y espera a que la TPU esté en estado READY. Esto puede tardar varios minutos.

  5. Vuelve a ejecutar el instructivo de MNIST.

  6. Si todavía tienes problemas para ejecutar el instructivo de MNIST, pide ayuda mediante uno de los mecanismos que se describen en Obtén asistencia.

Si el ejemplo de MNIST se ejecuta correctamente, pero tu modelo aún deja de responder, es probable que el problema esté relacionado con la canalización de entrenamiento. Para depurarlo, comienza reemplazando la estrategia de TPU de tu código con el estrategia predeterminada. Cuando usas la estrategia predeterminada, siempre que la uses strategy.scope() o strategy.run(), el modelo ejecuta en la CPU (o GPU si está presente) en la TPU. Si el modelo se ejecuta en la CPU y no en la TPU, debe haber un problema específico de la TPU. Si aún no se ejecuta, se recomienda depurar el problema en la CPU.

Pérdida de la conexión de ssh durante el entrenamiento

Es posible que se agote el tiempo de espera de tu conexión ssh a la Cloud TPU durante un entrenamiento de larga duración (en particular, si usas Cloud Shell). En ese momento, no hay resultados en la consola de TPU y es posible que parezca que la TPU dejó de entrenar. Para evitar esto, ejecuta la sesión de capacitación con un multiplexor de terminal o una herramienta de administración de sesiones, como tmux o screen. De esta manera, la conexión ssh permanecerá activa, independientemente de la duración del entrenamiento.

Depura errores comunes

En esta sección, se describe cómo solucionar errores comunes encontrar cuando se entrenan modelos en Cloud TPU.

No se puede crear una TPU

Cuando crees una Cloud TPU, es posible que veas el siguiente error:

googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."

Este es un problema de permisos y se puede resolver ejecutando el siguiente comando:

gcloud auth login --update-adc

Este comando actualiza las credenciales predeterminadas de la aplicación (ADC) y debería resolver el problema. Para obtener más información, consulta gcloud auth login.

Formas dinámicas no compatibles

Mensaje de error

ValueError: shape [Shape] must have a fixed size for dimension
d that is known at graph construction time.

Frameworks y parámetros de configuración afectados

Este mensaje solo ocurre durante la compilación de XLA con TensorFlow.

Detalles

Para ejecutar un modelo en la TPU, Cloud TPU compila el modelo con el compilador XLA. Mientras que este paso de compilación mejora significativamente la velocidad de entrenamiento y el uso de memoria, las formas (tamaños de las dimensiones) de todos los tensores en el gráfico deben conocerse en el momento de la compilación del grafo. Si no se pueden determinar algunas formas en el momento de la compilación, falla la compilación de la TPU con un error como el que se mostró anteriormente.

Una operación común que muestra una forma dinámica es dataset.batch(batch_size), ya que la cantidad de muestras restantes en una transmisión puede ser menor que el tamaño del lote. Por lo tanto, cuando entrenes en la TPU, configura drop remainder=True para dataset.batch. Esto podría eliminar las últimas muestras de un archivo para garantizar que cada lote tenga una forma estática de batch_size. Por ejemplo:

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

Operación de TensorFlow no disponible

Mensaje de error

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT
devices compatible with node

Frameworks y parámetros de configuración afectados

Este mensaje puede aparecer durante el entrenamiento con TensorFlow.

Detalles

El modelo usa una op de TensorFlow que no está disponible en la TPU.

Para obtener una lista de operaciones disponibles en la TPU, junto con los planes de asistencia y sugerencias para soluciones, consulta la guía para de operaciones de TensorFlow disponibles.

Mensaje de error de memoria insuficiente

Mensaje de error

ResourceExhaustedError: Ran out of memory in memory space hbm; used:
YYY; limit: 7.48G.

Frameworks y parámetros de configuración afectados

Este mensaje puede aparecer durante el entrenamiento con TensorFlow, PyTorch o JAX.

Detalles

Cada Cloud TPU está formada por ocho núcleos de TPU. Las TPU v2 tienen 8 GB y las TPU v3 tienen 16 GB de RAM (o memoria de alto ancho de banda, HBM). Esta memoria se utiliza para almacenar los tensores de peso (variable), así como los tensores de resultado intermedio necesarios en el cálculo de gradientes. Si el modelo es demasiado grande para la RAM de la TPU, falla la inicialización y se imprime el mensaje de error anterior. Consulta la sección sobre cómo reducir uso de memoria para obtener más ayuda.

Sugerencias para reducir el uso de la memoria:

Problemas para detener la ejecución

Si TensorFlow encuentra un error durante la ejecución de la TPU, la secuencia de comandos a veces parece interrumpirse en vez de salir de la shell. Si esto sucede, presiona CTRL+C en el teclado para activar un SIGQUIT, lo que hace que Python salga de inmediato.

De manera similar, si presionas CTRL+C durante la ejecución de la TPU, TensorFlow no se desactiva de inmediato, sino que espera hasta el final del bucle de iteración actual para cerrarse de forma correcta.

Si encuentras errores nuevos al volver a conectarte a la TPU después de salir de esta manera, restablece manualmente el servidor de TPU con los comandos:

gcloud compute tpus tpu-vm stop tpu-name --zone=zone
gcloud compute tpus tpu-vm start tpu-name --zone=zone

donde tpu-name se toma de la primera columna que muestra el El comando gcloud compute tpus tpu-vm list y zone es la zona que se muestra en la segunda columna.

Relleno excesivo del tensor

Causa posible de problema de la memoria

Se rellenan los tensores en la memoria de la TPU, es decir, la TPU redondea los tamaños de los tensores almacenados en la memoria para realizar cálculos de forma más eficaz. Este relleno sucede de manera transparente en el nivel del hardware y no afecta los resultados. No obstante, en ciertos casos el relleno puede provocar un aumento significativo del uso de la memoria y del tiempo de ejecución.

Cómo reducir el uso de la memoria

El software de la TPU intenta distribuir los tensores en la memoria para maximizar la eficiencia del cálculo y minimizar el relleno. El proceso de distribución de la memoria es complejo; sin embargo, para obtener mejores resultados, el modelo debe obedecer la siguiente regla general. Para minimizar la sobrecarga de la memoria y maximizar la eficiencia del cálculo, una de las siguientes condiciones debe ser verdadera:

  • El tamaño total del lote debe ser un múltiplo de 64 (8 por núcleo de la TPU) y las dimensiones de las funciones deben ser un múltiplo de 128.

    o

  • El tamaño total del lote debe ser un múltiplo de 1,024 (128 por núcleo de la TPU) y las dimensiones de las funciones deben ser un múltiplo de 8.

Se logra una mejor eficiencia con un tamaño de lote de 1,024 y dimensiones de las funciones que sean múltiplos de 128, aunque esto puede no ser posible para todos los modelos. Para evitar confusiones, “dimensión de las funciones” se refiere al tamaño oculto de una capa completamente conectada o a la cantidad de canales de salida en una convolución. No todas las capas pueden cumplir con esta regla, especialmente la primera y la última capa de la red. Esto es correcto y se espera que la mayoría de los modelos requieran de cierta cantidad de relleno.

Reduce el uso de la memoria

Si encuentras un error de memoria insuficiente cuando ejecutas tu modelo en la TPU, debes tomar medidas para reducir el uso de memoria del modelo.

Las formas más eficaces de reducir el uso de la memoria son las siguientes:

  • Reduce el relleno excesivo del tensor
  • Reduce el tamaño del lote

El tamaño del lote o el modelo es demasiado grande

Causa posible de problema de la memoria

Cuando entrenamos una red neuronal en una CPU, GPU o TPU, el uso de la memoria proviene de las siguientes dos fuentes:

  1. El uso de la memoria es proporcional a la cantidad de pesos del modelo.
  2. Almacenamiento de activaciones intermedias de la propagación necesaria para calcular la retropropagación. El uso de la memoria es directamente proporcional al tamaño del lote, a los tamaños de las capas y a las cantidades de capas.

Por lo tanto, la memoria requerida por un modelo depende en gran medida del tamaño del lote.

La memoria requerida por un modelo depende de la cantidad de capas la red.

El entorno de ejecución de TPU intenta optimizar los operadores para ajustar el modelo a la memoria (lo que se llama rematerialización, similar al punto de control de las gradientes), pero no siempre puede hacerlo.

Cómo reducir el uso de la memoria

Reduce lentamente el tamaño del lote hasta que se ajuste a la memoria y asegúrate de que el tamaño total del lote es un múltiplo de 64 (el tamaño del lote por núcleo debe ser un múltiplo de 8). Ten en cuenta que los tamaños de lotes más grandes son más eficaces en la TPU. Por lo general, un buen punto de partida es un tamaño total de lote de 1,024 (128 por núcleo).

Si no se puede ejecutar el modelo en la TPU incluso con un tamaño de lote pequeño (por ejemplo, 64), intenta reducir la cantidad de capas o sus tamaños.

Mejora la velocidad de entrenamiento

Si se puede ejecutar correctamente el modelo en la TPU, pero la velocidad de entrenamiento es inferior a la esperada, en esta sección se describen varios métodos que pueden mejorar la velocidad. Consulta la guía de rendimiento para obtener otras sugerencias sobre cómo mejorar el rendimiento del entrenamiento.

Muy pocos pasos por ejecución por bucle de entrenamiento

Descripción del problema de rendimiento

Pasa el argumento steps_per_execution a los controles de Model.compile cuántos pasos de entrenamiento se ejecutan entre devoluciones de llamada de host. Cada devolución de llamada del host requiere una comunicación significativa entre la CPU del host del servidor de TPU y el dispositivo de TPU, por lo que si steps_per_execution es demasiado pequeño, puede ralentizar el entrenamiento.

Cómo saber si tu modelo se ve afectado

Si un perfil de TPU revela devoluciones frecuentes de la CPU del host entre los pasos del dispositivo de TPU, tu entrenamiento puede beneficiarse con un valor de steps_per_execution mayor.

Cómo mitigar el problema

Establece steps_per_execution en un valor mayor. Ten en cuenta que steps_per_execution se puede establecer en un valor alto, pero ten en cuenta los mensajes de registro y guardar un punto de control solo pueden ocurrir después de que una cantidad específica de pasos ejecutados.

Cuello de botella del procesamiento de entrada

Descripción del problema de rendimiento

Mientras la TPU se entrena con un grupo específico de datos, la función de procesamiento de entrada prepara el siguiente grupo de datos en la CPU. Si tu entrada función tarda más que la función del modelo, la TPU se deja inactiva mientras la función de entrada recupera datos.

Cómo saber si tu modelo se ve afectado

Sigue las instrucciones de las Herramientas de Cloud TPU: analizador de canalización de entrada para visualizar el análisis de canalización de entrada en TensorBoard:

image

La página del análisis de canalización de entrada muestra un resumen claro que indica si tu modelo presenta un cuello de botella en el nivel del procesamiento de entrada. La misma página también muestra el tiempo de ejecución por operación, lo que te permite identificar las operaciones problemáticas.

Cómo mitigar el problema

Existen varias mitigaciones posibles cuando se cargan datos con la API de Dataset:

  1. Almacena tus datos como un conjunto de estructuras tf.train.Example en archivos TFRecord y cárgalos con TFRecordDataset. Consulta el instructivo de la API de Dataset o el instructivo de ResNet para ver ejemplos.
  2. Usa dataset.cache() o dataset.prefetch() para almacenar en búfer los datos de entrada. Esto evita que las demoras esporádicas en el acceso a archivos generen un cuello de botella.
  3. Especifica el parámetro num_parallel_calls de la función dataset.map() para habilitar operaciones map() de subprocesos múltiples. Una heurística para el valor de num_parallel_calls usa la cantidad de núcleos de CPU disponibles.
  4. Realiza el procesamiento previo de los datos costosos sin conexión como un costo único, en vez de que se aplique el costo en cada ciclo de cada entrenamiento.

Todo el procesamiento de entrada se realiza en las CPU ubicadas en el servidor de TPU, no en la máquina local; por lo que la velocidad de la máquina local no es un factor.

Tiempos de paso lentos y baja utilización de las MXU

Descripción del problema de rendimiento

Cloud TPU puede realizar convoluciones y multiplicaciones de matrices a velocidades increíblemente rápidas. La mayoría de las otras operaciones de TensorFlow tienen implementaciones eficientes en la TPU, pero no representan el valor principal de la TPU en comparación con otro hardware. En consecuencia, se debería dominar un modelo mediante las convoluciones o multiplicaciones de matrices para aprovechar la TPU al máximo.

Cómo saber si tu modelo se ve afectado

Los síntomas que verás en este caso son tiempos de paso lentos junto con una baja utilización de MXU que se muestra cuando generas perfiles del rendimiento.

Cómo mitigar el problema

Intenta reducir la cantidad de operaciones que no son multiplicaciones de matrices. Vuelve a comparar después de reducir la cantidad de multiplicaciones de matrices para ver si el rendimiento es aceptable en las TPU.

Relleno excesivo del tensor

Descripción del problema de rendimiento

La TPU rellena los tensores en la memoria para poder usar sus unidades de procesamiento eficientemente. El relleno puede aumentar el uso de la memoria y del ancho de banda de la memoria. Consulta la sección sobre el relleno del tensor para comprender y solucionar los problemas de relleno del tensor.

Capacidad de procesamiento lenta y bajo uso de memoria

Descripción del problema de rendimiento

Como regla general, el uso de tamaños de lote más grandes provoca una velocidad de entrenamiento mayor en la TPU, en términos de muestras por segundo.

Cómo saber si tu modelo se ve afectado

El tamaño del lote de cualquier modelo siempre debe ser al menos de 64 (8 por núcleo de la TPU), ya que la TPU siempre rellena los tensores hasta este tamaño. El tamaño de lote ideal durante el entrenamiento en la TPU es de 1,024 (128 por núcleo de la TPU), ya que esto elimina las ineficiencias relacionadas con el relleno y la transferencia de la memoria.

Cómo mitigar el problema

Se recomienda usar el tamaño de lote más grande que se ajuste a la memoria múltiplo de 64. La manera más sencilla de lograrlo es comenzar con 1,024 y, si esto causa un error de memoria insuficiente, intenta reducir el tamaño del lote hasta que el modelo se ejecute correctamente. Cambiar el tamaño del lote de un modelo puede requerir el ajuste de otros hiperparámetros para lograr la misma exactitud del modelo, como la tasa de aprendizaje; pero esto se debe evaluar caso por caso.

Capas demasiado pequeñas

Descripción del problema de rendimiento

Incluso cuando un modelo está dominado por convoluciones o multiplicaciones de matrices, es posible que la TPU no se ejecute con la eficiencia máxima si los tensores de entrada son pequeños. En comparación con otro hardware, la TPU se ejecuta con mayor eficiencia cuando tanto los lotes como las capas son grandes (por ejemplo, de dimensión >= 512).

Cómo saber si tu modelo se ve afectado

Como regla general, los tamaños de capa menores a 128 son poco eficientes en la TPU, ya que esa es la dimensión integrada de la unidad de multiplicación de matrices de la TPU. Se recomienda un tamaño oculto mínimo de 512 para capas completamente conectadas con el fin de obtener una eficiencia mayor. Ten en cuenta que las capas convolucionales en general no necesitan ser tan grandes como las capas completamente conectadas para lograr un nivel de eficacia equivalente.

Cómo mitigar el problema

Si la motivación principal para usar tamaños de capas pequeños en tu modelo es la velocidad de entrenamiento, vuelve a realizar comparativas de tus modelos con capas más grandes en la TPU. Por ejemplo, aumentar el tamaño de salida de una capa de 256 a 512 solo puede aumentar el tiempo de entrenamiento en un 20%, aunque el modelo realice el doble de cálculos.

Perfilado del modelo en el nivel de las operaciones

Generalmente, es útil medir el tiempo de ejecución y el uso de memoria en el nivel de las operaciones para identificar los cuellos de botella del rendimiento. Para obtener instrucciones sobre cómo hacerlo,
consulta la guía Herramientas de Cloud TPU: lector de seguimiento.

Depura degradaciones en la exactitud del modelo

Uno de los objetivos del ecosistema de Cloud TPU es que cualquier modelo que se entrene en una CPU o en una GPU obtenga una exactitud muy similar a la que se logra con el entrenamiento en la TPU, tal vez con menos ajustes a los hiperparámetros como el tamaño del lote y la tasa de aprendizaje. En ciertas ocasiones, sin embargo, los usuarios pueden notar una degradación en la precisión cuando entrenan modelos en la TPU. La depuración de estos problemas puede ser muy frustrante debido a la naturaleza aleatoria del entrenamiento de la red neuronal. En esta sección, se proporciona orientación sobre cómo identificar la causa raíz de cualquier degradación en la precisión del modelo cuando se porta un modelo a la TPU.

Comprende la fragmentación de datos (paralelismo de datos)

Uno de los objetivos principales de TensorFlow es que cada operación produzca resultados casi idénticos, ya sea que se ejecute en la CPU, GPU o TPU. Existen ciertas excepciones, como las operaciones aleatorias. En general, si encuentras una diferencia significativa entre la salida de las operaciones no aleatorias en la TPU y en la CPU, repórtalo como un error.

Sin embargo, para la canalización de entrenamiento en conjunto, hay una diferencia significativa entre el entrenamiento en la CPU/GPU y en la TPU. Cuando entrenas en una TPU, TensorFlow realiza la fragmentación de datos Cada Cloud TPU contiene 8 núcleos de TPU que funcionan como independientes unidades de procesamiento. Para cada paso del entrenamiento, cada núcleo de la TPU recibe un lote de datos, calcula los gradientes de peso, intercambia los gradientes con los otros núcleos de la TPU y, luego, calcula la actualización del peso. De manera predeterminada, la pérdida se promedia de todos los núcleos, pero se puede sumar si cambias el parámetro CrossShardOptimizer.

Si la pérdida total del modelo se puede calcular como el promedio (o el total) de las pérdidas independientes por muestra, este procedimiento equivale de manera matemática al entrenamiento en un lote grande único.

La op más común que no es independiente por muestra es la de lote normalización, que se ejecuta y cada lote por núcleo por separado. Por ejemplo, si el tamaño total del lote es 128, entonces el tamaño del lote por núcleo es de 16, y cada uno de los 8 núcleos realiza por lotes en sus propias 16 muestras. En algunos casos, se descubrió que la normalización de lotes pequeños (por ejemplo, inferiores a 32) provoca degradaciones en la exactitud. En una situación ideal, el tamaño total del lote debe ser grande (por ejemplo, de 256 a 1,024). Si el tamaño del lote es demasiado grande para ajustarse a la memoria, se debe evaluar el efecto de la fragmentación caso por caso.

Entrenamiento determinista

Un motivo por el que es difícil depurar las diferencias en la exactitud del modelo es que, en diferentes frameworks (TensorFlow, PyTorch, JAX), el software de entrenamiento usa una inicialización de peso y una redistribución de datos diferentes cada vez que se entrena un modelo. Es beneficioso modificar el procedimiento de entrenamiento para que sea determinista, para que múltiples ejecuciones produzcan modelos casi idénticos. En esta sección se demuestra cómo ejecutar el instructivo de MNIST de manera determinista:

  1. Genera un archivo de punto de control inicial con la ejecución de un solo paso en la CPU. Este paso se usa para lograr una inicialización de peso determinista. Además, asegúrate de usar un valor inicial aleatorio fijo para cualquier función aleatoria en el modelo.
# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1
  1. Modifica cualquier función de redistribución de datos en tu función de entrada para usar un valor de origen aleatorio. Esto ya se realizó en el instructivo de MNIST. Funciona para las operaciones de procesamiento de datos de entrada porque siempre se ejecutan en la CPU. Las operaciones aleatorias en la función de modelo pueden no ser deterministas entre la TPU y la CPU. Cada vez que llames a una operación aleatoria, pasa un valor inicial fijo para garantizar los mismos resultados entre las ejecuciones. Por ejemplo:
# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)
  1. Ejecuta el mismo modelo dos veces en la CPU para verificar que el entrenamiento sea deterministas. Ten en cuenta que el entrenamiento se debe ejecutar por una cantidad razonable de pasos (por ejemplo, 1,000), pero no es necesario hacerlo hasta llegar a la convergencia.

    Debido a que el entrenamiento en la CPU se compara con uno de núcleo único en la TPU, usa un tamaño de lote que pueda ajustarse a un núcleo de la TPU único (en general, el tamaño de lote completo dividido por 8). TensorFlow no garantiza un determinismo bit a bit entre las ejecuciones, pero la pérdida debe ser muy similar:

Copia los pesos iniciales

gcloud storage cp ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1/ --continue-on-error
gcloud storage cp ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2/ --continue-on-error

Serie 1

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

Salida 1

accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

Ejecución 2

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

Salida 2

accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

Entrenamiento de la TPU de núcleo único

Cuando puedas ejecutar el instructivo de MNIST de forma determinista, el paso siguiente consiste en replicar en la TPU los resultados entrenados en la CPU, con un núcleo único de TPU para identificar si el problema se relaciona con la fragmentación de datos o con el motor de ejecución de la TPU.

Estos son los pasos para ejecutar el entrenamiento de núcleo único y su evaluación con el instructivo de MNIST:

Usa la misma inicialización de peso que la CPU

gcloud storage cp ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output --continue-on-error

Ejecuta el entrenamiento durante 1,000 pasos

python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

Salida

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

La pérdida no coincidirá exactamente con el modelo entrenado en la CPU, pero debe ser similar. Si no es el caso para tu modelo, puede indicar un error en el motor de ejecución de la TPU. Antes de enviar un informe de errores, debes verificar los siguientes puntos:

  1. Estás pasando num_shards=1 a TPUConfig.

  2. No tienes operaciones aleatorias en la función de tu modelo, ni operaciones aleatorias en que tu función de entrada se genere de forma correcta.

  3. Estás usando el mismo archivo de punto de control inicial para el entrenamiento de la TPU y de la CPU.

Depura el entrenamiento de la TPU de varios núcleos

Si tu modelo alcanza efectivamente la misma pérdida en la CPU y en la TPU de núcleo único, es probable que el problema sea uno de los siguientes:

(a) La degradación se debe a la variación aleatoria natural cuando se entrenan modelos neuronales con diferentes inicializaciones.

(b) La degradación se debe a un problema relacionado con la fragmentación de datos en la TPU.

Para determinar si (a) es el problema, vuelve a entrenar el modelo completo en la CPU/GPU y en la TPU de varios núcleos con la misma inicialización de peso.

Si estás seguro de que la degradación en la precisión tiene importancia estadística, los problemas más probables relacionados con la fragmentación de datos son los siguientes:

  1. Si tu modelo usa la normalización por lotes, un tamaño de lote total inferior a 256 (por ejemplo, inferior a 32 por lote) puede reducir la exactitud.
  2. La fragmentación afecta a las funciones de pérdida por lotes. Estas funciones de pérdida suelen ser bastante especializados. Por ejemplo, Karras et al. 2017 usa un discriminante por lotes cuando se entrena una red generativa adversaria.

Solución de problemas de configuración de gcloud

Problema
gcloud components update muestra el siguiente mensaje de error:
ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager is
disabled for this installation.
Solución
Para usar gcloud, debes usar una instalación de gcloud que no se administre a través de un administrador de paquetes. Sigue estos pasos para instala gcloud desde el código fuente:
  sudo apt-get remove google-cloud-sdk
  curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  ./google-cloud-sdk/install.sh
  source ~/.bashrc
Problema

Comando gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} muestra el siguiente mensaje de error:

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
Solución

Es posible que haya un error en la propagación de la clave SSH. Intenta mover las claves generadas automáticamente a una ubicación de copia de seguridad para forzar a gcloud a volver a crearlas:

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

Registros de depuración

Los frameworks de Cloud TPU compatibles, JAX, PyTorch y TensorFlow acceder a las TPU con una biblioteca compartida llamada libtpu que está presente en cada VM de TPU. Esta biblioteca incluye el compilador XLA que se usa para compilar los programas de TPU, el tiempo de ejecución de TPU que se usa para ejecutar programas compilados. y el controlador de TPU que usa el entorno de ejecución para obtener acceso de bajo nivel a la TPU.

La biblioteca libtpu registra información que puede ser útil para la depuración. De forma predeterminada, estos registros se escriben en /tmp/tpu_logs en cada VM de Cloud TPU. Puedes configurar las siguientes variables de entorno antes de comenzar el entrenamiento para modificar el comportamiento de registro:

TPU_LOG_DIR: Es el directorio en el que se escriben los registros.
La ubicación del directorio es /tmp/tpu_logs de forma predeterminada. El directorio se crea si aún no existe, pero no se crean directorios superiores. Si se produce un error al encontrar o crear el directorio especificado, se imprime un mensaje en stderr, pero no se detendrá el programa y se inhabilitará el registro. Establece el nombre del directorio en "disabled" para inhabilitar el registro en el disco por completo.
TPU_MIN_LOG_LEVEL: la gravedad mínima que se registrará en el disco
Las opciones son 0 (INFORMACIÓN), 1 (ADVERTENCIA), 2 (ERROR) y 3 (FATAL). El valor predeterminado es 0.
TPU_STDERR_LOG_LEVEL: Es la gravedad mínima que se registrará en stderr, además del disco, si corresponde.
Las opciones son las mismas que las de TPU_MIN_LOG_LEVEL. El valor predeterminado es 3.
TPU_MAX_LOG_SIZE_MB: Es el tamaño máximo en megabytes de cada archivo de registro.
Se iniciará automáticamente un nuevo archivo de registro cuando el anterior alcance aproximadamente este tamaño. La configuración predeterminada es 1,024.