Solución de problemas de TensorFlow: TPU

En esta guía, junto con las Preguntas frecuentes proporciona ayuda para la solución de problemas a los usuarios que están entrenando 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 comunes que se encuentran con las Cloud TPU se clasifican 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 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 comando Instructivo de MNIST con la URL del servidor de TPU y verifica que funcione correctamente. Si aún hay de conexión con el instructivo de MNIST, esto confirma que es un problema con el servidor de TPU. En este caso, ocurre lo siguiente:

  1. Para ver una lista de las TPU disponibles, ejecuta el siguiente comando. Reemplazando zone y project-id por tu zona y el ID del proyecto.

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

    Esto imprime un resultado como el 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 en el ejemplo anterior) y que esta TPU aparece como READY.

  3. Si tu TPU no aparece como READY o sigues teniendo problemas conectando, reinicia manualmente el servidor con:

    (vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus 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 el 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 una Problema específico de TPU. Si aún no se ejecuta, la práctica recomendada es depurar el problema en la CPU.

Se perdió la conexión de ssh durante el entrenamiento

Es posible que se agote el tiempo de espera de tu conexión ssh a Cloud TPU durante un entrenamiento de larga duración (especialmente si usas Cloud Shell). En ese momento, no hay salida a la consola de TPU y podría aparecerá como si la TPU dejara de entrenarse. Para evitar esto, ejecuta de capacitación con un multiplexor de terminal o una herramienta de administración de sesiones, como como tmux o screen. De esta manera, se mantendrá la ssh. conexión activa sin importar 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 creas 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 configuraciones afectados

Este mensaje solo aparece 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 puede determinar alguna forma en el tiempo de compilación, fallará la compilación de la TPU. con un error como el que se mostró anteriormente.

Una op 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 por dataset.batch. Esta acción podría descartar 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 configuraciones 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 configuraciones 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, y 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 adaptarse a la RAM de la TPU, la inicialización falla y se imprime el mensaje de error. 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 una SIGQUIT, lo que causa Python se cierre de inmediato.

Del mismo modo, presionar CTRL+C durante la ejecución de TPU no cierra TensorFlow. de inmediato, pero espera hasta el final del bucle de iteración actual para salir de forma limpia.

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 stop tpu-name --zone=zone
gcloud compute tpus start tpu-name --zone=zone

donde tpu-name se toma de la primera columna que muestra el El comando gcloud compute tpus 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 memoria son las siguientes:

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

Tamaño del lote o modelo 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 memoria es proporcional a la cantidad de los pesos en el 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 en la memoria (llamada rematerialización, similar a gradiente controles de seguridad), 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 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 las 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

En este caso, los síntomas son tiempos de pasos lentos, junto con el bajo uso de MXU se muestra cuando generas un perfil 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 uso bajo 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 inferiores a 128 generan una eficiencia deficiente. en la TPU, ya que 128 es la dimensión integrada de la multiplicación de matrices de la TPU unidad. 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 generalmente no deben ser tan grandes como capas completamente conectadas para lograr una eficiencia equivalente a nivel de organización.

Cómo mitigar el problema

Si la motivación principal de los tamaños de capas pequeños en el modelo es la velocidad de entrenamiento volver a comparar los 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%, a pesar de que el rendimiento del modelo es el doble del 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: Visualizador de seguimiento.

Depura degradaciones en la exactitud del modelo

Uno de los objetivos del ecosistema de Cloud TPU es que cualquier modelo que se entrena con una CPU o GPU alcanzan una exactitud muy similar cuando se entrena en la TPU, tal vez con pequeños ajustes en 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 su conjunto, hay una diferencia diferencia entre el entrenamiento en la CPU/GPU y 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. Por cada paso del entrenamiento, cada núcleo de TPU recibe un lote de datos, calcula los gradientes de peso, los intercambia con los de otros núcleos de TPU, y, luego, calcula la actualización de peso. Por defecto, la pérdida se promedia entre pero se puede sumar si cambias el parámetro de 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, la ejecución por lotes se ha observado que la normalización en lotes pequeños (por ejemplo, menos de 32) disminuyen la exactitud. En una situación ideal, el tamaño total del lote debería ser grande (por ejemplo, de 256 a 1,024). Si un tamaño del lote es demasiado grande para caber memoria, el efecto de la fragmentación debe evaluarse caso por caso.

Entrenamiento determinista

Una razón por la que es difícil depurar diferencias en la exactitud del modelo es que, entre 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. Esta sección 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 una op aleatoria, pasa un valor inicial fijo para garantizar los mismos resultados entre 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 son deterministas. Ten en cuenta que el entrenamiento debe ejecutarse durante una cantidad razonable de pasos (por ejemplo, 1,000), pero no es necesario ejecutarlo hasta 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

gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

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

Serie 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

gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

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 debería 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 una 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. de 2017 usa un discriminante de lotes cuando entrena una red generativa adversaria.

Solución de problemas relacionados con la 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 utilizar una gcloud instalación que no se gestiona 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 problema con la propagación de la clave SSH. Intenta mover el claves generadas automáticamente para una ubicación de copia de seguridad para forzar la recreación de gcloud con ellos:

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 mediante 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. Las siguientes variables de entorno se pueden configurar antes de comenzar el entrenamiento para modificar el comportamiento de registro:

TPU_LOG_DIR: El directorio en el que se escriben los registros
La ubicación predeterminada del directorio es /tmp/tpu_logs. El directorio es si aún no existe, pero no se crean directorios superiores. Si se produce un error cuando se busca o crea el directorio especificado, Se imprime un mensaje en stderr, pero no detendrá el programa y el registro está inhabilitado. Establece el nombre del directorio como "inhabilitado" a inhabilitar por completo el registro en el disco.
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: La gravedad mínima que se registrará en stderr, además del disco, si corresponde.
Las opciones son las mismas que TPU_MIN_LOG_LEVEL. El valor predeterminado es 3.
TPU_MAX_LOG_SIZE_MB: el tamaño máximo en megabytes de cada archivo de registro
Se iniciará automáticamente un nuevo archivo de registro cuando alcance el anterior más o menos de este tamaño. La configuración predeterminada es 1,024.