Solución de problemas de TensorFlow - TPU

En esta guía, junto con las Preguntas frecuentes, se proporciona ayuda para la solución de problemas para los usuarios que entrenan modelos de TensorFlow en Cloud TPU. Si quieres solucionar problemas de entrenamiento de Pytorch o JAX, puedes consultar los documentos de solución de problemas de esos frameworks:

• Solución de problemas de PyTorch: TPU

• Solución de problemas de JAX: TPU

  Importante: Cloud TPU es compatible con dos arquitecturas de TPU: la VM de TPU o el nodo TPU. Las dos arquitecturas de VM se describen en Arquitectura del sistema. Los procedimientos de solución de problemas de este documento se aplican a ambas configuraciones, pero también hay una sección independiente de VM de TPU, específicamente para solucionar problemas que solo ocurren con las configuraciones de VM de TPU.

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

• Usa TPU de Colab
• Guías de inicio rápido de Cloud TPU
• Instructivo de MNIST
• Entrena modelos de AA en unidades de procesamiento tensorial (TPU) de Cloud.

Nota: Si no puedes resolver el problema con esta guía, consulta Obtén asistencia para obtener más ayuda.

Descripción general

Los problemas comunes que se encuentran con las Cloud TPU se clasifican en las siguientes categorías:

1. Problemas de conexión a la TPU

2. Cómo depurar errores comunes

3. Cómo reducir el uso de memoria

4. Mejora la velocidad de entrenamiento

5. Cómo depurar disminuciones en la exactitud del modelo

Problemas para establecer la conexión al servidor de TPU

