12.6. Cách lập trình Súc tích đa GPU

Lập trình từ đầu việc song song hóa cho từng mô hình mới khá mất công. Hơn nữa, việc tối ưu các công cụ đồng bộ hóa sẽ cho hiệu suất cao. Sau đây chúng tôi sẽ giới thiệu cách thực hiện điều này bằng Gluon. Phần lý thuyết toán và các thuật toán giống trong Section 12.5. Như trước đây, ta bắt đầu bằng cách nhập các mô-đun cần thiết (tất nhiên là ta sẽ cần ít nhất hai GPU để chạy notebook này).

from d2l import mxnet as d2l
from mxnet import autograd, gluon, init, np, npx
from mxnet.gluon import nn
npx.set_np()

12.6.1. Ví dụ Đơn giản

Hãy sử dụng một mạng có ý nghĩa hơn một chút so với LeNet ở phần trước mà vẫn có thể huấn luyện dễ dàng và nhanh chóng. Chúng tôi chọn một biến thể của ResNet-18 [He et al., 2016a]. Vì hình ảnh đầu vào rất nhỏ nên ta sửa đổi nó một chút. Cụ thể, điểm khác biệt so với ở Section 7.6 là ở phần đầu, ta sử dụng hạt nhân tích chập có kích thước, sải bước và đệm nhỏ hơn, và cũng loại bỏ đi tầng gộp cực đại.

#@save
def resnet18(num_classes):
    """A slightly modified ResNet-18 model."""
    def resnet_block(num_channels, num_residuals, first_block=False):
        blk = nn.Sequential()
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.add(d2l.Residual(
                    num_channels, use_1x1conv=True, strides=2))
            else:
                blk.add(d2l.Residual(num_channels))
        return blk

    net = nn.Sequential()
    # This model uses a smaller convolution kernel, stride, and padding and
    # removes the maximum pooling layer
    net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1),
            nn.BatchNorm(), nn.Activation('relu'))
    net.add(resnet_block(64, 2, first_block=True),
            resnet_block(128, 2),
            resnet_block(256, 2),
            resnet_block(512, 2))
    net.add(nn.GlobalAvgPool2D(), nn.Dense(num_classes))
    return net

12.6.2. Khởi tạo Tham số và Công việc phụ trợ

Phương thức initialize cho phép ta thiết lập giá trị mặc định ban đầu cho các tham số trên thiết bị được chọn. Với độc giả mới, có thể tham khảo Section 4.8. Một điều rất thuận tiện là nó cũng cho phép ta khởi tạo mạng trên nhiều thiết bị cùng một lúc. Hãy thử xem cách nó hoạt động trong thực tế.

net = resnet18(10)
# get a list of GPUs
ctx = d2l.try_all_gpus()
# initialize the network on all of them
net.initialize(init=init.Normal(sigma=0.01), ctx=ctx)

Sử dụng hàm split_and_load được giới thiệu trong phần trước, chúng ta có thể phân chia một minibatch dữ liệu và sao chép các phần dữ liệu vào danh sách các thiết bị được cung cấp bởi biến ngữ cảnh. Mạng sẽ tự động sử dụng GPU thích hợp để tính giá trị của lượt truyền xuôi. Ta tạo ra 4 mẫu dữ liệu và phân chia chúng trên các GPU như trước đây.

x = np.random.uniform(size=(4, 1, 28, 28))
x_shards = gluon.utils.split_and_load(x, ctx)
net(x_shards[0]), net(x_shards[1])

