tensorflow_cloud と Google Cloud を使用した分散トレーニング NasNet

TensorFlow.org で見る Google Colab で実行するGitHub で見るノートブックをダウンロードKaggleのロゴKaggle で実行する

この例は、EfficientNet による微調整による画像分類に基づいており、分散トレーニングを使用してtensorflow_cloudと Google Cloud Platform を使用してNasNetMobileモデルを大規模にトレーニングする方法を示しています。

必要なモジュールをインポートする

import tensorflow as tf
tf.version.VERSION

'2.6.0'

! pip install -q tensorflow-cloud

import tensorflow_cloud as tfc
tfc.__version__

import sys

プロジェクト構成

プロジェクトパラメータを設定します。 Google Cloud 固有のパラメータについては、 「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")

Google Cloud プロジェクトを使用するためのノートブックの認証

Kaggle ノートブックの場合は、下のセルを実行する前に、[アドオン] -> [Google Cloud SDK] をクリックします。

# 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

データをロードして準備する

生データを読み取り、データセットをトレーニングおよびテストするために分割します。

(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

画像拡張のための前処理レイヤー API を追加します。

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

モデルをロードしてトレーニングの準備をする

NASNetMobile の事前トレーニング済みモデル (重み付き) をロードし、データセットとよりよく一致するようにモデルを微調整するためにいくつかのレイヤーをフリーズ解除します。

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)

リモートトレーニングを開始します

このステップでは、このノートブックのコードをリモート実行用に準備し、Google Cloud Platform 上でリモートで分散トレーニングを開始してモデルをトレーニングします。ジョブが送信されたら、次のステップに進み、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}
)

研修結果

Colab インスタンスを再接続します

ほとんどのリモート トレーニング ジョブは長時間実行されます。 Colab を使用している場合、トレーニング結果が利用可能になる前にタイムアウトになる可能性があります。その場合は、次のセクションを再実行して再接続し、トレーニング結果にアクセスできるように Colab インスタンスを設定します。次のセクションを順番に実行します。

  1. 必要なモジュールをインポートする
  2. プロジェクト構成
  3. Google Cloud プロジェクトを使用するためのノートブックの認証

テンソルボードをロードする

トレーニングの進行中に、Tensorboard を使用して結果を表示できます。結果はトレーニングの開始後にのみ表示されることに注意してください。これには数分かかる場合があります。 

%load_ext tensorboard
%tensorboard --logdir $TENSORBOARD_LOGS_DIR

トレーニングされたモデルをロードする

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