En esta sección, se describe cómo solucionar situaciones en las que TensorFlow deja de responder o imprime un error cuando se conecta a la TPU. El paso de compilación del grafo de TPU puede llevar mucho tiempo para los modelos grandes, por lo que debes dejar que la secuencia de comandos se ejecute durante al menos 5 minutos antes de concluir que 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 aún hay problemas de conexión con el instructivo de MNIST, esto confirma que se trata de un problema con el servidor de TPU. En este caso, haz lo siguiente:

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

   (vm)$ gcloud compute tpus 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 el siguiente comando:

   (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. Este proceso puede tardar varios minutos en completarse.

4. Vuelve a ejecutar el comando ... tpus list anterior 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 de forma correcta, pero tu modelo aún deja de responder, es probable que el problema sea de tu canalización de entrenamiento. Para depurar esto, comienza por reemplazar TPUStrategy en tu código por la estrategia predeterminada. Cuando usas la estrategia predeterminada, siempre que uses strategy.scope() o strategy.run(), el modelo se ejecuta en la CPU (o GPU, si está presente) en lugar de en la TPU. Si el modelo se ejecuta en CPU y no en TPU, debe haber un problema específico con la TPU. Si aún no se ejecuta, la práctica recomendada es 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 Cloud TPU durante un entrenamiento de larga duración (en especial si usas Cloud Shell). En ese momento, no hay resultados en la consola de la TPU y puede parecer que la TPU dejó de entrenarse. Para evitar esto, ejecuta la sesión de entrenamiento con un multiplexor de terminal o una herramienta de administración de sesiones como tmux o screen. Esto mantendrá activa la conexión ssh sin importar la duración del entrenamiento.

Depura errores comunes

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 con el siguiente comando:

gcloud auth login --update-adc

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

No se puede usar el sistema de archivos local

Mensaje de error

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

Frameworks y configuraciones afectados

Este mensaje puede aparecer durante el entrenamiento con TensorFlow mediante la arquitectura de nodo TPU.

Detalles

Todos los archivos de entrada y el directorio del modelo deben usar una ruta de bucket de almacenamiento en la nube (gs://bucket-name/...), y este bucket debe ser accesible desde el servidor de TPU. Ten en cuenta que el punto de control de modelos y el procesamiento de datos se realizan en el servidor de TPU, no en la máquina local. Si deseas obtener información sobre cómo configurar el almacenamiento en la nube de forma correcta para su uso con la TPU, consulta la guía Conéctate a depósitos de Cloud Storage.

Tipo de datos no compatible

Mensaje de error

TypeError: DataType is not a supported TPU infeed type.

Frameworks y configuraciones afectados

Este mensaje puede aparecer durante el entrenamiento con TensorFlow mediante la arquitectura de nodo TPU.

Detalles

Hoy en día, solo los tipos de datos tf.float32, tf.int32, tf.bfloat16 y tf.bool son compatibles con la TPU. Otros tipos de datos comunes, como tf.uint8, tf.string y tf.int64, deben convertirse en uno de los tipos de datos admitidos durante el procesamiento previo de los datos (es decir, en la canalización tf.data.Dataset).

Consulta un ejemplo de la conversión en la función decode_image que se usa en el entrenamiento de MNIST.

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, TensorFlow compila el modelo con el compilador de XLA. Si bien este paso de compilación mejora de manera significativa la velocidad de entrenamiento y el uso de memoria, en el momento de la compilación del grafo, se deben conocer las formas (tamaños de dimensión) de todos los tensores 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 anterior.

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 elimina potencialmente 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 operación de TensorFlow que no se encuentra actualmente disponible en la TPU.

Para obtener una lista de operaciones disponibles en la TPU, junto con los planes de asistencia futura y sugerencias para una solución alternativa, consulta la guía 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, las TPU v2 tienen 8 GB y las TPU v3 tienen 16 GB de RAM (o HBM, memoria de alto ancho de banda). 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 el uso de memoria para obtener más ayuda.

Sugerencias para reducir el uso de la memoria:

• Comprueba si hay exceso relleno en el tensor.
• Usa el formato bfloat16
• Si los tamaños o el modelo de entrada son demasiado grandes, es posible que puedas usar el paralelismo del modelo experimental de TensorFlow para abordar el tamaño del modelo.

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+\ en el teclado para activar un SIGQUIT, lo que hace que Python salga de inmediato.

Del mismo modo, cuando se presiona CTRL+C durante la ejecución de la TPU, TensorFlow no se cierra de inmediato, sino que se espera hasta el final del bucle de la iteración actual para salir de forma correcta.

Si encuentras algún error nuevo cuando te vuelves a conectar a la TPU después de salir de esta manera, restablece el servidor de TPU de forma manual con los siguientes comandos:

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

En el ejemplo anterior, tpu-name se toma de la primera columna que muestra 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 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 que requiere un modelo depende de la cantidad de capas en la red.

El entorno de ejecución de TPU intenta optimizar los operadores para ajustar el modelo a la memoria (llamada rematerialización, similar al 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 sea 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

Pasar el argumento steps_per_execution a Model.compile controla cuántos pasos de entrenamiento se ejecutan entre las devoluciones de llamada del host. Cada devolución de llamada de 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 de llamada frecuentes de la CPU del host entre los pasos del dispositivo de la TPU, el entrenamiento puede beneficiarse de 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 grande, pero ten en cuenta que el registro de mensajes y el guardado de un punto de control solo pueden ocurrir después de que se haya ejecutado la cantidad especificada de pasos.

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 la función de entrada tarda más que la función de modelo, la TPU se deja inactiva mientras la función de entrada recupera los 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:

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, 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 simple para el valor de num_parallel_calls es usar 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 tanto, la velocidad de la máquina local no es un factor.

Tiempos de pasos lentos y uso bajo de 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 que verás son tiempos de pasos lentos junto con el bajo uso 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. Después de reducir la cantidad de multiplicaciones de matrices, vuelve a evaluar 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 de memoria bajo

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

La práctica recomendada es usar el tamaño de lote más grande que se ajuste a la memoria y sea un 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 nativa 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 los tamaños de capa pequeños en tu modelo es la velocidad de entrenamiento, vuelve a comparar 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 esté realizando 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: visualizador de seguimiento.

Depura degradaciones en la exactitud del modelo

Uno de los objetivos del ecosistema de Cloud TPU es que cualquier modelo que esté en entrenamiento 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 su conjunto, hay una diferencia significativa entre el entrenamiento en la CPU/GPU y en la TPU. Cuando entrenas en una TPU, TensorFlow realiza fragmentación de datos. Cada Cloud TPU contiene 8 núcleos de TPU que operan como unidades de procesamiento independientes. Para cada paso del entrenamiento, cada núcleo de TPU recibe un lote de datos, calcula los gradientes de peso, los intercambia con los otros núcleos de TPU y, luego, calcula la actualización del peso. De forma predeterminada, la pérdida se promedia entre los núcleos, 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 operación más común que no es independiente por muestra es la normalización por lotes, que se ejecuta en cada lote por núcleo de forma separada. Por ejemplo, si el tamaño total del lote es 128, el tamaño del lote por núcleo es 16, y cada uno de los 8 núcleos realiza la normalización 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) disminuye la exactitud. En la 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 caber en la memoria, el efecto de la fragmentación debe evaluarse caso por caso.

Entrenamiento determinista

Una de las razones por las que es difícil depurar las diferencias en la precisión del modelo es que, entre distintos frameworks (TensorFlow, PyTorch y JAX), el software de entrenamiento usa diferentes inicializaciones de peso y redistribución de datos cada vez que se entrena un modelo. Es conveniente modificar el procedimiento de entrenamiento para que sea determinista, de modo que varias ejecuciones produzcan modelos casi idénticos. En esta sección, se muestra 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 fijo aleatorio 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

2. 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 op, 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)

3. Ejecuta el mismo modelo dos veces en la CPU para verificar que el entrenamiento sea determinista. Ten en cuenta que el entrenamiento se debe ejecutar por una cantidad razonable de pasos (por ejemplo, 1,000), pero no es necesario que se ejecute 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 las ponderaciones 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

Resultado 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

Resultado 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

Ejecutar el entrenamiento de 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 modelo y cualquier operación aleatoria en tu función de entrada se inicializa 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, como se indicó anteriormente.

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. Las funciones de pérdida por lotes se ven afectadas por la fragmentación. Estas funciones de pérdida suelen ser bastante especializadas. Por ejemplo, Karras et al. 2017 usa un discriminante por lotes cuando se entrena una red generativa adversaria.

Solución de problemas de VMs de TPU

Los siguientes problemas y soluciones solo se aplican a la configuración de VM de TPU.

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 con la VM de TPU, deberás usar una instalación de gcloud que no se administre a través de un administrador de paquetes. Sigue estos pasos para instalar 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

El 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 las claves generadas de forma automática 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 acceden a las TPU a través de una biblioteca compartida llamada libtpu que está presente en cada VM de TPU. Esta biblioteca incluye el compilador XLA que se usa para compilar 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 acceder 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 del directorio predeterminada es /tmp/tpu_logs. El directorio se crea si aún no existe, pero no se crean directorios superiores. Si se produce un error para encontrar o crear el directorio especificado, se mostrará un mensaje en stderr, pero no se detendrá el programa y se inhabilitará el registro. Establece el nombre del directorio como “inhabilitado” 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 (INFO), 1 (WARNING), 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

Un archivo de registro nuevo se iniciará automáticamente cuando el anterior alcance este tamaño. La configuración predeterminada es 1,024.