(array([[ 2.2610202e-06,  2.2045995e-06, -5.4046805e-06,  1.2869957e-06,
          5.1373154e-06, -3.8297981e-06,  1.4339025e-07,  5.4683442e-06,
         -2.8279189e-06, -3.9651104e-06],
        [ 2.0698674e-06,  2.0084674e-06, -5.6382496e-06,  1.0498472e-06,
          5.5506421e-06, -4.1065487e-06,  6.0830212e-07,  5.4521779e-06,
         -3.7365019e-06, -4.1891653e-06]], ctx=gpu(0)),
 array([[ 2.4629783e-06,  2.6015541e-06, -5.4362622e-06,  1.2938226e-06,
          5.6387880e-06, -4.1360095e-06,  3.5758796e-07,  5.5125233e-06,
         -3.1957338e-06, -4.2976326e-06],
        [ 1.9431684e-06,  2.2600436e-06, -5.2698174e-06,  1.4807399e-06,
          5.4830948e-06, -3.9678889e-06,  7.5751359e-08,  5.6764356e-06,
         -3.2530243e-06, -4.0943960e-06]], ctx=gpu(1)))

Khi dữ liệu được truyền qua mạng, các tham số tương ứng sẽ được khởi tạo trên thiết bị mà dữ liệu được truyền qua. Điều này có nghĩa là việc khởi tạo xảy ra theo từng thiết bị. Do ta lựa chọn việc khởi tạo trên GPU 0 và GPU 1, mạng chỉ được khởi tạo trên hai thiết bị này chứ trên CPU thì không. Trong thực tế, các tham số này thậm chí còn không tồn tại trên CPU. Ta có thể kiểm chứng điều này bằng cách in các tham số ra và theo dõi xem liệu có lỗi nào xảy ra hay không.

weight = net[0].params.get('weight')

try:
    weight.data()
except RuntimeError:
    print('not initialized on cpu')
weight.data(ctx[0])[0], weight.data(ctx[1])[0]

not initialized on cpu

(array([[[ 0.01382882, -0.01183044,  0.01417866],
         [-0.00319718,  0.00439528,  0.02562625],
         [-0.00835081,  0.01387452, -0.01035946]]], ctx=gpu(0)),
 array([[[ 0.01382882, -0.01183044,  0.01417866],
         [-0.00319718,  0.00439528,  0.02562625],
         [-0.00835081,  0.01387452, -0.01035946]]], ctx=gpu(1)))

Cuối cùng, hãy cùng thay đổi đoạn mã đánh giá độ chính xác để có thể chạy song song trên nhiều thiết bị. Hàm này được viết lại từ hàm evaluate_accuracy_gpu ở Section 6.6. Điểm khác biệt lớn nhất nằm ở việc ta tách một batch ra trước khi truyền vào mạng. Các phần còn lại gần như là giống hệt.

#@save
def evaluate_accuracy_gpus(net, data_iter, split_f=d2l.split_batch):
    # Query the list of devices
    ctx = list(net.collect_params().values())[0].list_ctx()
    metric = d2l.Accumulator(2)  # num_corrected_examples, num_examples
    for features, labels in data_iter:
        X_shards, y_shards = split_f(features, labels, ctx)
        # Run in parallel
        pred_shards = [net(X_shard) for X_shard in X_shards]
        metric.add(sum(float(d2l.accuracy(pred_shard, y_shard)) for
                       pred_shard, y_shard in zip(
                           pred_shards, y_shards)), labels.size)
    return metric[0] / metric[1]

12.6.3. Huấn luyện

Như phần trên, đoạn mã huấn luyện cần thực hiện một số hàm cơ bản để quá trình song song hóa đạt hiệu quả:

  • Các tham số của mạng cần được khởi tạo trên tất cả các thiết bị.
  • Trong suốt quá trình lặp trên tập dữ liệu, các minibatch được chia nhỏ cho tất cả các thiết bị.
  • Ta tính toán song song hàm mất mát và gradient của nó trên tất cả các thiết bị.
  • Mất mát được tích luỹ (bởi phương thức huấn luyện trainer) và các tham số được cập nhật tương ứng.

