Pelatihan terdistribusi NasNet dengan tensorflow_cloud dan Google Cloud

Lihat di TensorFlow.org Lihat di GitHub Unduh buku catatan Logo KaggleJalankan di Kaggle

Contoh ini didasarkan pada klasifikasi Gambar melalui penyesuaian dengan EfficientNet untuk menunjukkan cara melatih model NasNetMobile menggunakan tensorflow_cloud dan Google Cloud Platform dalam skala besar menggunakan pelatihan terdistribusi.

Impor modul yang diperlukan 

import tensorflow as tf
tf.version.VERSION

'2.6.0'

! pip install -q tensorflow-cloud

import tensorflow_cloud as tfc
tfc.__version__

import sys

Konfigurasi Proyek

Tetapkan parameter proyek. Untuk parameter khusus Google Cloud, lihat Petunjuk Penyiapan Project 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")

Mengautentikasi notebook untuk menggunakan Proyek Google Cloud Anda

Untuk Kaggle Notebooks, klik "Add-on"->"Google Cloud SDK" sebelum menjalankan sel di bawah.

# 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

Memuat dan menyiapkan data

Membaca data mentah dan membaginya untuk melatih dan menguji kumpulan data.

(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

Tambahkan API lapisan prapemrosesan untuk augmentasi gambar.

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

Muat model dan bersiap untuk pelatihan

Kami akan memuat model NASNetMobile yang telah dilatih sebelumnya (dengan bobot) dan mencairkan beberapa lapisan untuk menyempurnakan model agar lebih cocok dengan kumpulan data.

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)

Mulai pelatihan jarak jauh

Langkah ini akan menyiapkan kode Anda dari notebook ini untuk eksekusi jarak jauh dan memulai pelatihan terdistribusi dari jarak jauh di Google Cloud Platform untuk melatih model. Setelah tugas dikirimkan, Anda dapat melanjutkan ke langkah berikutnya untuk memantau kemajuan tugas melalui 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}
)

Hasil Pelatihan

Hubungkan kembali instance Colab Anda

Sebagian besar pekerjaan pelatihan jarak jauh berlangsung lama. Jika Anda menggunakan Colab, mungkin waktu habis sebelum hasil pelatihan tersedia. Jika demikian, jalankan kembali bagian berikut untuk menyambungkan kembali dan mengonfigurasi instance Colab Anda untuk mengakses hasil pelatihan. Jalankan bagian berikut secara berurutan:

  1. Impor modul yang diperlukan
  2. Konfigurasi Proyek
  3. Mengautentikasi notebook untuk menggunakan Proyek Google Cloud Anda

Muat Papan Tensor

Saat pelatihan sedang berlangsung, Anda dapat menggunakan Tensorboard untuk melihat hasilnya. Perhatikan bahwa hasilnya hanya akan terlihat setelah pelatihan Anda dimulai. Ini mungkin memakan waktu beberapa menit. 

%load_ext tensorboard
%tensorboard --logdir $TENSORBOARD_LOGS_DIR

Muat model terlatih Anda 

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