Keras와 함께 DTensor 사용

TensorFlow.org에서 보기

Google Colab에서 실행하기

GitHub에서 소스 보기

노트북 다운로드하기

개요

이 튜토리얼에서는 Keras와 함께 DTensor를 사용하는 방법을 배웁니다.

Keras와 DTensor의 통합을 통해 기존 Keras 레이어와 모델을 재사용하여 분산 머신 러닝 모델을 구축하고 훈련할 수 있습니다.

MNIST 데이터로 다중 레이어 분류 모델을 훈련합니다. 하위 클래스화 모델, 순차형 모델, 함수형 모델에 대한 레이아웃 설정에 대해 설명합니다.

이 튜토리얼은 여러분이 이미 DTensor 프로그래밍 가이드를 읽었고 Mesh 및 Layout과 같은 기본 DTensor 개념에 익숙하다고 가정합니다.

이 튜토리얼은 https://www.tensorflow.org/datasets/keras_example을 기반으로 합니다.

설정

DTensor는 TensorFlow 2.9.0 릴리스의 일부입니다.

pip install --quiet --upgrade --pre tensorflow tensorflow-datasets

다음으로 tensorflow 및 tensorflow.experimental.dtensor를 가져오고 8개의 가상 CPU를 사용하도록 TensorFlow를 구성합니다.

이 예제에서는 CPU를 사용하지만 DTensor는 CPU, GPU 또는 TPU 장치에서 동일한 방식으로 작동합니다.

import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow.experimental import dtensor

2022-12-15 01:46:36.893005: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-15 01:46:36.893123: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-15 01:46:36.893133: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

def configure_virtual_cpus(ncpu):
  phy_devices = tf.config.list_physical_devices('CPU')
  tf.config.set_logical_device_configuration(
        phy_devices[0], 
        [tf.config.LogicalDeviceConfiguration()] * ncpu)

configure_virtual_cpus(8)
tf.config.list_logical_devices('CPU')

devices = [f'CPU:{i}' for i in range(8)]

결정성 있는 의사 난수 생성기

한 가지 주의해야 할 점은 DTensor API에서 실행 중인 각 클라이언트가 동일한 임의의 시드를 갖도록 요구하므로 가중치 초기화에 대한 동작이 결정성 있게 작동할 수 있다는 것입니다. 이를 위해 tf.keras.utils.set_random_seed()를 통해 keras에 전역 시드를 설정할 수 있습니다.

tf.keras.backend.experimental.enable_tf_random_generator()
tf.keras.utils.set_random_seed(1337)

데이터 병렬 메쉬 생성하기

이 튜토리얼은 데이터 병렬 훈련을 보여줍니다. 모델 병렬 훈련 및 공간 병렬 훈련에 맞게 조정하려면 다른 Layout 객체 집합으로 전환하기만 하면 됩니다. 데이터 병렬 이외의 분산 훈련에 대한 자세한 내용은 DTensor 심층 ML 튜토리얼을 참조하세요.

데이터 병렬 훈련은 일반적으로 사용되는 병렬 훈련 방식이며 예를 들어 tf.distribute.MirroredStrategy에서도 사용됩니다.

DTensor를 사용하면 데이터 병렬 훈련 루프가 단일 '배치' 차원으로 구성된 Mesh를 사용합니다. 여기서 각 장치는 전역 배치에서 샤드를 수신하는 모델의 복제본을 실행합니다.

mesh = dtensor.create_mesh([("batch", 8)], devices=devices)

각 장치가 모델의 전체 복제본을 실행할 때 모델 변수는 메쉬 전체에 완전히 복제되어야 합니다(샤딩되지 않음). 예를 들어 이 Mesh의 랭크 2 가중치에 대해 완전히 복제된 레이아웃은 다음과 같습니다.

example_weight_layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh)  # or
example_weight_layout = dtensor.Layout.replicated(mesh, rank=2)

이 Mesh의 랭크 2 데이터 텐서에 대한 레이아웃은 첫 번째 차원을 따라 샤딩됩니다(때로 batch_sharded라고도 함).

example_data_layout = dtensor.Layout(['batch', dtensor.UNSHARDED], mesh)  # or
example_data_layout = dtensor.Layout.batch_sharded(mesh, 'batch', rank=2)

레이아웃으로 Keras 레이어 생성하기

데이터 병렬 방식에서는 일반적으로 모델의 각 복제본이 샤딩된 입력 데이터로 계산을 수행할 수 있도록 완전히 복제된 레이아웃으로 모델 가중치를 생성합니다.

