tensorflow_cloud 및 Google Cloud를 사용한 분산 교육 NasNet

TensorFlow.org에서 보기 Google Colab에서 실행 GitHub에서 보기 노트북 다운로드 캐글 로고캐글에서 실행

이 예는 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 Notebooks의 경우 아래 셀을 실행하기 전에 "추가 기능"->"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()