Đào tạo một bác sĩ X quang AI bằng cách sử dụng Học sâu phân tán với GPU Nvidia.

Rakshith Vasudev và John Lockman – Phòng thí nghiệm đổi mới AI của Dell HPC

Khi nhiều tổ chức bắt đầu áp dụng AI trong doanh nghiệp của họ, mong muốn đào tạo các mô hình nhanh hơn sẽ chỉ tăng lên. Đây là lúc tiềm năng tăng tốc, cụ thể là từ GPU, trở nên rõ ràng. GPU cung cấp khả năng đào tạo các mô hình học sâu nhanh hơn, đôi khi theo thứ tự độ lớn, so với điện toán không tăng tốc.

Như đã giới thiệu trước đây, CheXNet là mô hình trợ lý bác sĩ X-quang AI sử dụng DenseNet để xác định tối đa 14 bệnh lý từ một hình ảnh X-quang ngực nhất định. Một số phương pháp đã được khám phá để mở rộng quy mô đào tạo một mô hình có thể hoạt động tốt bằng hoặc tốt hơn CheXNet-121 ban đầu với ResNet-50 thể hiện lời hứa về cả khả năng mở rộng và tăng độ chính xác đào tạo (AUROC tích cực). Các tác giả đã chứng minh khả năng mở rộng trên các hệ thống CPU, tuy nhiên, chúng tôi quan tâm đến việc khai thác tính song song của GPU để đẩy nhanh quá trình đào tạo. Trong bài viết này, chúng tôi mô tả cách chúng tôi có thể đạt được quá trình đào tạo CheXNet mở rộng quy mô phân tán bằng cách sử dụng GPU Nvidia V100 SXM2 trong các máy chủ Dell EMC C4140. Dell EMC PowerEdge C4140 cung cấp cả mật độ và hiệu suất với bốn GPU Nvidia V100 trong cấu hình SXM2.

Bộ dữ liệu chứa khoảng 112.000 hình ảnh chụp X-quang ngực phía trước, mỗi hình ảnh được gắn nhãn với một hoặc nhiều trong số mười bốn bệnh lý lồng ngực. Tập dữ liệu rất mất cân bằng với hơn một nửa hình ảnh tập dữ liệu không có bệnh lý được liệt kê.

Mục tiêu đào tạo CheXNet không chỉ là đào tạo mỗi lần nói. Chúng tôi muốn đẩy nhanh quá trình đào tạo từ vài ngày thành vài giờ và điều đó có thể được thực hiện với đào tạo phân tán. Chúng tôi phát hiện ra rằng việc đào tạo trên một CPU có thể mất vài ngày. Đối với dự án này, chúng tôi đã sử dụng horovod để thực hiện đào tạo phân tán trên nhiều GPU. Horovod sử dụng MPI để liên lạc.

https://eng.uber.com/horovod/

Horovod sử dụng phương pháp tiếp cận song song dữ liệu, đây là ý chính nhanh về cách thức hoạt động của nó:

1. Chạy nhiều bản sao của tập lệnh đào tạo và mỗi bản sao:
   • • a) đọc một đoạn dữ liệu
     • b) chạy nó qua mô hình
     • c) tính toán cập nhật mô hình (độ dốc)
2. Độ dốc trung bình giữa nhiều bản sao đó
3. Cập nhật mô hình
4. Lặp lại (từ Bước 1a)

Horovod sử dụng NCCL của Nvidia để cung cấp phiên bản tối ưu hóa của vòng giảm tất cả cho giao tiếp tập thể. MPI hỗ trợ cho horovod để kích hoạt giao tiếp.

Cảm hứng:

Bài báo CheXNet từ Stanford đã đề xuất sử dụng 121 lớp DensNet khi chúng cải thiện luồng thông tin và độ dốc qua mạng, làm cho việc tối ưu hóa các mạng rất sâu có thể thực hiện được. Lớp được kết nối đầy đủ được thay thế bằng lớp có một đầu vào duy nhất áp dụng phi tuyến tính sigmoid. Các trọng số được khởi tạo với trọng số đó từ một mô hình được đào tạo trước trên ImageNet. Mạng được đào tạo bằng cách sử dụng trình tối ưu hóa Adam, kích thước lô là 16. Tốc độ học bắt đầu ở 0,001 và giảm dần theo hệ số 10 mỗi khi mất xác thực ổn định sau một kỷ nguyên. Mô hình mất xác thực thấp nhất được chọn.

Dữ liệu đầu vào :