레이어 가중치에 대한 레이아웃 정보를 구성하기 위해 Keras는 대부분의 내장 레이어에 대해 레이어 생성자에 추가 매개변수를 노출했습니다.

다음 예제는 완전히 복제된 가중치 레이아웃으로 작은 이미지 분류 모델을 빌드합니다. kernel_layout 및 bias_layout 인수를 통해 tf.keras.layers.Dense에서 kernel 및 bias의 레이아웃 정보를 지정할 수 있습니다. 대부분의 내장 Keras 레이어는 레이어 가중치에 대한 Layout을 명시적으로 지정할 준비가 되어 있습니다.

unsharded_layout_2d = dtensor.Layout.replicated(mesh, 2)
unsharded_layout_1d = dtensor.Layout.replicated(mesh, 1)

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, 
                        activation='relu',
                        name='d1',
                        kernel_layout=unsharded_layout_2d, 
                        bias_layout=unsharded_layout_1d),
  tf.keras.layers.Dense(10,
                        name='d2',
                        kernel_layout=unsharded_layout_2d, 
                        bias_layout=unsharded_layout_1d)
])

가중치의 layout 속성을 확인하여 레이아웃 정보를 확인할 수 있습니다.

for weight in model.weights:
  print(f'Weight name: {weight.name} with layout: {weight.layout}')
  break

Weight name: d1/kernel:0 with layout: Layout(sharding_specs=['unsharded', 'unsharded'], mesh=<Mesh object with dims=[('batch', 8)], device_type="CPU", num_local_devices=8), size=8>)

데이터세트 로드 및 입력 파이프라인 구축하기

MNIST 데이터세트를 로드하고 이에 대한 일부 전처리 입력 파이프라인을 구성합니다. 데이터세트 자체는 DTensor 레이아웃 정보와 연결되어 있지 않습니다. 향후 TensorFlow 릴리스에서 tf.data와 DTensor Keras의 통합을 개선할 계획이 있습니다.

