Đào Tạo Phân Tán với Horovod và TF Distributed trong Handson-ML2

Đào Tạo Phân Tán với Horovod và TF Distributed trong Handson-ML2

Trong thế học máy hiện đại, việc huấn luyện mô hình trên một GPU đơn thường trở thành nút thắt cổ chai. Khi bộ dữ liệu quá lớn hoặc mô hình phức tạp, thời gian huấn luyện có thể kéo dài hàng tuần. Bài viết này sẽ giới thiệu hai giải pháp đào tạo phân tán hiệu quả từ dự án handson-ml2: Horovod và TF Distributed, giúp tăng tốc độ huấn luyện lên nhiều lần.

Nguyên Lý Đào Tạo Phân Tán

Đào tạo phân tán (Distributed Training) là phương pháp chia sẻ khối lượng công việc huấn luyện mô hình học máy trên nhiều thiết bị tính toán (GPU hoặc CPU) cùng lúc. Có hai phương pháp chính:

  • Phân Tán Dữ Liệu (Data Parallelism): Mỗi thiết bị có một bản sao mô hình và huấn luyện trên một phần dữ liệu riêng. Sau mỗi vòng lặp, các gradient được tổng hợp để cập nhật tham số mô hình.
  • Phân Tán Mô Hình (Model Parallelism): Mô hình được chia thành nhiều phần, mỗi phần được xử lý trên một thiết bị khác nhau. Phương pháp này hữu ích với các mô hình rất lớn không thể chứa trên một GPU duy nhất.

Dự án handson-ml2, một tài nguyên học tập mã nguồn mở, cung cấp hướng dẫn chi tiết về các kỹ thuật này trong notebook 19_training_and_deploying_at_scale.ipynb.

Tổng Quan về TF Distributed

TF Distributed là framework phân tán chính thức của TensorFlow, cung cấp nhiều chiến lược phân phối khác nhau:

  • MirroredStrategy: Dùng cho đào tạo trên nhiều GPU trong một máy
  • MultiWorkerMirroredStrategy: Mở rộng cho nhiều máy
  • ParameterServerStrategy: Sử dụng server tham số cho đào tạo quy mô lớn

Cấu Hình Môi Trường

Để bắt đầu với TF Distributed, hãy đảm bảo bạn đã cài đặt TensorFlow và các phụ thuộc cần thiết:

# Cài đặt TensorFlow Serving
echo "deb http://storage.googleapis.com/tensorflow-serving-apt stable tensorflow-model-server tensorflow-model-server-universal" | tee /etc/apt/sources.list.d/tensorflow-serving.list
curl https://storage.googleapis.com/tensorflow-serving-apt/tensorflow-serving.release.pub.gpg | apt-key add -
apt update && apt-get install -y tensorflow-model-server
pip install -q -U tensorflow-serving-api

Áp Dụng TF Distributed cho Mô Hình CNN

Dưới đây là ví dụ về cách sử dụng MirroredStrategy để huấn luyện mô hình phân loại hình ảnh:

import tensorflow as tf
from tensorflow.keras import layers, models

# Khởi tạo chiến lược phân phối
distribute_strategy = tf.distribute.MirroredStrategy()

print(f'Số lượng GPU có sẵn: {distribute_strategy.num_replicas_in_sync}')

with distribute_strategy.scope():
    # Xây dựng mô hình CNN
    cnn_model = models.Sequential([
        layers.Rescaling(1./255, input_shape=(150, 150, 3)),
        layers.Conv2D(32, 3, activation='relu'),
        layers.MaxPooling2D(),
        layers.Conv2D(64, 3, activation='relu'),
        layers.MaxPooling2D(),
        layers.Conv2D(128, 3, activation='relu'),
        layers.MaxPooling2D(),
        layers.Flatten(),
        layers.Dense(256, activation='relu'),
        layers.Dense(5, activation='softmax')  # Giả sử 5 lớp
    ])
    
    # Biên dịch mô hình
    cnn_model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

# Tải dữ liệu
train_dataset = tf.keras.utils.image_dataset_from_directory(
    'train_data',
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(150, 150),
    batch_size=32
)

validation_dataset = tf.keras.utils.image_dataset_from_directory(
    'train_data',
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(150, 150),
    batch_size=32
)

# Huấn luyện mô hình
history = cnn_model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=15
)

Giải Pháp Horovod

Horovod là framework phân tán được Uber phát triển, hỗ trợ TensorFlow, PyTorch và MXNet. Horovod sử dụng giao tiếp MPI (Message Passing Interface) để đồng bộ hóa giữa các tiến trình đào tạo.

Cài Đặt Horovod

# Cài đặt Horovod với hỗ trợ TensorFlow
pip install horovod[tensorflow]

Triển Khai Đào Tạo với Horovod

Dưới đây là ví dụ triển khai Horovod cho bài toán phân loại văn bản:

import tensorflow as tf
import horovod.tensorflow.keras as hvd
from tensorflow.keras import layers, models
import numpy as np

# Khởi tạo Horovod
hvd.init()

# Cấu hình GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

# Xây dựng mô hình LSTM cho phân loại văn bản
text_model = models.Sequential([
    layers.Embedding(input_dim=10000, output_dim=128),
    layers.Bidirectional(layers.LSTM(64)),
    layers.Dense(64, activation='relu'),
    layers.Dense(1, activation='sigmoid')
])