Hình ảnh X-quang ngực ban đầu từ NIH đã được chuyển đổi thành bản ghi tf. Dữ liệu đầu vào cho mô hình CheXNet là một chuỗi các bản ghi tf đã phân mảnh được phân bổ đều sao cho mỗi bản ghi tf có 256 hình ảnh. Con số này chỉ là bội số của kích thước lô mà trong trường hợp của chúng tôi là 64. Do đó, đối với mỗi lô, dữ liệu được đọc đầy đủ và không phải đọc lại cho đến khi hoàn thành.

Các bản ghi tf được chuyển đổi phải tuân theo đường dẫn dữ liệu tf giúp GPU chạy hiệu quả với thời gian nhàn rỗi rất ít, do đó làm cho quá trình đào tạo nhanh hơn. Nó tải dữ liệu từ đĩa, áp dụng các chuyển đổi được tối ưu hóa, tạo các lô và gửi dữ liệu đó tới GPU. Không có đường dẫn dữ liệu, GPU sẽ đợi CPU tải dữ liệu, dẫn đến các vấn đề về hiệu suất và thiếu dữ liệu. Bài viết này mô tả chi tiết hơn về thiết lập đường dẫn dữ liệu tf mà chúng tôi đã sử dụng.

Mô tả về mô hình:

                                             Cấu trúc liên kết DenseNet (nguồn: pytorchhub)

Chúng tôi đã sử dụng thickNet làm mô hình cơ sở được đào tạo trước với trọng số ImageNet, trình tối ưu hóa Adam được bao quanh bằng trình tối ưu hóa horovod để hỗ trợ đào tạo phân tán, kích thước lô nhỏ cục bộ là 64. Mô hình được đào tạo trước được cung cấp cho lớp tổng hợp trung bình toàn cầu nơi dữ liệu không gian được tính trung bình và được gộp lại, có đầu ra được cung cấp cho một lớp dày đặc được kết nối đầy đủ với 14 nơ-ron có chức năng kích hoạt sigmoid. Quá trình học bắt đầu ở mức 0,001 và bị giảm dần theo hệ số 10 mỗi khi các nếp gấp mất xác thực xuất hiện sau một kỷ nguyên. Mô hình đã được đào tạo trong 10 kỷ nguyên. Tệp trọng số có giá trị AUC Trung bình cao nhất được chọn.

Tăng tốc quy trình đào tạo với horovod:

Để thêm horovod, chúng tôi thực hiện các sửa đổi sau đối với mô hình tf.keras:

1. Khởi tạo môi trường MPI.
2. Phát các trọng số ngẫu nhiên ban đầu hoặc trọng số điểm kiểm tra cho tất cả công nhân,
3. Gói chức năng trình tối ưu hóa để cho phép tính tổng độ dốc nhiều nút,
4. Số liệu trung bình giữa các công nhân và
5. Hạn chế viết điểm kiểm tra cho một công nhân.

Khởi tạo môi trường MPI

Khởi tạo môi trường MPI trong Horovod chỉ yêu cầu gọi phương thức init:

import horovod.tensorflow.keras as hvd

hvd.init()

Điều này sẽ đảm bảo rằng hàm MPI_Init được gọi, thiết lập cấu trúc truyền thông và chỉ định cấp bậc cho tất cả nhân viên.

Trọng lượng phát sóng

Việc phát các trọng số nơ ron được thực hiện bằng cách gọi lại phương thức model.fit Keras. Trên thực tế, nhiều tính năng của horovod được triển khai dưới dạng gọi lại Model.fit, do đó, đáng để xác định một đối tượng danh sách gọi lại để giữ tất cả các gọi lại.

callbacks = [

hvd.callbacks.BroadcastGlobalVariablesCallback(0),

]

Bạn sẽ nhận thấy rằng BroadcastGlobalVariablesCallback nhận một đối số duy nhất được đặt thành 0. Đây là nhân viên gốc, sẽ chịu trách nhiệm đọc các tệp điểm kiểm tra hoặc tạo trọng số ban đầu mới, trọng số phát sóng khi bắt đầu chạy đào tạo và viết điểm kiểm tra định kỳ để công việc không bị mất nếu công việc đào tạo không thành công hoặc chấm dứt.

Bao hàm chức năng tối ưu hóa

Chức năng trình tối ưu hóa phải được bao bọc để nó có thể tổng hợp thông tin lỗi từ tất cả các công nhân trước khi thực hiện. Hàm DistributedOptimizer của Horovod có thể bọc bất kỳ trình tối ưu hóa nào kế thừa lớp Trình tối ưu hóa cơ sở của tf.Keras, bao gồm SGD, Adam, Adadelta, Adagrad và các loại khác.

from tensorflow.keras import optimizers

hvd_opt = hvd.DistributedOptimizer(optimizers.Adam(lr=0.001))