(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label

batch_size = 128

ds_train = ds_train.map(
    normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(batch_size)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

ds_test = ds_test.map(
    normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_test = ds_test.batch(batch_size)
ds_test = ds_test.cache()
ds_test = ds_test.prefetch(tf.data.AUTOTUNE)

모델에 대한 훈련 로직 정의하기

다음으로 모델에 대한 훈련 및 평가 로직을 정의합니다.

TensorFlow 2.9부터 DTensor 지원 Keras 모델에 대한 사용자 정의 훈련 루프를 작성해야 합니다. 이것은 Keras의 표준 tf.keras.Model.fit() 또는 tf.keras.Model.eval() 함수와 통합되지 않은 적절한 레이아웃 정보로 입력 데이터를 패킹하기 위한 것입니다. 향후 릴리스에서는 더 많은 tf.data 지원이 제공될 예정입니다.

@tf.function
def train_step(model, x, y, optimizer, metrics):
  with tf.GradientTape() as tape:
    logits = model(x, training=True)
    # tf.reduce_sum sums the batch sharded per-example loss to a replicated
    # global loss (scalar).
    loss = tf.reduce_sum(tf.keras.losses.sparse_categorical_crossentropy(
        y, logits, from_logits=True))

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

  for metric in metrics.values():
    metric.update_state(y_true=y, y_pred=logits)

  loss_per_sample = loss / len(x)
  results = {'loss': loss_per_sample}
  return results

@tf.function
def eval_step(model, x, y, metrics):
  logits = model(x, training=False)
  loss = tf.reduce_sum(tf.keras.losses.sparse_categorical_crossentropy(
        y, logits, from_logits=True))

  for metric in metrics.values():
    metric.update_state(y_true=y, y_pred=logits)

  loss_per_sample = loss / len(x)
  results = {'eval_loss': loss_per_sample}
  return results

def pack_dtensor_inputs(images, labels, image_layout, label_layout):
  num_local_devices = image_layout.mesh.num_local_devices()
  images = tf.split(images, num_local_devices)
  labels = tf.split(labels, num_local_devices)
  images = dtensor.pack(images, image_layout)
  labels = dtensor.pack(labels, label_layout)
  return  images, labels

메트릭 및 옵티마이저

Keras Metric 및 Optimizer와 함께 DTensor API를 사용할 때 내부 상태 변수와 텐서가 모델의 변수와 함께 작동할 수 있도록 추가 메쉬 정보를 제공해야 합니다.

옵티마이저를 위해 DTensor는 새로운 실험적 네임스페이스인 keras.dtensor.experimental.optimizers를 도입했습니다. 여기서 기존의 많은 Keras 옵티마이저가 추가 mesh 인수를 받도록 확장됩니다. 향후 릴리스에서는 Keras 코어 옵티마이저와 병합될 수 있습니다.
메트릭의 경우 생성자에 대한 mesh를 인수로 직접 지정하여 DTensor 호환 Metric으로 만들 수 있습니다.

optimizer = tf.keras.dtensor.experimental.optimizers.Adam(0.01, mesh=mesh)
metrics = {'accuracy': tf.keras.metrics.SparseCategoricalAccuracy(mesh=mesh)}
eval_metrics = {'eval_accuracy': tf.keras.metrics.SparseCategoricalAccuracy(mesh=mesh)}

모델 훈련하기

다음 예제는 배치 차원에서 입력 파이프라인의 데이터를 샤딩하고 완전히 복제된 가중치가 있는 모델로 훈련합니다.

3개의 epoch에서 모델은 약 97%의 정확도를 달성해야 합니다.

num_epochs = 3

image_layout = dtensor.Layout.batch_sharded(mesh, 'batch', rank=4)
label_layout = dtensor.Layout.batch_sharded(mesh, 'batch', rank=1)

for epoch in range(num_epochs):
  print("============================") 
  print("Epoch: ", epoch)
  for metric in metrics.values():
    metric.reset_state()
  step = 0
  results = {}
  pbar = tf.keras.utils.Progbar(target=None, stateful_metrics=[])
  for input in ds_train:
    images, labels = input[0], input[1]
    images, labels = pack_dtensor_inputs(
        images, labels, image_layout, label_layout)

    results.update(train_step(model, images, labels, optimizer, metrics))
    for metric_name, metric in metrics.items():
      results[metric_name] = metric.result()

    pbar.update(step, values=results.items(), finalize=False)
    step += 1
  pbar.update(step, values=results.items(), finalize=True)

  for metric in eval_metrics.values():
    metric.reset_state()
  for input in ds_test:
    images, labels = input[0], input[1]
    images, labels = pack_dtensor_inputs(
        images, labels, image_layout, label_layout)
    results.update(eval_step(model, images, labels, eval_metrics))

  for metric_name, metric in eval_metrics.items():
    results[metric_name] = metric.result()

  for metric_name, metric in results.items():
    print(f"{metric_name}: {metric.numpy()}")

============================
Epoch:  0
    469/Unknown - 39s 76ms/step - loss: 0.2907 - accuracy: 0.8308
    469/Unknown - 35s 75ms/step - loss: 0.1285 - accuracy: 0.9595
    469/Unknown - 35s 75ms/step - loss: 0.1010 - accuracy: 0.9682
loss: 0.044021397829055786
accuracy: 0.9682833552360535
eval_loss: 0.05413995310664177
eval_accuracy: 0.9656000137329102

기존 모델 코드에 대한 레이아웃 지정하기

사용 사례에 잘 맞는 모델이 있는 경우가 많습니다. 모델 내의 각 개별 레이어에 Layout 정보를 지정하려면 많은 편집이 요구되므로 과중한 작업이 될 것입니다.

기존 Keras 모델을 DTensor API와 작동하도록 쉽게 변환할 수 있도록 새로운 dtensor.LayoutMap API를 사용하여 전역 관점에서 Layout을 지정할 수 있습니다.

먼저, 모델 가중치에 대해 지정하려는 모든 Layout을 포함하는 사전과 같은 객체인 LayoutMap 인스턴스를 만들어야 합니다.

LayoutMap은 초기화 시 Mesh 인스턴스가 필요하며, 이는 레이아웃이 구성되지 않은 모든 가중치에 대해 기본 복제 Layout을 제공하는 데 사용할 수 있습니다. 모든 모델 가중치를 완전히 복제하려는 경우 빈 LayoutMap을 제공하면 됩니다. 그러면 기본 메쉬가 복제된 Layout을 생성하는 데 사용됩니다.

LayoutMap은 문자열을 키로 사용하고 Layout을 값으로 사용합니다. 일반 Python dict와 이 클래스 사이에는 동작의 차이가 있습니다. 문자열 키는 값을 검색할 때 정규식으로 처리됩니다.

하위 클래스화된 모델

Keras 하위 클래스화 모델 구문을 사용하여 정의된 다음 모델을 생각해 보세요.

class SubclassedModel(tf.keras.Model):

  def __init__(self, name=None):
    super().__init__(name=name)
    self.feature = tf.keras.layers.Dense(16)
    self.feature_2 = tf.keras.layers.Dense(24)
    self.dropout = tf.keras.layers.Dropout(0.1)

  def call(self, inputs, training=None):
    x = self.feature(inputs)
    x = self.dropout(x, training=training)
    return self.feature_2(x)

이 모델에는 두 개의 Dense 레이어에 대한 kernel 및 bias인 4개의 가중치가 있습니다. 각각은 객체 경로를 기반으로 매핑됩니다.

model.feature.kernel
model.feature.bias
model.feature_2.kernel
model.feature_2.bias

참고: 하위 클래스화된 모델의 경우 레이어의 .name 속성이 아닌 속성 이름이 매핑에서 레이아웃을 검색하는 키로 사용됩니다. 이것은 tf.Module 체크포인트가 따르는 규칙과 일치합니다. 여러 레이어가 있는 복잡한 모델의 경우 체크포인트를 수동으로 검사하여 속성 매핑을 볼 수 있습니다.

이제 다음 LayoutMap을 정의하고 모델에 적용합니다.

layout_map = tf.keras.dtensor.experimental.LayoutMap(mesh=mesh)

layout_map['feature.*kernel'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=2)
layout_map['feature.*bias'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=1)

with tf.keras.dtensor.experimental.layout_map_scope(layout_map):
  subclassed_model = SubclassedModel()

WARNING:tensorflow:From /tmpfs/tmp/ipykernel_846541/1504549613.py:6: layout_map_scope (from keras.dtensor.layout_map) is deprecated and will be removed in a future version.
Instructions for updating:
use tf.keras.dtensor.experimental.LayoutMap.scope() instead.
WARNING:tensorflow:From /tmpfs/tmp/ipykernel_846541/1504549613.py:6: layout_map_scope (from keras.dtensor.layout_map) is deprecated and will be removed in a future version.
Instructions for updating:
use tf.keras.dtensor.experimental.LayoutMap.scope() instead.

모델 가중치는 첫 번째 호출에서 생성되므로 DTensor 입력으로 모델을 호출하고 가중치에 예상 레이아웃이 있는지 확인합니다.

dtensor_input = dtensor.copy_to_mesh(tf.zeros((16, 16)), layout=unsharded_layout_2d)
# Trigger the weights creation for subclass model
subclassed_model(dtensor_input)

print(subclassed_model.feature.kernel.layout)

Layout(sharding_specs=['batch', 'unsharded'], mesh=<Mesh object with dims=[('batch', 8)], device_type="CPU", num_local_devices=8), size=8>)

이를 통해 기존 코드를 업데이트하지 않고도 Layout을 모델에 빠르게 매핑할 수 있습니다.

순차 및 기능 모델

Keras 기능 및 순차 모델의 경우 LayoutMap도 사용할 수 있습니다.

참고: 기능 및 순차 모델의 경우 매핑이 약간 다릅니다. 모델의 레이어에는 모델에 연결된 공용 속성이 없습니다(model.layers를 통해 목록으로 액세스할 수는 있음). 이 경우 문자열 이름을 키로 사용합니다. 문자열 이름은 모델 내에서 고유하도록 보장됩니다.

layout_map = tf.keras.dtensor.experimental.LayoutMap(mesh=mesh)

layout_map['feature.*kernel'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=2)
layout_map['feature.*bias'] = dtensor.Layout.batch_sharded(mesh, 'batch', rank=1)

with tf.keras.dtensor.experimental.layout_map_scope(layout_map):
  inputs = tf.keras.Input((16,), batch_size=16)
  x = tf.keras.layers.Dense(16, name='feature')(inputs)
  x = tf.keras.layers.Dropout(0.1)(x)
  output = tf.keras.layers.Dense(32, name='feature_2')(x)
  model = tf.keras.Model(inputs, output)

print(model.layers[1].kernel.layout)

Layout(sharding_specs=['batch', 'unsharded'], mesh=<Mesh object with dims=[('batch', 8)], device_type="CPU", num_local_devices=8), size=8>)

with tf.keras.dtensor.experimental.layout_map_scope(layout_map):
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(16, name='feature', input_shape=(16,)),
      tf.keras.layers.Dropout(0.1),
      tf.keras.layers.Dense(32, name='feature_2')
  ])

print(model.layers[2].kernel.layout)

Layout(sharding_specs=['batch', 'unsharded'], mesh=<Mesh object with dims=[('batch', 8)], device_type="CPU", num_local_devices=8), size=8>)