Guía avanzada de Inception v3 para Cloud TPU

En este documento, se analizan los aspectos del modelo Inception y cómo se relacionan para lograr que este se ejecute de manera eficiente en Cloud TPU. Es una vista avanzada de la guía para ejecutar Inception v3 en Cloud TPU. Los cambios específicos en el modelo que dieron lugar a mejoras significativas se analizan detalladamente. Este documento es un complemento del instructivo de Inception v3.

Las ejecuciones de entrenamiento de la TPU de Inception v3 coinciden con las curvas de exactitud que producen los trabajos de GPU de configuración similar. El modelo se entrenó correctamente con las configuraciones v2-8, v2-128 y v2-512. Alcanzó una exactitud superior al 78.1% en, aproximadamente, 170 ciclos de entrenamiento en cada una de esas configuraciones.

Los ejemplos de código que se muestran en este documento son ilustrativos, es decir que, representan una imagen de alto nivel de lo que sucede en la implementación real. El código de trabajo se puede encontrar en GitHub.

Introducción

Inception v3 es un modelo de reconocimiento de imágenes muy utilizado, el cual se demostró que alcanza una exactitud superior al 78.1% en el conjunto de datos de ImageNet. El modelo representa la culminación de muchas ideas que desarrollaron varios investigadores durante años. Se basa en el documento original: Reformulación de la arquitectura de Inception para la visión artificial de Szegedy y otros.

El modelo está formado por bloques de construcción simétricos y asimétricos que incluyen convoluciones, reducción promedio, reducción máxima, concatenaciones, retirados y capas completamente conectadas. La normalización por lotes se usa con frecuencia en todo el modelo y se aplica a las entradas de activación. Las pérdidas se calculan a través de Softmax.

A continuación, se muestra un diagrama de alto nivel del modelo:

image

El modelo Inception README tiene más información sobre la arquitectura de Inception.

API de Estimator

La versión de la TPU de Inception v3 se escribe con TPUEstimator, una API diseñada para facilitar el desarrollo y ayudarte a que te enfoques en los modelos en lugar de en los detalles del hardware subyacente. La API realiza la mayor parte del trabajo arduo de bajo nivel necesario para ejecutar modelos en las TPU en segundo plano, al tiempo que automatiza funciones comunes, como guardar y restablecer puntos de control.

La API de Estimator aplica la separación del modelo y las partes de entrada del código. Debes definir las funciones model_fn y input_fn, que corresponden a la definición del modelo y a las etapas de la canalización de entrada o procesamiento previo del grafo de TensorFlow, respectivamente. A continuación, se proporciona un ejemplo de la estructura de estas funciones:

def model_fn(features, labels, mode, params):
     …
  return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)

def input_fn(params):
    def parser(serialized_example):
          …
        return image, label

          …
   images, labels = dataset.make_one_shot_iterator().get_next()
   return images, labels

La API proporciona dos funciones clave, train() y evaluate(), las cuales se utilizan para entrenar y evaluar, respectivamente. Estas funciones se suelen llamar desde alguna parte de la función principal. A continuación, se brinda un ejemplo del proceso:

def main(unused_argv):
  …
  run_config = tpu_config.RunConfig(
      master=FLAGS.master,
      model_dir=FLAGS.model_dir,
      session_config=tf.ConfigProto(
          allow_soft_placement=True, log_device_placement=True),
      tpu_config=tpu_config.TPUConfig(FLAGS.iterations, FLAGS.num_shards),)

  estimator = tpu_estimator.TPUEstimator(
      model_fn=model_fn,
      use_tpu=FLAGS.use_tpu,
      train_batch_size=FLAGS.batch_size,
      eval_batch_size=FLAGS.batch_size,
      config=run_config)

  estimator.train(input_fn=input_fn, max_steps=FLAGS.train_steps)

  eval_results = inception_classifier.evaluate(
      input_fn=imagenet_eval.input_fn, steps=eval_steps)

Conjunto de datos de ImageNet

Se debe entrenar al modelo antes de utilizarlo para reconocer imágenes. Generalmente, el entrenamiento se realiza a través del aprendizaje supervisado con un conjunto extenso de imágenes etiquetadas. Aunque Inception v3 se puede entrenar a partir de conjuntos de imágenes etiquetadas diferentes, ImageNet es un conjunto de datos común a elección.

ImageNet tiene más de diez millones de URL de imágenes etiquetadas. Cerca de un millón de las imágenes también incluyen cuadros de límites que especifican una ubicación más precisa para los objetos etiquetados.

Con respecto a este modelo, el conjunto de datos de ImageNet está compuesto por 1,331,167 imágenes, las cuales se dividen en conjuntos de datos de entrenamiento y evaluación que contienen 1,281,167 y 50,000 imágenes, respectivamente.

Los conjuntos de datos de entrenamiento y evaluación se mantienen separados intencionalmente. Para entrenar el modelo, solo se utilizan las imágenes del conjunto de datos de entrenamiento, mientras que las imágenes del conjunto de datos de evaluación se utilizan únicamente con el fin de evaluar la exactitud del modelo.