Trình tối ưu hóa phân tán hiện sẽ sử dụng tập hợp MPI_Allgather để tổng hợp thông tin lỗi từ các lô đào tạo cho tất cả nhân viên, thay vì chỉ thu thập chúng cho nhân viên gốc. Điều này cho phép các worker cập nhật mô hình của họ một cách độc lập thay vì đợi root phát lại các trọng số đã cập nhật trước khi bắt đầu đợt huấn luyện tiếp theo.

Số liệu trung bình

Số liệu lỗi giữa các bước cần được tính trung bình để tính tổn thất toàn cầu. Horovod cung cấp một chức năng gọi lại khác để thực hiện điều này được gọi là MetricAverageCallback.

callbacks = [ hvd.callbacks.BroadcastGlobalVariablesCallback(0),

hvd.callbacks.MetricAverageCallback()

]

Điều này sẽ đảm bảo rằng các hoạt động tối ưu hóa được thực hiện trên các chỉ số toàn cầu, chứ không phải các chỉ số cục bộ cho từng nhân viên.

Viết điểm kiểm tra từ một công nhân duy nhất

Khi sử dụng học sâu phân tán, điều quan trọng là chỉ một nhân viên ghi các tệp điểm kiểm tra để đảm bảo rằng nhiều nhân viên ghi vào cùng một tệp không tạo ra tình trạng tương tranh, điều này có thể dẫn đến lỗi điểm kiểm tra.

Tính năng ghi điểm kiểm tra trong tf.keras được bật bằng một lệnh gọi lại khác tới Model.fit. Tuy nhiên, chúng tôi chỉ muốn gọi cuộc gọi lại này từ một công nhân thay vì tất cả công nhân. Theo quy ước, chúng tôi sử dụng nhân viên 0 cho tác vụ này, nhưng về mặt kỹ thuật, chúng tôi có thể sử dụng bất kỳ nhân viên nào cho tác vụ này. Một điều tốt về worker 0 là ngay cả khi bạn quyết định chạy công việc học sâu phân tán của mình chỉ với 1 worker, thì worker đó sẽ là worker 0.

callbacks = [ ... ]

if hvd.rank() == 0:

callbacks.append(tf.keras.callbacks.ModelCheckpoint(‘./checkpoint-{epoch].h5’))

Kết quả: Đào tạo deep learning phân tán giúp cải thiện thời gian đào tạo và thông lượng Hình ảnh theo mức độ lớn.

Biểu đồ sau đây giải thích cách phân phối công việc có thể đẩy nhanh quá trình đào tạo trên GPU. Bảy thử nghiệm được hiển thị là tốc độ đào tạo của mô hình tf.keras Densenet 121 trên GPU 1 x V100, SXM2, 32GB đến GPU 16 x V100, SXM2, 32GB. Bằng cách sử dụng 4 nút C4140 có 16 GPU được phân phối, deep learning có thể cải thiện tốc độ đào tạo gấp 10,5 lần, kéo dài thời gian đào tạo cho 10 kỷ nguyên trên bộ dữ liệu ChestXray 14 từ 70 phút xuống còn 6,7 phút .

Hình 1: So sánh hiệu suất thông lượng của CheXNet với học sâu phân tán bằng cách sử dụng Horovod.

Mặc dù tỷ lệ thông lượng gần như tuyến tính, nhưng có thể nhận thấy rằng thời gian đào tạo cũng tuyến tính. Có một chút ảnh hưởng đến hiệu suất khi chúng tôi thoát khỏi nút. Tuy nhiên, sử dụng vải nhanh giúp giảm tắc nghẽn giao tiếp.

Hình 2: Thời gian đào tạo CheXNet bằng DenseNet121 với đào tạo phân tán giảm 10,5 lần với 16GPU.

Độ chính xác phân loại mô hình:

Vì CheXNet là một vấn đề phân loại nhiều nhãn và nhiều lớp và là một tập dữ liệu không cân bằng, nên số liệu chính xác để đánh giá mô hình sẽ là AUCROC. Biểu đồ phân tán sau đây hiển thị các giá trị AUC của tất cả 14 loại khác nhau được dự đoán bởi mô hình được đào tạo.

　 Hình 3: Giá trị AUC của CheXNet sử dụng DenseNet121 với đào tạo phân tán.

Phần kết luận

Trong blog này, chúng ta đã thấy cách đẩy nhanh quá trình đào tạo khi phát triển các mô hình dựa trên mạng thần kinh với học sâu phân tán. Blog này cho thấy quá trình chuyển đổi mô hình tf.keras để tận dụng nhiều nút bằng cách sử dụng khung horovod và cách một vài thay đổi mã đơn giản với một số cơ sở hạ tầng điện toán bổ sung có thể giảm thời gian cần thiết để đào tạo mô hình từ vài giờ xuống còn vài phút.