Capacitación distribuida NasNet con tensorflow_cloud y Google Cloud

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver en GitHub

Descargar libreta

Corre en Kaggle

Este ejemplo se basa en la clasificación de imágenes a través de ajustes con EfficientNet para demostrar cómo entrenar un modelo de NasNetMobile usando tensorflow_cloud y Google Cloud Platform a escala usando entrenamiento distribuido.

Importar módulos requeridos

import tensorflow as tf
tf.version.VERSION

'2.6.0'

! pip install -q tensorflow-cloud

import tensorflow_cloud as tfc
tfc.__version__

import sys

Configuraciones del proyecto

Establecer los parámetros del proyecto. Para conocer los parámetros específicos de Google Cloud, consulte las Instrucciones de configuración del proyecto de Google Cloud .

# Set Google Cloud Specific parameters

# TODO: Please set GCP_PROJECT_ID to your own Google Cloud project ID.
GCP_PROJECT_ID = 'YOUR_PROJECT_ID'

# TODO: set GCS_BUCKET to your own Google Cloud Storage (GCS) bucket.
GCS_BUCKET = 'YOUR_GCS_BUCKET_NAME'

# DO NOT CHANGE: Currently only the 'us-central1' region is supported.
REGION = 'us-central1'

# OPTIONAL: You can change the job name to any string.
JOB_NAME = 'nasnet'

# Setting location were training logs and checkpoints will be stored
GCS_BASE_PATH = f'gs://{GCS_BUCKET}/{JOB_NAME}'
TENSORBOARD_LOGS_DIR = os.path.join(GCS_BASE_PATH,"logs")
MODEL_CHECKPOINT_DIR = os.path.join(GCS_BASE_PATH,"checkpoints")
SAVED_MODEL_DIR = os.path.join(GCS_BASE_PATH,"saved_model")

Autenticar el cuaderno para usar su proyecto de Google Cloud

Para Kaggle Notebooks, haga clic en "Complementos"->"Google Cloud SDK" antes de ejecutar la celda a continuación.

# Using tfc.remote() to ensure this code only runs in notebook
if not tfc.remote():

    # Authentication for Kaggle Notebooks
    if "kaggle_secrets" in sys.modules:
        from kaggle_secrets import UserSecretsClient
        UserSecretsClient().set_gcloud_credentials(project=GCP_PROJECT_ID)

    # Authentication for Colab Notebooks
    if "google.colab" in sys.modules:
        from google.colab import auth
        auth.authenticate_user()
        os.environ["GOOGLE_CLOUD_PROJECT"] = GCP_PROJECT_ID

Cargar y preparar datos

Lea datos sin procesar y divídalos para entrenar y probar conjuntos de datos.

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

# Setting input specific parameters
# The model expects input of dimension (INPUT_IMG_SIZE, INPUT_IMG_SIZE, 3)
INPUT_IMG_SIZE = 32
NUM_CLASSES = 10

Agregue API de capas de preprocesamiento para el aumento de imágenes.

from tensorflow.keras.layers.experimental import preprocessing
from tensorflow.keras.models import Sequential


img_augmentation = Sequential(
    [
        # Resizing input to better match ImageNet size
        preprocessing.Resizing(256, 256),
        preprocessing.RandomRotation(factor=0.15),
        preprocessing.RandomFlip(),
        preprocessing.RandomContrast(factor=0.1),
    ],
    name="img_augmentation",
)

Cargue el modelo y prepárese para el entrenamiento

Cargaremos un modelo preentrenado de NASNetMobile (con pesos) y descongelaremos algunas capas para ajustar el modelo para que coincida mejor con el conjunto de datos.

from tensorflow.keras import layers