El modelo espera que las imágenes se almacenen como TFRecords. Para convertir las imágenes de los archivos JPEG sin procesar en TFRecords, utiliza esta secuencia de comandos por lotes de código abierto: download_and_preprocess_imagenet.sh. La secuencia de comandos debería producir una serie de archivos (para el entrenamiento y la validación) con la siguiente forma:

${DATA_DIR}/train-00000-of-01024
${DATA_DIR}/train-00001-of-01024
 ...
${DATA_DIR}/train-01023-of-01024
and
${DATA_DIR}/validation-00000-of-00128
S{DATA_DIR}/validation-00001-of-00128
 ...
${DATA_DIR}/validation-00127-of-00128

DATA_DIR representa la ubicación en la cual reside el conjunto, por ejemplo: DATA_DIR=$ HOME/imagenet-data.

La sección de introducción del documento Modelo Inception README incluye instrucciones detalladas sobre cómo compilar y ejecutar esta secuencia de comandos.

Canalización de entrada

Cada dispositivo Cloud TPU tiene 8 núcleos y está conectado a un host (CPU). Las partes de mayor tamaño contienen varios hosts. Otras configuraciones más amplias interactúan con hosts diferentes. Por ejemplo, v2-256 se comunica con 16 hosts.

Los hosts recuperan los datos del sistema de archivos o de la memoria local, llevan a cabo el procesamiento previo necesario de los datos y, luego, transfieren los datos preprocesados a los núcleos de la TPU. Analizamos estas tres fases del control de datos que realiza el host a nivel individual y nos referimos a esas fases de la siguiente manera: 1) Almacenamiento, 2) Procesamiento previo y 3) Transferencia. En la siguiente representación, se proporciona una imagen de alto nivel del diagrama:

image

Para obtener un rendimiento adecuado, el sistema debe estar balanceado. Independientemente del tiempo que invierte una CPU del host en recuperar imágenes, decodificarlas y llevar a cabo el procesamiento previo relevante, idealmente, debería tardar un tiempo menor o casi igual al que tarda la TPU en realizar los cálculos. Si la CPU del host tarda más que la TPU en completar las tres fases de control de datos, la ejecución estará vinculada al host. (Nota: Debido a que las TPU son tan rápidas, es inevitable que ocurra esto en algunos modelos muy simples). Ambos casos se presentan en el siguiente diagrama.

image

La implementación actual de Inception v3 se encuentra justo en el límite de estar vinculada a la entrada. Las imágenes deben recuperarse del sistema de archivos, decodificarse y se debe realizar el procesamiento previo. Hay diferentes clases de etapas de procesamiento previo disponibles, que varían desde moderadas hasta complejas. Si usamos las etapas de procesamiento previo más complejas, la gran cantidad de operaciones costosas que ejecuta la etapa de procesamiento previo expondrá el sistema al límite, y la canalización de entrada se vinculará al procesamiento previo. Sin embargo, no es necesario recurrir a ese nivel de complejidad para alcanzar una exactitud superior al 78.1%. En cambio, usamos una etapa de procesamiento previo moderadamente compleja que inclina la balanza en la otra dirección y mantiene el modelo vinculado a la TPU. Esto se analiza con más detalle en la siguiente sección.

El modelo usa tf.data.Dataset para controlar todas las necesidades relacionadas con la canalización de entrada. Consulta la guía de rendimiento de los conjuntos de datos para obtener más información sobre cómo optimizar las canalizaciones de entrada con la API tf.data.

Aunque simplemente puedes definir una función y pasarla a la API de Estimator, en el caso de Inception, la creación de la clase InputPipeline encapsula la funcionalidad requerida y define un método __call__.

La API de Estimator facilita mucho el uso de esta clase. Uno simplemente tiene que pasarlo al parámetro input_fn de las funciones train() y evaluate(), como se muestra en el siguiente fragmento de código de muestra:

def main(unused_argv):

          …

  inception_classifier = tpu_estimator.TPUEstimator(
      model_fn=inception_model_fn,
      use_tpu=FLAGS.use_tpu,
      config=run_config,
      params=params,
      train_batch_size=FLAGS.train_batch_size,
      eval_batch_size=eval_batch_size,
      batch_axis=(batch_axis, 0))

          …

  for cycle in range(FLAGS.train_steps // FLAGS.train_steps_per_eval):
    tf.logging.info('Starting training cycle %d.' % cycle)
    inception_classifier.train(
        input_fn=InputPipeline(True), steps=FLAGS.train_steps_per_eval)

    tf.logging.info('Starting evaluation cycle %d .' % cycle)
    eval_results = inception_classifier.evaluate(
        input_fn=InputPipeline(False), steps=eval_steps, hooks=eval_hooks)
    tf.logging.info('Evaluation results: %s' % eval_results)

Los elementos principales de la clase InputPipeline se muestran en el siguiente fragmento de código, en el cual destacamos las tres fases con diferentes colores. El método __call__ crea el conjunto de datos con tf.data.Dataset y, luego, realiza varias llamadas a la API para utilizar las capacidades integradas de captación previa, intercalación y redistribución del conjunto de datos.

class InputPipeline(object):

  def __init__(self, is_training):
    self.is_training = is_training

  def __call__(self, params):
    # Storage
    file_pattern = os.path.join(
        FLAGS.data_dir, 'train-*' if self.is_training else 'validation-*')
    dataset = tf.data.Dataset.list_files(file_pattern)
    if self.is_training and FLAGS.initial_shuffle_buffer_size > 0:
      dataset = dataset.shuffle(
          buffer_size=FLAGS.initial_shuffle_buffer_size)
    if self.is_training:
      dataset = dataset.repeat()

    def prefetch_dataset(filename):
      dataset = tf.data.TFRecordDataset(
          filename, buffer_size=FLAGS.prefetch_dataset_buffer_size)
      return dataset

    dataset = dataset.apply(
        tf.contrib.data.parallel_interleave(
            prefetch_dataset,
            cycle_length=FLAGS.num_files_infeed,
            sloppy=True))
    if FLAGS.followup_shuffle_buffer_size > 0:
      dataset = dataset.shuffle(
          buffer_size=FLAGS.followup_shuffle_buffer_size)

    # Preprocessing
    dataset = dataset.map(
        self.dataset_parser,
        num_parallel_calls=FLAGS.num_parallel_calls)

    dataset = dataset.prefetch(batch_size)
    dataset = dataset.apply(
        tf.contrib.data.batch_and_drop_remainder(batch_size))
    dataset = dataset.prefetch(2)  # Prefetch overlaps in-feed with training
    images, labels = dataset.make_one_shot_iterator().get_next()

    # Transfer
    return images, labels

La sección storage comienza con la creación de un conjunto de datos que incluye la lectura de TFRecords desde el almacenamiento (con tf.data.TFRecordDataset). Las funciones con un propósito especial, repeat() y shuffle(), se usan según sea necesario. La función tf.contrib.data.parallel_interleave() asigna la función prefetch_dataset() a través de su entrada para producir conjuntos de datos anidados y muestra sus elementos intercalados. Obtiene elementos de los conjuntos de datos anidados de cycle_length en paralelo, lo que aumenta la capacidad de procesamiento. El argumento sloppy flexibiliza el requisito de que los resultados se produzcan en un orden determinista y permite que la implementación omita los conjuntos de datos anidados cuyos elementos no estén disponibles cuando se soliciten.

La sección procesamiento previo llama a dataset.map(parser), que a su vez llama a la función del analizador en la que se procesan las imágenes. Los detalles de la etapa de procesamiento previo se analizan en la siguiente sección.

En la sección transferencia (al final de la función), se incluye la línea return images, labels. TPUEstimator toma los valores mostrados y los transfiere automáticamente al dispositivo.

En la siguiente figura, se muestra un ejemplo de seguimiento del rendimiento de Cloud TPU de Inception v3. El tiempo de procesamiento de TPU, si se descuentan las paradas de entrada, se encuentra actualmente en 815 ms o menos.

image

A continuación, se muestra el almacenamiento del host que también se observa en el seguimiento.

image

A continuación, se muestra el procesamiento previo del host, que incluye la decodificación de imágenes y una serie de funciones de distorsión de imagen.

image

La transferencia de host/TPU se puede ver aquí:

image

Etapa de procesamiento previo

El procesamiento previo de imágenes es una parte crucial del sistema y puede influir, en gran medida, en la exactitud máxima que el modelo alcanza durante el entrenamiento. Como mínimo, se deben decodificar las imágenes y cambiar el tamaño para adaptarlas al modelo. En el caso de Inception, las imágenes deben ser de 299 x 299 x 3 píxeles.

Sin embargo, la decodificación y el cambio de tamaño no serán suficientes para lograr una exactitud adecuada. El conjunto de datos de entrenamiento de ImageNet contiene 1,281,167 imágenes. Un ciclo de entrenamiento se define como un pase a través del conjunto de imágenes de entrenamiento. Durante el entrenamiento, el modelo requerirá varios pases a través del conjunto de datos de entrenamiento para mejorar las capacidades de reconocimiento de imágenes. En el caso de Inception v3, la cantidad necesaria de ciclos de entrenamiento estará en un rango de 140 a 200, según el tamaño general del lote.

Es beneficioso alterar continuamente las imágenes antes de enviarlas al modelo y, para eso, hay que hacerlo de manera tal que una imagen en particular sea ligeramente diferente en cada ciclo de entrenamiento. La mejor manera de realizar este procesamiento previo de imágenes no solo es una ciencia, sino también un arte. Por un lado, una etapa de procesamiento previo bien diseñada puede aumentar las capacidades de reconocimiento de un modelo de forma significativa. Por otro lado, una etapa de procesamiento previo demasiado simple puede establecer un límite artificial sobre la exactitud máxima que ese modelo es capaz de alcanzar durante el entrenamiento.

Inception v3 ofrece diferentes opciones para la etapa de procesamiento previo, que varían desde relativamente simples y económicas en términos de procesamiento hasta muy complejas y costosas desde el punto de vista de los recursos informáticos. Se pueden encontrar dos tipos distintos de opciones en los archivos inception_preprocessing.py y vgg_preprocessing.py.

El archivo vgg_preprocessing.py define una etapa de procesamiento previo que se utilizó con éxito para entrenar resnet con una exactitud del 75%, pero genera resultados deficientes cuando se aplica en Inception v3.

El archivo inception_preprocessing.py contiene una etapa de procesamiento previo de varias opciones con diferentes niveles de complejidad, que se empleó satisfactoriamente para entrenar a Inception v3 con una exactitud del 78.1% al 78.5% cuando se ejecutó en la TPU. En esta sección se analiza la canalización de procesamiento previo.

Las variaciones en el procesamiento previo se basan en si el modelo se está entrenando o se está utilizando para realizar una inferencia o evaluación.

En el momento de la evaluación, el procesamiento previo es bastante sencillo: recorta una región central de la imagen y, luego, cambia el tamaño a la medida predeterminada de 299 x 299. A continuación, se muestra el código de fragmento que realiza esta acción:

def preprocess_for_eval(image, height, width, central_fraction=0.875):
  with tf.name_scope(scope, 'eval_image', [image, height, width]):
    if image.dtype != tf.float32:
      image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    image = tf.image.central_crop(image, central_fraction=central_fraction)
    image = tf.expand_dims(image, 0)
    image = tf.image.resize_bilinear(image, [height, width], align_corners=False)
    image = tf.squeeze(image, [0])
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)
    image.set_shape([height, width, 3])
    return image

