Formation distribuée NasNet avec tensorflow_cloud et Google Cloud

Voir sur TensorFlow.org Voir sur GitHub Télécharger le carnet Logo KaggleCourir à Kaggle

Cet exemple est basé sur la classification d'images via un réglage fin avec EfficientNet pour montrer comment entraîner un modèle NasNetMobile à l'aide de tensorflow_cloud et de Google Cloud Platform à grande échelle à l'aide d'une formation distribuée.

Importer les modules requis 

import tensorflow as tf
tf.version.VERSION

'2.6.0'

! pip install -q tensorflow-cloud

import tensorflow_cloud as tfc
tfc.__version__

import sys

Configurations de projet

Définir les paramètres du projet. Pour les paramètres spécifiques à Google Cloud, reportez-vous aux instructions de configuration du projet 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")

Authentifier le notebook pour utiliser votre projet Google Cloud

Pour les notebooks Kaggle, cliquez sur "Modules complémentaires" -> "SDK Google Cloud" avant d'exécuter la cellule ci-dessous.

# 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

Charger et préparer les données

Lisez les données brutes et divisez-les pour former et tester des ensembles de données.

(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

Ajoutez des API de couches de prétraitement pour l’augmentation d’image.

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",
)

Charger le modèle et préparer la formation

Nous allons charger un modèle pré-entraîné NASNetMobile (avec des pondérations) et débloquer quelques couches pour affiner le modèle afin de mieux correspondre à l'ensemble de données.

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)

Commencer la formation à distance

Cette étape préparera votre code à partir de ce notebook pour une exécution à distance et démarrera une formation distribuée à distance sur Google Cloud Platform pour entraîner le modèle. Une fois le travail soumis, vous pouvez passer à l'étape suivante pour suivre la progression des travaux via 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}
)

Résultats de la formation

Reconnectez votre instance Colab

La plupart des emplois de formation à distance durent de longue durée. Si vous utilisez Colab, celui-ci peut expirer avant que les résultats de l'entraînement ne soient disponibles. Dans ce cas, réexécutez les sections suivantes pour vous reconnecter et configurer votre instance Colab afin d'accéder aux résultats de la formation. Exécutez les sections suivantes dans l'ordre :

  1. Importer les modules requis
  2. Configurations de projet
  3. Authentifier le notebook pour utiliser votre projet Google Cloud

Charger le tableau tenseur

Pendant que la formation est en cours, vous pouvez utiliser Tensorboard pour afficher les résultats. Notez que les résultats ne s'afficheront qu'après le début de votre entraînement. Cela peut prendre quelques minutes. 

%load_ext tensorboard
%tensorboard --logdir $TENSORBOARD_LOGS_DIR

Chargez votre modèle entraîné 

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