Cuối cùng ta tính toán (vẫn song song) độ chính xác và báo cáo giá trị cuối cùng của mạng. Quá trình huấn luyện ở đây khá giống với chương trước, trừ việc ta cần chia nhỏ và tổng hợp lại dữ liệu.

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    ctx = [d2l.try_gpu(i) for i in range(num_gpus)]
    net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)
    trainer = gluon.Trainer(net.collect_params(), 'sgd',
                            {'learning_rate': lr})
    loss = gluon.loss.SoftmaxCrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        timer.start()
        for features, labels in train_iter:
            X_shards, y_shards = d2l.split_batch(features, labels, ctx)
            with autograd.record():
                losses = [loss(net(X_shard), y_shard) for X_shard, y_shard
                          in zip(X_shards, y_shards)]
            for l in losses:
                l.backward()
            trainer.step(batch_size)
        npx.waitall()
        timer.stop()
        animator.add(epoch + 1, (evaluate_accuracy_gpus(net, test_iter),))
    print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
          f'on {str(ctx)}')

12.6.4. Thử nghiệm

Hãy cùng xem cách hoạt động trong thực tế. Để khởi động, ta huấn luyện mạng này trên một GPU đơn.

train(num_gpus=1, batch_size=256, lr=0.1)

test acc: 0.93, 31.8 sec/epoch on [gpu(0)]
../_images/output_multiple-gpus-concise_vn_9b044d_15_1.svg

Tiếp theo, ta sử dụng 2 GPU để huấn luyện. Mô hình ResNet-18 phức tạp hơn đáng kể so với LeNet. Đây chính là cơ hội để song song hóa chứng tỏ lợi thế của nó, vì thời gian dành cho việc tính toán lớn hơn đáng kể so với thời gian đồng bộ hóa các tham số. Điều này giúp cải thiện khả năng mở rộng do tổng chi phí song song hóa không quá đáng kể.

train(num_gpus=2, batch_size=512, lr=0.2)

test acc: 0.92, 33.9 sec/epoch on [gpu(0), gpu(1)]
../_images/output_multiple-gpus-concise_vn_9b044d_17_1.svg

12.6.5. Tóm tắt

  • Gluon cung cấp các hàm để khởi tạo mô hình trên nhiều thiết bị bằng cách cung cấp một danh sách ngữ cảnh.
  • Dữ liệu được tự động đánh giá trên các thiết bị mà nó được lưu trữ.
  • Chú ý việc khởi tạo mạng trên mỗi thiết bị trước khi thử truy cập vào các tham số trên thiết bị đó. Nếu không khả năng cao sẽ có lỗi xảy ra.
  • Các thuật toán tối ưu tự động tổng hợp kết quả trên nhiều GPU.

12.6.6. Bài tập

  1. Phần này ta sử dụng ResNet-18. Hãy thử với số epoch, kích thước batch và tốc độ học khác. Thử sử dụng nhiều GPU hơn để tính toán. Chuyện gì sẽ xảy ra nếu ta chạy mô hình này trên máy chủ p2.16xlarge với 16 GPU?
  2. Đôi khi mỗi thiết bị khác nhau cung cấp khả năng tính toán khác nhau. Ta có thể sử dụng GPU và CPU cùng lúc. Vậy ta nên phân chia công việc thế nào? Liệu việc phân chia có đáng hay không? Tại sao?
  3. Chuyện gì sẽ xảy ra nếu ta bỏ hàm npx.waitall()? Bạn sẽ thay đổi quá trình huấn luyện thế nào để có thể xử lý song song tối đa 2 bước cùng lúc?

12.6.7. Thảo luận

12.6.8. Những người thực hiện

Bản dịch trong trang này được thực hiện bởi:

  • Đoàn Võ Duy Thanh
  • Trần Yến Thy
  • Lê Khắc Hồng Phúc
  • Nguyễn Văn Cường
  • Đỗ Trường Giang
  • Nguyễn Lê Quang Nhật
  • Phạm Hồng Vinh