Durante el entrenamiento, el recorte es aleatorio: se elige un cuadro de límite de manera aleatoria para seleccionar una región de la imagen a la que, posteriormente, se le cambia el tamaño. La imagen redimensionada se gira opcionalmente y los colores se distorsionan. A continuación, se muestra el código de fragmento que realiza esta acción:

def preprocess_for_train(image, height, width, bbox, fast_mode=True, scope=None):
  with tf.name_scope(scope, 'distort_image', [image, height, width, bbox]):
    if bbox is None:
      bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4])
    if image.dtype != tf.float32:
      image = tf.image.convert_image_dtype(image, dtype=tf.float32)

    distorted_image, distorted_bbox = distorted_bounding_box_crop(image, bbox)
    distorted_image.set_shape([None, None, 3])

    num_resize_cases = 1 if fast_mode else 4
    distorted_image = apply_with_random_selector(
        distorted_image,
        lambda x, method: tf.image.resize_images(x, [height, width], method),
        num_cases=num_resize_cases)

    distorted_image = tf.image.random_flip_left_right(distorted_image)

    if FLAGS.use_fast_color_distort:
      distorted_image = distort_color_fast(distorted_image)
    else:
      num_distort_cases = 1 if fast_mode else 4
      distorted_image = apply_with_random_selector(
          distorted_image,
          lambda x, ordering: distort_color(x, ordering, fast_mode),
          num_cases=num_distort_cases)

    distorted_image = tf.subtract(distorted_image, 0.5)
    distorted_image = tf.multiply(distorted_image, 2.0)
    return distorted_image

La función distort_color se encarga de modificar el color. Ofrece una forma rápida en la que solo se modifican el brillo y la saturación. El modo completo modifica el brillo, la saturación y el tono, y altera el orden en que se modifican estos aspectos de forma aleatoria. A continuación, se muestra un fragmento de código de esta función:

def distort_color(image, color_ordering=0, fast_mode=True, scope=None):
  with tf.name_scope(scope, 'distort_color', [image]):
    if fast_mode:
      if color_ordering == 0:
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
      else:
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
    else:
      if color_ordering == 0:
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
      elif color_ordering == 1:
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
      elif color_ordering == 2:
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
      elif color_ordering == 3:
        image = tf.image.random_hue(image, max_delta=0.2)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)

    return tf.clip_by_value(image, 0.0, 1.0)

La función distort_color es costosa en términos de procesamiento, en parte debido a las conversiones no lineales de RGB a HSV y de HSV a RGB que se requieren para acceder al tono y la saturación. Tanto el modo rápido como el completo requieren estas conversiones y, aunque el modo rápido es menos costoso desde el punto de vista del procesamiento, expone el modelo a la región vinculada al procesamiento de la CPU, siempre que esté habilitada.

Como alternativa, se agregó una nueva función distort_color_fast a la lista de opciones. Esta función asigna la imagen de RGB a YCrCb con el esquema de conversión JPEG y altera aleatoriamente el brillo y los Chromas/Cb antes de asignarlos a RGB. La función se muestra en el siguiente fragmento de código:

def distort_color_fast(image, scope=None):
  with tf.name_scope(scope, 'distort_color', [image]):
    br_delta = random_ops.random_uniform([], -32./255., 32./255., seed=None)
    cb_factor = random_ops.random_uniform(
        [], -FLAGS.cb_distortion_range, FLAGS.cb_distortion_range, seed=None)
    cr_factor = random_ops.random_uniform(
        [], -FLAGS.cr_distortion_range, FLAGS.cr_distortion_range, seed=None)

    channels = tf.split(axis=2, num_or_size_splits=3, value=image)
    red_offset = 1.402 * cr_factor + br_delta
    green_offset = -0.344136 * cb_factor - 0.714136 * cr_factor + br_delta
    blue_offset = 1.772 * cb_factor + br_delta
    channels[0] += red_offset
    channels[1] += green_offset
    channels[2] += blue_offset
    image = tf.concat(axis=2, values=channels)
    image = tf.clip_by_value(image, 0., 1.)
    return image

Aquí hay una imagen de muestra después del procesamiento previo. Se seleccionó una región de la imagen elegida al azar y se modificaron los colores con la función distort_color_fast.

image

La función distort_color_fast es eficiente en términos de procesamiento y, aun así, permite que el entrenamiento esté vinculado al tiempo de ejecución de la TPU. Además, muestra resultados positivos y se utilizó satisfactoriamente para entrenar el modelo Inception v3 con una exactitud superior al 78.1%, con lotes en el rango de 1,024 a 16,384. Se utiliza como opción predeterminada de Inception v3.

Optimizador

El modelo actual muestra tres tipos de optimizadores: SGD, momentum y RMSProp.

Stochastic gradient descent (SGD) es el tipo de actualización más simple: los pesos se desplazan en la dirección negativa del gradiente. A pesar de la simplicidad, es posible obtener resultados positivos en algunos modelos. La dinámica de las actualizaciones se puede escribir de la siguiente manera:

$$w_{k+1}=w_k-\alpha∇f(w_k)$$

El momentum es un optimizador popular que, con frecuencia, conduce a una convergencia más rápida que la que alcanza el SGD. Este optimizador actualiza los pesos como SGD, pero también agrega un componente en la dirección de la actualización anterior. La dinámica de la actualización se ilustra de la siguiente manera:

$$z_{k+1}=\beta z_k+∇f(w_k)$$
$$w_{k+1}=w_k-\alpha z_{k+1}$$

que puede escribirse de esta manera:

$$w_{k+1}=w_k-\alpha ∇f(w_k)+\beta \left(w_k-w_{k-1}\right)$$

El último término es el componente de la dirección de la actualización anterior. Esto se representa gráficamente en la siguiente imagen:

image

Para el momentum \({\beta}\), utilizamos el valor frecuente de 0.9.

RMSprop es un optimizador popular que propuso por primera vez Geoff Hinton en una de sus clases. La dinámica de actualización se ilustra de la siguiente manera:

$$g_{k+1}^{-2} = \alpha g_{k}^{-2} + (1-\alpha) g_{k}^2$$ $$w_{k+1}=\beta w_k + \frac{\eta}{\sqrt {g_{k+1^{\mathbf{+{\epsilon}}}}^{-2}}} ∇f(w_k)$$

En el caso de Inception v3, las pruebas muestran que RMSProp ofrece los mejores resultados en términos de máxima exactitud y de tiempo suficiente para lograrla, con un momentum de casi un segundo. Por lo tanto, RMSprop se configura como optimizador predeterminado. Los parámetros utilizados son los siguientes: disminución \({\alpha}\)=0.9, momentum \({\beta}\) 0.9 y ({\epsilon}\)=1.0.

A continuación, se muestra el fragmento de código con las opciones de optimizador.

if FLAGS.optimizer == 'sgd':
  tf.logging.info('Using SGD optimizer')
  optimizer = tf.train.GradientDescentOptimizer(
      learning_rate=learning_rate)
elif FLAGS.optimizer == 'momentum':
  tf.logging.info('Using Momentum optimizer')
  optimizer = tf.train.MomentumOptimizer(
      learning_rate=learning_rate, momentum=0.9)
elif FLAGS.optimizer == 'RMS':
  tf.logging.info('Using RMS optimizer')
  optimizer = tf.train.RMSPropOptimizer(
      learning_rate,
      RMSPROP_DECAY,
      momentum=RMSPROP_MOMENTUM,
      epsilon=RMSPROP_EPSILON)
else:
  tf.logging.fatal('Unknown optimizer:', FLAGS.optimizer)

Cuando se ejecuta en las TPU y se utiliza la API de Estimator, el optimizador debe incluirse en una función CrossShardOptimizer para garantizar la sincronización entre las réplicas (junto con cualquier comunicación cruzada que sea necesaria). A continuación, se proporciona el fragmento de código en el que se realiza esta acción en Inception v3:

if FLAGS.use_tpu:
    optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
  train_op = optimizer.minimize(loss, global_step=global_step)

Media móvil exponencial

El procedimiento normal en el entrenamiento tiene el objetivo de lograr que los parámetros entrenables se actualicen durante la propagación inversa, de acuerdo con las reglas de actualización del optimizador. Este tema se explicó en la sección anterior y se repite aquí para mayor practicidad:

$${\theta_{k+1}} = {\theta_k}-{\alpha ∇f(\theta_k)} \qquad(SGD)$$
$${\theta_{k+1}}={\theta_k}-{\alpha z_{k+1}} \qquad(momentum)$$
$${\theta_{k+1}}= {\beta \theta_k}+\frac{\eta}{\sqrt {g_{k+1^{\mathbf+{\epsilon}}}}^{-2}} ∇f(\theta_k) \qquad(RMSprop)$$

La media móvil exponencial (también conocida como suavizamiento exponencial lineal) es un paso opcional del procesamiento previo que se aplica a los pesos actualizados y, a veces, puede llevar a mejoras notables en el rendimiento. Inception v3 obtiene beneficios importantes gracias a este paso adicional. TensorFlow proporciona la función tf.train.ExponentialMovingAverage, que calcula el EMA \({\hat{\theta}}\) de peso \({\theta}\) mediante la fórmula:

$${\hat{\theta_t}}={\alpha {\hat{\theta}{_{t-1}}}}+{(1-\alpha)}{\theta_t}$$

según la cual, \({\alpha}\) representa un factor de disminución (cercano a 1.0). En el caso de Inception v3, \({\alpha}\) se configura en 0.995.

Aunque este sea un filtro de respuesta de impulso infinito (IIR), el factor de disminución establece un intervalo efectivo en el que reside la mayor parte de la energía (o muestras relevantes), como se muestra en el siguiente diagrama:

image

Para ver esto con mayor claridad, reformulamos la ecuación del filtro de la siguiente manera:

$${\hat{\theta}_{t+T+1}}={\alpha(1-\alpha)}({\theta_{t+T}}+{\alpha \theta_{t+T-1}}+...+{\alpha^{t+T}}{\theta_0})$$

en el mismo espacio en que habíamos utilizado \({\hat\theta_{-1}}=0\).

Los valores de \({\alpha}^k\) disminuyen a medida que aumenta k, por lo que, efectivamente, solo un subconjunto de las muestras tendrá una influencia considerable en \(\hat {\theta}_{t+T+1}\). La regla general de la duración de ese intervalo es \(\frac {1} {1-\alpha}\), que corresponde a \({\alpha}\)=200 para =0.995.

Primero, obtenemos una colección de variables entrenables y, posteriormente, usamos el método apply() a fin de crear variables de sombra pertenecientes a cada variable entrenada (y agregamos las operaciones correspondientes a fin de mantener las medias móviles de estas en sus instantáneas). A continuación, se muestra un fragmento del código que realiza esta acción en Inception v3:

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
  train_op = optimizer.minimize(loss, global_step=global_step)

if FLAGS.moving_average:
  ema = tf.train.ExponentialMovingAverage(
      decay=MOVING_AVERAGE_DECAY, num_updates=global_step)
  variables_to_average = (tf.trainable_variables() +
                          tf.moving_average_variables())
  with tf.control_dependencies([train_op]), tf.name_scope('moving_average'):
    train_op = ema.apply(variables_to_average)

Nos gustaría usar las variables ema durante la evaluación. Para lograr esto, definimos la clase LoadEMAHookque aplica el método variables_to_restore() al archivo de punto de control a fin de evaluar mediante el uso de los nombres de variables de sombra:

class LoadEMAHook(tf.train.SessionRunHook):
  def __init__(self, model_dir):
    super(LoadEMAHook, self).__init__()
    self._model_dir = model_dir

  def begin(self):
    ema = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY)
    variables_to_restore = ema.variables_to_restore()
    self._load_ema = tf.contrib.framework.assign_from_checkpoint_fn(
        tf.train.latest_checkpoint(self._model_dir), variables_to_restore)

  def after_create_session(self, sess, coord):
    tf.logging.info('Reloading EMA...')
    self._load_ema(sess)

La función hooks se pasa a evaluate(), como se muestra en el siguiente fragmento:

if FLAGS.moving_average:
    eval_hooks = [LoadEMAHook(FLAGS.model_dir)]
else:
    eval_hooks = []

    …

eval_results = inception_classifier.evaluate(
    input_fn=InputPipeline(False), steps=eval_steps, hooks=eval_hooks)

Normalización por lotes

La normalización por lotes es una técnica muy utilizada para normalizar los atributos de entrada en modelos que pueden ofrecer una reducción sustancial en el tiempo de convergencia. Es una de las mejoras algorítmicas más populares y útiles en el aprendizaje automático durante los últimos años; además, se utiliza en una amplia gama de modelos, incluido Inception v3.

Las entradas de activación, en primer lugar, se normalizan restando la media del lote y dividiéndola por la desviación estándar del lote; pero la normalización por lotes también cumple otras funciones. Para mantener el sistema balanceado en presencia de la propagación inversa, se ingresan dos parámetros entrenables en cada capa. Las salidas normalizadas \({\hat{x}}\) experimentan una operación posterior \({\gamma\hat {x}}+\beta\), en la cual \({\gamma}\) y \({\beta}\) constituyen una especie de desviación estándar y media; pero estas son procesadas por el propio modelo.

El conjunto completo de ecuaciones está en el documento y se repite aquí para mayor practicidad:

Entrada: valores de x sobre un minilote: \(\Phi=\) { \({x_{1..m}\\} \) }. Parámetros que hay que aprender: \({\gamma}\), \({\beta}\)

Salida: { \({y_i}=BN_{\gamma,\beta}{(x_i)}\) }

\[{\mu_\phi} \leftarrow {\frac{1}{m}}{\sum_{i=1}^m}x_i \qquad \mathsf(mini-batch\ mean)\]

\[{\sigma_\phi}^2 \leftarrow {\frac{1}{m}}{\sum_{i=1}^m} {(x_i - {\mu_\phi})^2} \qquad \mathbf(mini-batch\ variance)\]

\[{\hat{x_i}} \leftarrow {\frac{x_i-{\mu_\phi}}{\sqrt {\sigma^2_\phi}+{\epsilon}}}\qquad \mathbf(normalize)\]

\[{y_i}\leftarrow {\gamma \hat{x_i}} + \beta \equiv BN_{\gamma,\beta}{(x_i)}\qquad \mathbf(scale \ and \ shift)\]

La normalización ocurre satisfactoriamente durante el entrenamiento, pero en el momento de la evaluación, nos gustaría que el modelo se comporte de una manera determinista: El resultado de la clasificación de una imagen debe depender únicamente de la imagen de entrada y no del conjunto de imágenes que se están enviando al modelo. Por lo tanto, necesitamos corregir \({\mu} \) y \({\ sigma}^2\), y estas fórmulas deben representar las estadísticas de propagación de imágenes.

Para lograr este objetivo, el modelo calcula los promedios móviles de la media y la varianza con respecto a los minilotes:

\[{\hat\mu_i} = {\alpha \hat\mu_{t-1}}+{(1-\alpha)\mu_t}\]

\[{\hat\sigma_t}^2 = {\alpha{\hat\sigma^2_{t-1}}} + {(1-\alpha) {\sigma_t}^2}\]

En el caso particular de Inception v3, se obtuvo un factor de disminución razonable (a través del ajuste de hiperparámetro) para utilizarlo en las GPU. También nos gustaría usar este valor en la TPU; pero para hacer eso, necesitamos hacer algunos ajustes.

La media móvil y la varianza de la normalización por lotes se calculan mediante un filtro de paso con pérdidas, cuya ecuación canónica se muestra a continuación (aquí: \({y_t}\), y representa la media o varianza móvil):

\[{y_t}={\alpha y_{t-1}}+{(1-\alpha)}{x_t} \]

(1)

En un trabajo de GPU (síncrono) de 8 x 1, cada réplica lee la media móvil actual y la actualiza. Las actualizaciones son secuenciales, ya que la réplica actual debe escribir la nueva variable en movimiento antes de que la próxima pueda leerla:

Cuando hay 8 réplicas, el conjunto de operaciones de una actualización de ensamble es el siguiente:

\[{y_t}={\alpha y_{t-1}}+{(1-\alpha)}{x_t} \]

\[{y_{t+1}}={\alpha y_{t}}+{(1-\alpha)}{x_{t+1}} \]

\[{y_{t+2}}={\alpha y_{t+1}}+{(1-\alpha)}{x_{t+2}} \]

\[{y_{t+3}}={\alpha y_{t+2}}+{(1-\alpha)}{x_{t+3}} \]

\[{y_{t+4}}={\alpha y_{t+3}}+{(1-\alpha)}{x_{t+4}} \]

\[{y_{t+5}}={\alpha y_{t+4}}+{(1-\alpha)}{x_{t+5}} \]

\[{y_{t+6}}={\alpha y_{t+5}}+{(1-\alpha)}{x_{t+6}} \]

\[{y_{t+7}}={\alpha y_{t+6}}+{(1-\alpha)}{x_{t+7}} \]

Este conjunto de 8 actualizaciones secuenciales se puede escribir de la siguiente manera:

\[{y_{t+7}}={\alpha^8y_{t-1}}+(1-\alpha){\sum_{k=0}^7} {\alpha^{7-k}}{x_{t+k}}\]

(2)

En la implementación de cálculo de momento móvil actual en la TPU, cada fragmento realiza cálculos de forma independiente, y no hay comunicación de fragmentos cruzados. Los lotes se distribuyen a cada fragmento, y cada uno de ellos procesa 1/8 de la cantidad total de lotes (cuando hay 8 fragmentos).

Aunque cada fragmento pasa por los movimientos y calcula los momentos móviles (es decir, la media y la varianza), solo los resultados del fragmento 0 se comunican a la CPU del host. Entonces, efectivamente, solo una réplica se encarga de actualizar la media o varianza móviles:

\[{z_t}={\beta {z_{t-1}}}+{(1-\beta)u_t}\]

(3)

y esta actualización ocurre a 1/8 respecto de la tasa de su contraparte secuencial.

Para comparar las ecuaciones de actualización de la GPU y la TPU, necesitamos alinear las escalas de tiempo respectivas. Específicamente, el conjunto de operaciones que comprende un conjunto de 8 actualizaciones secuenciales en la GPU debe compararse con una única actualización en la TPU. Esto se ilustra en el siguiente diagrama:

image

Veamos las ecuaciones con los índices de tiempo modificados:

\[{y_t}={\alpha^8y_{t-1}}+(1-\alpha){\sum_{k=0}^7} {\alpha^{7-k}}{x_{t-k/8}} \qquad \mathsf(GPU)\]

\[{z_t}={\beta {z_{t-1}}}+{(1-\beta)u_t}\qquad \mathsf(TPU) \]

Si suponemos, de forma simplificada, que los 8 minilotes (normalizados en todas las dimensiones relevantes) generan valores similares dentro de la actualización secuencial de 8 minilotes de la GPU, podemos aproximar estas ecuaciones de la siguiente manera:

\[{y_t}={\alpha^8y_{t-1}}+(1-\alpha){\sum_{k=0}^7} {\alpha^{7-k}}{\hat{x_t}}={\alpha^8y_{t-1}+(1-\alpha^8){\hat{x_t}}} \qquad \mathsf(GPU)\]

\[{z_t}={\beta {z_{t-1}}}+{(1-\beta)u_t}\qquad \mathsf(TPU) \]

Por lo tanto, para hacer coincidir el efecto de un factor de disminución dado en la GPU, debemos modificar el factor de disminución en la TPU según corresponda. Específicamente, necesitamos configurar \({\beta}\)=\({\alpha}^8\).

Para Inception v3, el valor de disminución utilizado en la GPU es \({\alpha}\)=0.9997, que se traduce en un valor de disminución de la TPU de \({\beta}\)=0.9976.

Adaptación de la tasa de aprendizaje

A medida que aumenta el tamaño de los lotes, el entrenamiento se vuelve más difícil. Se siguen proponiendo diferentes técnicas a fin de lograr un entrenamiento eficiente para los lotes de mayor tamaño (por ejemplo, puedes consultar aquí, aquí y aquí).

Una de esas técnicas es el aumento gradual de la tasa de aprendizaje. Esta herramienta se usó con el fin de entrenar el modelo con una exactitud superior al 78.1% para tamaños de lotes que van desde 4,096 a 16,384. En el caso de Inception v3, la tasa de aprendizaje primero se fija en aproximadamente el 10% del total de lo que sería la tasa de aprendizaje inicial. Esta tasa permanece constante en este valor bajo para una cantidad específica (pequeña) de "ciclos de entrenamiento fríos" y, luego, comienza un aumento lineal para una cantidad determinada de "ciclos de entrenamiento calientes", al final de los cuales se interseca con el valor de lo que habría sido la tasa de aprendizaje si se hubiera utilizado una disminución exponencial normal. Este ejemplo se ilustra en la siguiente imagen.

image

La parte del código que lleva a cabo este proceso se muestra en el siguiente fragmento de código:

initial_learning_rate = FLAGS.learning_rate * FLAGS.train_batch_size / 256
if FLAGS.use_learning_rate_warmup:
  warmup_decay = FLAGS.learning_rate_decay**(
    (FLAGS.warmup_epochs + FLAGS.cold_epochs) /
    FLAGS.learning_rate_decay_epochs)
  adj_initial_learning_rate = initial_learning_rate * warmup_decay

final_learning_rate = 0.0001 * initial_learning_rate

train_op = None
if training_active:
  batches_per_epoch = _NUM_TRAIN_IMAGES / FLAGS.train_batch_size
  global_step = tf.train.get_or_create_global_step()
  current_epoch = tf.cast(
    (tf.cast(global_step, tf.float32) / batches_per_epoch), tf.int32)

  learning_rate = tf.train.exponential_decay(
    learning_rate=initial_learning_rate,
    global_step=global_step,
    decay_steps=int(FLAGS.learning_rate_decay_epochs * batches_per_epoch),
    decay_rate=FLAGS.learning_rate_decay,
    staircase=True)

  if FLAGS.use_learning_rate_warmup:
    wlr = 0.1 * adj_initial_learning_rate
    wlr_height = tf.cast(
      0.9 * adj_initial_learning_rate /
      (FLAGS.warmup_epochs + FLAGS.learning_rate_decay_epochs - 1),
      tf.float32)
    epoch_offset = tf.cast(FLAGS.cold_epochs - 1, tf.int32)
    exp_decay_start = (FLAGS.warmup_epochs + FLAGS.cold_epochs +
                   FLAGS.learning_rate_decay_epochs)
    lin_inc_lr = tf.add(
      wlr, tf.multiply(
        tf.cast(tf.subtract(current_epoch, epoch_offset), tf.float32),
        wlr_height))
    learning_rate = tf.where(
      tf.greater_equal(current_epoch, FLAGS.cold_epochs),
      (tf.where(tf.greater_equal(current_epoch, exp_decay_start),
              learning_rate, lin_inc_lr)),
       wlr)

  # Set a minimum boundary for the learning rate.
  learning_rate = tf.maximum(
      learning_rate, final_learning_rate, name='learning_rate')