Formazione distribuita NasNet con tensorflow_cloud e Google Cloud

Visualizza su TensorFlow.org Visualizza su GitHub Scarica il taccuino Logo KaggleCorri a Kaggle

Questo esempio si basa sulla classificazione delle immagini tramite messa a punto con EfficientNet per dimostrare come addestrare un modello NasNetMobile utilizzando tensorflow_cloud e Google Cloud Platform su larga scala utilizzando la formazione distribuita.

Importa i moduli richiesti 

import tensorflow as tf
tf.version.VERSION

'2.6.0'

! pip install -q tensorflow-cloud

import tensorflow_cloud as tfc
tfc.__version__

import sys

Configurazioni di progetto

Imposta i parametri del progetto. Per i parametri specifici di Google Cloud, fare riferimento alle istruzioni di configurazione del progetto 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")

Autenticazione del notebook per utilizzare il tuo progetto Google Cloud

Per i notebook Kaggle, fai clic su "Componenti aggiuntivi" -> "Google Cloud SDK" prima di eseguire la cella sottostante.

# 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

Caricare e preparare i dati

Leggi i dati grezzi e dividili per addestrare e testare set di dati.

(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

Aggiungi API dei livelli di preelaborazione per l'aumento delle immagini.

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

Carica il modello e preparati per l'addestramento

Caricheremo un modello preaddestrato NASNetMobile (con pesi) e sbloccheremo alcuni livelli per ottimizzare il modello in modo che corrisponda meglio al set di dati.

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)

Inizia la formazione a distanza

Questo passaggio preparerà il codice da questo notebook per l'esecuzione remota e avvierà un addestramento distribuito in remoto su Google Cloud Platform per addestrare il modello. Una volta inviato il lavoro, puoi andare al passaggio successivo per monitorare l'avanzamento dei lavori tramite 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}
)

Risultati della formazione

Riconnetti la tua istanza Colab

La maggior parte dei lavori di formazione a distanza sono di lunga durata. Se utilizzi Colab, potrebbe verificarsi un timeout prima che i risultati della formazione siano disponibili. In tal caso, esegui nuovamente le seguenti sezioni per riconnetterti e configurare la tua istanza Colab per accedere ai risultati della formazione. Esegui le seguenti sezioni in ordine:

  1. Importa i moduli richiesti
  2. Configurazioni di progetto
  3. Autenticazione del notebook per utilizzare il tuo progetto Google Cloud

Carica Tensorboard

Mentre la formazione è in corso puoi utilizzare Tensorboard per visualizzare i risultati. Tieni presente che i risultati verranno visualizzati solo dopo l'inizio dell'allenamento. L'operazione potrebbe richiedere alcuni minuti. 

%load_ext tensorboard
%tensorboard --logdir $TENSORBOARD_LOGS_DIR

Carica il tuo modello addestrato 

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