TensorFlow クラウド

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

TensorFlow Cloud は、Google Cloud 上で Keras モデルのトレーニングとハイパーパラメータ調整を簡単に行うことができるライブラリです。

TensorFlow Cloud のrun API を使用すると、モデル コードを Google Cloud アカウントに直接送信でき、(コンソールでプロジェクトを設定した後は) ログインして Cloud UI を操作することなく、Google Cloud のコンピューティング リソースを使用できます。

これは、Python ノートブック内から Google Cloud コンピューティング リソースを直接使用できることを意味します。これは、まさにこのようなノートブックです。プレーンな.py Python スクリプトからモデルを Google Cloud に送信することもできます。

簡単な例

これは、 TensorFlow Cloudと Google Cloud を使用してモデルをリモートでトレーニングする方法を示す簡単な入門例です。

これを読んで仕組みを理解することも、Google Colab でノートブックを実行することもできます。ノートブックを実行するには、Google Cloud アカウントに接続し、認証情報とプロジェクト ID を入力する必要があります。まだアカウントをお持ちでない場合、またはコンソールでプロジェクトを設定する方法がわからない場合は、「Google Cloud アカウントの設定と接続」を参照してください。

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

import tensorflow as tf
tf.version.VERSION
'2.6.0'
! pip install -q tensorflow-cloud
import tensorflow_cloud as tfc
print(tfc.__version__)
import sys

プロジェクト構成

プロジェクトパラメータを設定します。 GCP_PROJECT_IDまたはGCS_BUCKETが何であるべきかわからない場合は、 「Google Cloud アカウントの設定と接続」を参照してください。

JOB_NAMEはオプションであり、任意の文字列に設定できます。大規模なプロジェクトの一部として (たとえば) 複数のトレーニング実験を行っている場合は、それぞれに一意のJOB_NAMEを与えることができます。

# 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 = 'mnist'

# 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 プロジェクトを使用するためのノートブックの認証

このコードはノートブックを認証し、有効な Google Cloud 認証情報と ID を確認します。これはif not tfc.remote()ブロック内にあり、ノートブック内でのみ実行され、ノートブック コードが Google Cloud に送信されるときには実行されないようにします。

# 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

モデルとデータのセットアップ

ここからは、MNIST データセットで分類を実行するための単純な Keras モデルをセットアップする基本手順に従います。

データのロードと分割

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

(x_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()

x_train = x_train.reshape((60000, 28 * 28))
x_train = x_train.astype('float32') / 255

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

単純なモデルを作成し、それにいくつかのコールバックを設定します。

from tensorflow.keras import layers

model = tf.keras.Sequential([
  tf.keras.layers.Dense(512, activation='relu', input_shape=(28 * 28,)),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(loss='sparse_categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=['accuracy'])

クイック検証トレーニング

すべてが正しく設定されていることを確認するために、モデルを 1 エポックだけトレーニングします。また、トレーニング コマンドがtfc.remote if notラップし、ランタイム環境でのみ実行されるようにします。 Google Cloud に送信されるときではなく、これを読んでいるのです。

if not tfc.remote():
    # 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 に送信されるときにのみ実行され、これを読んでいるランタイム内では実行されません。

まず、次のようなコールバックを設定します。

  • TensorBoardのログを作成します。
  • チェックポイントを作成し、上で指定したチェックポイント ディレクトリに保存します。
  • 損失が十分に改善されない場合は、モデルのトレーニングを停止します。

次に、 model.fitmodel.saveを呼び出します。これにより (このコードが Google Cloud で実行されている場合)、実際に完全なトレーニング (100 エポック) が実行され、トレーニングされたモデルが上で定義された GCS バケットとディレクトリに保存されます。

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)

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

TensorFlow Cloud は、ローカル実行環境 (このノートブック) からすべてのコードを取得してラップし、実行のために Google Cloud に送信します。 (そのため、 ifおよびif not tfc.remoteラッパーが重要です。)

このステップでは、このノートブックのコードをリモート実行用に準備し、Google Cloud Platform でリモート トレーニング ジョブを開始してモデルをトレーニングします。

まず、 tensorflow-cloud Python パッケージをrequirements.txtファイルに追加します。このファイルは、このノートブックのコードと一緒に送信されます。必要に応じて、ここに他のパッケージを追加できます。

次に、GPU と CPU イメージを指定します。どちらかを指定するだけで済みます。 GPU は次のコードで使用されます。

最後に、TensorFlow クラウドの中心であるtfc.runの呼び出しです。これがこのノートブック内で実行されると、このノートブックのすべてのコードとこのディレクトリ内の残りのファイルがパッケージ化され、実行のために Google Cloud に送信されます。 runメソッドのパラメータでは、GPU CPU イメージの詳細を指定します。どちらかを指定するだけで済みます。 GPU は次のコードで使用されます。

最後に、TensorFlow クラウドの中心であるtfc.runの呼び出しです。これがこのノートブック内で実行されると、このノートブックのすべてのコードとこのディレクトリ内の残りのファイルがパッケージ化され、実行のために Google Cloud に送信されます。 runメソッドのパラメータで GPU の詳細を指定し、CPU イメージを指定します。どちらかを指定するだけで済みます。 GPU は次のコードで使用されます。

最後に、TensorFlow クラウドの中心であるtfc.runの呼び出しです。これがこのノートブック内で実行されると、このノートブックのすべてのコードとこのディレクトリ内の残りのファイルがパッケージ化され、実行のために Google Cloud に送信されます。 runメソッドのパラメーターは、使用する実行環境と配布戦略 (存在する場合) の詳細を指定します。

ジョブが送信されたら、次のステップに進み、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 single node training job using GPU.
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'],
    job_labels={'job': JOB_NAME}
)

研修結果

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

ほとんどのリモート トレーニング ジョブは長時間実行されます。 Colab を使用している場合、トレーニング結果が利用可能になる前にタイムアウトになる可能性があります。

その場合は、次のセクションを再実行して、 Colab インスタンスを再接続し、トレーニング結果にアクセスできるように設定します。

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

コードの残りの部分を再実行しないでください

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

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

%load_ext tensorboard
%tensorboard --logdir $TENSORBOARD_LOGS_DIR

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

トレーニングが完了すると、上で指定した GCS バケットからモデルを取得できます。

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