# Điều chỉnh learning rate theo số lượng quy trình
base_learning_rate = 0.001
optimizer = tf.keras.optimizers.Adam(base_learning_rate * hvd.size())
optimizer = hvd.DistributedOptimizer(optimizer)

text_model.compile(
    optimizer=optimizer,
    loss='binary_crossentropy',
    metrics=['accuracy']
)

# Tạo dữ liệu mẫu
x_train = np.random.randint(10000, size=(1000, 50))
y_train = np.random.randint(2, size=(1000,))
x_val = np.random.randint(10000, size=(200, 50))
y_val = np.random.randint(2, size=(200,))

# Định nghĩa callback
callbacks = [
    hvd.callbacks.BroadcastGlobalVariablesCallback(0),
    hvd.callbacks.MetricAverageCallback(),
    tf.keras.callbacks.TensorFlowLogger(log_dir='./logs'),
    tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', patience=3)
]

# Huấn luyện mô hình
history = text_model.fit(
    x_train, y_train,
    batch_size=32,
    epochs=10,
    validation_data=(x_val, y_val),
    callbacks=callbacks,
    verbose=1 if hvd.rank() == 0 else 0
)

Chạy Đào Tạo Phân Tán

Để khởi chạy đào tạo với Horovod trên nhiều GPU:

# Sử dụng 4 GPU trên localhost
horovodrun -np 4 -H localhost:4 python train_text_model.py

So Sánh TF Distributed và Horovod

<>Dựa trên MPI, đa framework
Đặc Điểm TF Distributed Horovod
Nền tảng Tích hợp sâu với TensorFlow
Độ phức tạp cấu hình Trung bình đến cao Thấp
Hiệu năng Tốt Xuất sắc
Hỗ trợ framework Chỉ TensorFlow TensorFlow, PyTorch, MXNet
Tích hợp với hệ sinh thái Hoàn hảo Tốt

Ứng Dụng Thực Tế: Phân Loại Hình Ảnh CIFAR-10

Triển Khai với TF Distributed

import tensorflow as tf
from tensorflow.keras import datasets, layers, models

# Tải dữ liệu CIFAR-10
(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()

# Chuẩn hóa dữ liệu
train_images, test_images = train_images / 255.0, test_images / 255.0

# Định nghĩa chiến lược phân phối
strategy = tf.distribute.MultiWorkerMirroredStrategy()

print(f'Đào tạo trên {strategy.num_replicas_in_sync} thiết bị')

with strategy.scope():
    # Xây dựng mô hình CNN
    cifar_model = models.Sequential([
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(10)
    ])
    
    # Biên dịch mô hình
    cifar_model.compile(
        optimizer='adam',
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=['accuracy']
    )

# Huấn luyện mô hình
history = cifar_model.fit(
    train_images, train_labels,
    epochs=10,
    validation_data=(test_images, test_labels)
)

Triển Khai với Horovod

import tensorflow as tf
import horovod.tensorflow.keras as hvd
from tensorflow.keras import datasets, layers, models
import numpy as np

# Khởi tạo Horovod
hvd.init()

# Cấu hình GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

# Tải dữ liệu CIFAR-10
(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()

# Chuẩn hóa dữ liệu
train_images, test_images = train_images / 255.0, test_images / 255.0

# Xây dựng mô hình CNN
cifar_model = models.Sequential([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dense(10)
])

# Điều chỉnh learning rate
optimizer = tf.keras.optimizers.Adam(0.001 * hvd.size())
optimizer = hvd.DistributedOptimizer(optimizer)

cifar_model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

# Định nghĩa callback
callbacks = [
    hvd.callbacks.BroadcastGlobalVariablesCallback(0),
    tf.keras.callbacks.ModelCheckpoint(
        filepath='cifar_model_{epoch:02d}.h5',
        save_weights_only=True,
        save_freq='epoch'),
    tf.keras.callbacks.TensorBoard(log_dir='./logs')
]

# Huấn luyện mô hình
history = cifar_model.fit(
    train_images, train_labels,
    epochs=10,
    batch_size=64,
    validation_data=(test_images, test_labels),
    callbacks=callbacks,
    verbose=1 if hvd.rank() == 0 else 0
)

Kết Luận

TF Distributed và Horovod đều là những giải pháp đào tạo phân tán mạnh mẽ cho dự án handson-ml2. TF Distributed cung cấp sự tích hợp sâu với hệ sinh thái TensorFlow, trong khi Horovod mang lại hiệu năng vượt trội và khả năng tương thích với nhiều framework khác nhau. Việc lựa chọn giữa hai giải pháp phụ thuộc vào yêu cầu cụ thể của dự án và cơ sở hạ tầng sẵn có.

Dự án handson-ml2 cung cấp nhiều ví dụ thực tế về cách triển khai các kỹ thuật này. Bằng cách áp dụng đào tạo phân tán, bạn có thể rút ngắn đáng kể thời gian huấn luyện các mô hình học máy phức tạp, từ đó tăng tốc độ phát triển và thử nghiệm ý tưởng mới.

Thẻ: tensorflow Horovod Distributed Training Machine Learning deep learning

Đăng vào ngày 6 tháng 6 lúc 02:30