def build_model(num_classes, input_image_size):
    inputs = layers.Input(shape=(input_image_size, input_image_size, 3))
    x = img_augmentation(inputs)

    model = tf.keras.applications.NASNetMobile(
        input_shape=None,
        include_top=False,
        weights="imagenet",
        input_tensor=x,
        pooling=None,
        classes=num_classes,
    )

    # Freeze the pretrained weights
    model.trainable = False

    # We unfreeze the top 20 layers while leaving BatchNorm layers frozen
    for layer in model.layers[-20:]:
        if not isinstance(layer, layers.BatchNormalization):
            layer.trainable = True

    # Rebuild top
    x = layers.GlobalAveragePooling2D(name="avg_pool")(model.output)
    x = layers.BatchNormalization()(x)

    x = layers.Dense(128, activation="relu")(x)
    x = layers.Dense(64, activation="relu")(x)
    outputs = layers.Dense(num_classes, activation="softmax", name="pred")(x)

    # Compile
    model = tf.keras.Model(inputs, outputs, name="NASNetMobile")
    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-4)
    model.compile(
        optimizer=optimizer,
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"]
    )
    return model

model = build_model(NUM_CLASSES, INPUT_IMG_SIZE)

if tfc.remote():
    # Configure Tensorboard logs
    callbacks=[
        tf.keras.callbacks.TensorBoard(log_dir=TENSORBOARD_LOGS_DIR),
        tf.keras.callbacks.ModelCheckpoint(
            MODEL_CHECKPOINT_DIR,
            save_best_only=True),
        tf.keras.callbacks.EarlyStopping(
            monitor='loss',
            min_delta =0.001,
            patience=3)]

    model.fit(x=x_train, y=y_train, epochs=100,
              validation_split=0.2, callbacks=callbacks)

    model.save(SAVED_MODEL_DIR)

else:
    # Run the training for 1 epoch and a small subset of the data to validate setup
    model.fit(x=x_train[:100], y=y_train[:100], validation_split=0.2, epochs=1)

Iniciar el entrenamiento a distancia

Este paso preparará su código desde este cuaderno para la ejecución remota e iniciará un entrenamiento distribuido de forma remota en Google Cloud Platform para entrenar el modelo. Una vez que se envía el trabajo, puede ir al siguiente paso para monitorear el progreso de los trabajos a través de Tensorboard.

# If you are using a custom image you can install modules via requirements
# txt file.
with open('requirements.txt','w') as f:
    f.write('tensorflow-cloud\n')

# Optional: Some recommended base images. If you provide none the system
# will choose one for you.
TF_GPU_IMAGE= "tensorflow/tensorflow:latest-gpu"
TF_CPU_IMAGE= "tensorflow/tensorflow:latest"

# Submit a distributed training job using GPUs.
tfc.run(
    distribution_strategy='auto',
    requirements_txt='requirements.txt',
    docker_config=tfc.DockerConfig(
        parent_image=TF_GPU_IMAGE,
        image_build_bucket=GCS_BUCKET
        ),
    chief_config=tfc.COMMON_MACHINE_CONFIGS['K80_1X'],
      worker_config=tfc.COMMON_MACHINE_CONFIGS['K80_1X'],
      worker_count=3,
    job_labels={'job': JOB_NAME}
)

Resultados de entrenamiento

Vuelva a conectar su instancia de Colab

La mayoría de los trabajos de capacitación remota son de larga duración. Si está utilizando Colab, es posible que se agote el tiempo de espera antes de que los resultados de la capacitación estén disponibles. En ese caso, vuelva a ejecutar las siguientes secciones para volver a conectar y configurar su instancia de Colab para acceder a los resultados de la capacitación. Ejecute las siguientes secciones en orden:

Importar módulos requeridos
Configuraciones del proyecto
Autenticar el cuaderno para usar su proyecto de Google Cloud

Tensor de carga

Mientras la capacitación está en progreso, puede usar Tensorboard para ver los resultados. Tenga en cuenta que los resultados se mostrarán solo después de que haya comenzado su entrenamiento. Esto puede tomar unos pocos minutos.

%load_ext tensorboard
%tensorboard --logdir $TENSORBOARD_LOGS_DIR

Cargue su modelo entrenado

trained_model = tf.keras.models.load_model(SAVED_MODEL_DIR)
trained_model.summary()