Lưu ngày! Google I / O hoạt động trở lại từ ngày 18 đến 20 tháng 5 Đăng ký ngay
Trang này được dịch bởi Cloud Translation API.
Switch to English

Đào tạo phân tán với TensorFlow

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

Tổng quat

tf.distribute.Strategy là một API TensorFlow để phân phối đào tạo trên nhiều GPU, nhiều máy hoặc TPU. Sử dụng API này, bạn có thể phân phối các mô hình và mã đào tạo hiện có của mình với những thay đổi mã tối thiểu.

tf.distribute.Strategy đã được thiết kế với những mục tiêu chính sau:

  • Dễ sử dụng và hỗ trợ nhiều phân khúc người dùng, bao gồm các nhà nghiên cứu, kỹ sư ML, v.v.
  • Cung cấp hiệu suất tốt.
  • Dễ dàng chuyển đổi giữa các chiến lược.

tf.distribute.Strategy có thể được sử dụng với API cấp cao như Keras và cũng có thể được sử dụng để phân phối các vòng huấn luyện tùy chỉnh (và nói chung, bất kỳ tính toán nào sử dụng TensorFlow).

Trong TensorFlow 2.x, bạn có thể thực thi các chương trình của mình một cách háo hức hoặc trong một biểu đồ bằng cách sử dụng tf.function . tf.distribute.Strategy dự định hỗ trợ cả hai chế độ thực thi này, nhưng hoạt động tốt nhất với tf.function . tf.function . Chế độ háo hức chỉ được khuyến nghị cho mục đích gỡ lỗi và không được hỗ trợ cho TPUStrategy . Mặc dù đào tạo là trọng tâm của hướng dẫn này, API này cũng có thể được sử dụng để phân phối đánh giá và dự đoán trên các nền tảng khác nhau.

Bạn có thể sử dụng tf.distribute.Strategy với rất ít thay đổi đối với mã của mình, bởi vì chúng tôi đã thay đổi các thành phần cơ bản của TensorFlow để trở nên biết chiến lược. Điều này bao gồm các biến, lớp, mô hình, trình tối ưu hóa, số liệu, tóm tắt và điểm kiểm tra.

Trong hướng dẫn này, chúng tôi giải thích các loại chiến lược khác nhau và cách bạn có thể sử dụng chúng trong các tình huống khác nhau. Để tìm hiểu cách gỡ lỗi các vấn đề về hiệu suất, hãy xem hướng dẫn Tối ưu hóa Hiệu suất GPU TensorFlow .

# Import TensorFlow
import tensorflow as tf

Các loại chiến lược

tf.distribute.Strategy dự định sẽ bao gồm một số trường hợp sử dụng dọc theo các trục khác nhau. Một số kết hợp này hiện đang được hỗ trợ và những kết hợp khác sẽ được thêm vào trong tương lai. Một số trục này là:

  • Đào tạo đồng bộ vs không đồng bộ: Đây là hai cách phổ biến để phân phối đào tạo với song song dữ liệu. Trong đào tạo đồng bộ, tất cả nhân viên đào tạo đồng bộ các phần dữ liệu đầu vào khác nhau và tổng hợp các độ dốc ở mỗi bước. Trong đào tạo không đồng bộ, tất cả nhân viên được đào tạo độc lập về dữ liệu đầu vào và cập nhật các biến một cách không đồng bộ. Thông thường, đào tạo đồng bộ hóa được hỗ trợ thông qua tất cả giảm và không đồng bộ thông qua kiến ​​trúc máy chủ tham số.
  • Nền tảng phần cứng: Bạn có thể muốn mở rộng quy mô đào tạo của mình lên nhiều GPU trên một máy hoặc nhiều máy trong mạng (với 0 hoặc nhiều GPU mỗi máy) hoặc trên Cloud TPU.

Để hỗ trợ các trường hợp sử dụng này, có sáu chiến lược có sẵn. Phần tiếp theo giải thích cái nào trong số này được hỗ trợ trong các tình huống trong TF. Dưới đây là tổng quan nhanh:

API đào tạo MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy CentralStorageStrategy ParameterServerStrategy
API Keras Được hỗ trợ Được hỗ trợ Được hỗ trợ Hỗ trợ thử nghiệm Bài đăng được lên kế hoạch được hỗ trợ 2.4
Vòng đào tạo tùy chỉnh Được hỗ trợ Được hỗ trợ Được hỗ trợ Hỗ trợ thử nghiệm Hỗ trợ thử nghiệm
API công cụ ước tính Hỗ trợ có giới hạn Không được hỗ trợ Hỗ trợ có giới hạn Hỗ trợ có giới hạn Hỗ trợ có giới hạn

MirroredStrategy

tf.distribute.MirroredStrategy hỗ trợ đào tạo phân tán đồng bộ trên nhiều GPU trên một máy. Nó tạo ra một bản sao cho mỗi thiết bị GPU. Mỗi biến trong mô hình được sao chép trên tất cả các bản sao. Cùng với nhau, các biến này tạo thành một biến khái niệm duy nhất được gọi là MirroredVariable . Các biến này được giữ đồng bộ với nhau bằng cách áp dụng các bản cập nhật giống hệt nhau.

Các thuật toán giảm thiểu hiệu quả được sử dụng để thông báo các bản cập nhật có thể thay đổi trên các thiết bị. Giảm tất cả tổng hợp các lực căng trên tất cả các thiết bị bằng cách thêm chúng và cung cấp chúng trên mỗi thiết bị. Đó là một thuật toán hợp nhất rất hiệu quả và có thể giảm chi phí đồng bộ hóa đáng kể. Có rất nhiều thuật toán giảm thiểu và triển khai có sẵn, tùy thuộc vào loại giao tiếp có sẵn giữa các thiết bị. Theo mặc định, nó sử dụng NVIDIA NCCL làm phần thực hiện giảm toàn bộ. Bạn có thể chọn từ một số tùy chọn khác hoặc viết của riêng bạn.

Đây là cách đơn giản nhất để tạo MirroredStrategy :

mirrored_strategy = tf.distribute.MirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

Điều này sẽ tạo ra một phiên bản MirroredStrategy sẽ sử dụng tất cả các GPU hiển thị cho TensorFlow và sử dụng NCCL làm giao tiếp thiết bị chéo.

Nếu bạn chỉ muốn sử dụng một số GPU trên máy của mình, bạn có thể làm như sau:

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')

Nếu bạn muốn ghi đè giao tiếp thiết bị chéo, bạn có thể làm như vậy bằng cách sử dụng đối số cross_device_ops bằng cách cung cấp một phiên bản của tf.distribute.CrossDeviceOps . Hiện tại, tf.distribute.HierarchicalCopyAllReducetf.distribute.ReductionToOneDevice là hai tùy chọn khác với tf.distribute.NcclAllReduce là tùy chọn mặc định.

mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

TPUStrategy

tf.distribute.TPUStrategy cho phép bạn chạy chương trình đào tạo TensorFlow của mình trên các Đơn vị xử lý Tensor (TPU). TPU là ASIC chuyên dụng của Google được thiết kế để tăng tốc đáng kể khối lượng công việc học máy. Chúng có sẵn trên Google Colab, Đám mây nghiên cứu TensorFlowĐám mây TPU .

Về kiến ​​trúc đào tạo phân tán, TPUStrategy cũng giống như MirroredStrategy - nó thực hiện đào tạo phân tán đồng bộ. TPU cung cấp việc triển khai hiệu quả tất cả và các hoạt động tập thể khác trên nhiều lõi TPU, được sử dụng trong TPUStrategy .

Đây là cách bạn tạo TPUStrategy :

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)

Thể hiện TPUClusterResolver giúp định vị các TPU. Trong Colab, bạn không cần chỉ định bất kỳ đối số nào cho nó.

Nếu bạn muốn sử dụng điều này cho Cloud TPUs:

  • Bạn phải chỉ định tên tài nguyên TPU của mình trong đối số tpu .
  • Bạn phải khởi tạo hệ thống tpu một cách rõ ràng khi bắt đầu chương trình. Điều này là bắt buộc trước khi TPU có thể được sử dụng để tính toán. Việc khởi tạo hệ thống tpu cũng xóa sạch bộ nhớ TPU, vì vậy điều quan trọng là phải hoàn thành bước này trước để tránh bị mất trạng thái.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy rất giống với MirroredStrategy . Nó triển khai đào tạo phân tán đồng bộ trên nhiều nhân viên, mỗi nhân viên có nhiều GPU tiềm năng. Tương tự như tf.distribute.MirroredStrategy , nó tạo bản sao của tất cả các biến trong mô hình trên mỗi thiết bị trên tất cả các công nhân.

Đây là cách đơn giản nhất để tạo MultiWorkerMirroredStrategy :

strategy = tf.distribute.MultiWorkerMirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO

MultiWorkerMirroredStrategy có hai cách triển khai cho truyền thông thiết bị chéo. CommunicationImplementation.RING dựa trên RPC và hỗ trợ cả CPU và GPU. CommunicationImplementation.NCCL sử dụng NCCL của Nvidia và cung cấp hiệu suất hiện đại trên GPU, nhưng nó không hỗ trợ CPU. CollectiveCommunication.AUTO xác định sự lựa chọn đối với Tensorflow. Bạn có thể chỉ định chúng theo cách sau:

communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.NCCL

Một trong những điểm khác biệt chính để bắt đầu đào tạo nhiều nhân viên, so với đào tạo đa GPU, là thiết lập nhiều nhân viên. TF_CONFIG môi trường TF_CONFIG là cách tiêu chuẩn trong TensorFlow để chỉ định cấu hình cụm cho từng công nhân là một phần của cụm. Tìm hiểu thêm về cách thiết lập TF_CONFIG .

ParameterServerStrategy

Đào tạo máy chủ tham số là một phương pháp song song dữ liệu phổ biến để mở rộng quy mô đào tạo mô hình trên nhiều máy. Một cụm đào tạo máy chủ tham số bao gồm công nhân và máy chủ tham số. Các biến được tạo trên máy chủ tham số và chúng được công nhân đọc và cập nhật trong mỗi bước. Vui lòng xem hướng dẫn đào tạo máy chủ tham số để biết chi tiết.

Đào tạo máy chủ tham số TensorFlow 2 sử dụng kiến ​​trúc dựa trên điều phối viên trung tâm thông qua lớp tf.distribute.experimental.coordinator.ClusterCoordinator .

Trong việc triển khai này, parameter server tác vụ parameter serverworker chạy tf.distribute.Server lắng nghe các tác vụ từ bộ điều phối. Người điều phối tạo tài nguyên, gửi các nhiệm vụ đào tạo, viết các điểm kiểm tra và xử lý các lỗi nhiệm vụ.

Trong chương trình chạy trên bộ điều phối, bạn sẽ sử dụng một đối tượng ParameterServerStrategy để xác định một bước đào tạo và sử dụng một ClusterCoordinator để gửi các bước đào tạo đến những người làm việc từ xa. Đây là cách đơn giản nhất để tạo chúng:

strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)

Lưu ý rằng bạn sẽ cần phải định cấu hình biến môi trường TF_CONFIG nếu bạn sử dụng TFConfigClusterResolver . Nó tương tự như TF_CONFIG trong MultiWorkerMirroredStrategy nhưng có các cảnh báo bổ sung.

Trong TF 1, ParameterServerStrategy chỉ khả dụng với công cụ ước tính thông qua ký hiệu tf.compat.v1.distribute.experimental.ParameterServerStrategy .

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy cũng thực hiện đào tạo đồng bộ. Các biến không được sao chép, thay vào đó chúng được đặt trên CPU và các hoạt động được sao chép trên tất cả các GPU cục bộ. Nếu chỉ có một GPU, tất cả các biến và hoạt động sẽ được đặt trên GPU đó.

Tạo một phiên bản của CentralStorageStrategy bằng cách:

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:GPU:0'

Điều này sẽ tạo ra một phiên bản CentralStorageStrategy sẽ sử dụng tất cả các GPU và CPU có thể nhìn thấy được. Cập nhật cho các biến trên bản sao sẽ được tổng hợp trước khi áp dụng cho các biến.

Các chiến lược khác

Ngoài các chiến lược trên, có hai chiến lược khác có thể hữu ích cho việc tạo mẫu và gỡ lỗi khi sử dụng các API tf.distribute .

Chiến lược mặc định

Chiến lược mặc định là chiến lược phân phối hiện diện khi không có chiến lược phân phối rõ ràng trong phạm vi. Nó triển khai giao diện tf.distribute.Strategy nhưng là một phương thức truyền và không cung cấp phân phối thực tế. Ví dụ, strategy.run(fn) sẽ đơn giản gọi fn . Mã được viết bằng chiến lược này sẽ hoạt động chính xác như mã được viết mà không có bất kỳ chiến lược nào. Bạn có thể coi nó như một chiến lược "no-op".

Chiến lược mặc định là một chiến lược đơn lẻ - và người ta không thể tạo thêm các phiên bản của nó. Nó có thể được lấy bằng cách sử dụng tf.distribute.get_strategy() bên ngoài phạm vi của chiến lược rõ ràng bất kỳ (cùng một API có thể được sử dụng để lấy chiến lược hiện tại bên trong phạm vi của chiến lược rõ ràng).

default_strategy = tf.distribute.get_strategy()

Chiến lược này phục vụ hai mục đích chính:

  • Nó cho phép viết mã thư viện nhận biết phân phối vô điều kiện. Ví dụ: trong tf.optimizer s có thể sử dụng tf.distribute.get_strategy() và sử dụng chiến lược đó để giảm độ dốc - nó sẽ luôn trả về một đối tượng chiến lược mà chúng ta có thể gọi là API giảm.
# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values
1.0
  • Tương tự như mã thư viện, nó có thể được sử dụng để viết các chương trình của người dùng cuối để làm việc với và không có chiến lược phân phối, mà không yêu cầu logic có điều kiện. Đoạn mã mẫu minh họa điều này:
if tf.config.list_physical_devices('GPU'):
  strategy = tf.distribute.MirroredStrategy()
else:  # use default strategy
  strategy = tf.distribute.get_strategy() 

with strategy.scope():
  # do something interesting
  print(tf.Variable(1.))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>
}

OneDeviceStrategy

tf.distribute.OneDeviceStrategy là một chiến lược để đặt tất cả các biến và tính toán trên một thiết bị được chỉ định duy nhất.

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

Chiến lược này khác với chiến lược mặc định theo một số cách. Trong chiến lược mặc định, logic vị trí thay đổi vẫn không thay đổi khi so sánh với việc chạy TensorFlow mà không có bất kỳ chiến lược phân phối nào. Nhưng khi sử dụng OneDeviceStrategy , tất cả các biến được tạo trong phạm vi của nó đều được đặt rõ ràng trên thiết bị được chỉ định. Hơn nữa, bất kỳ chức năng nào được gọi qua OneDeviceStrategy.run cũng sẽ được đặt trên thiết bị được chỉ định.

Đầu vào được phân phối thông qua chiến lược này sẽ được tìm nạp trước vào thiết bị được chỉ định. Trong chiến lược mặc định, không có phân phối đầu vào.

Tương tự như chiến lược mặc định, chiến lược này cũng có thể được sử dụng để kiểm tra mã của bạn trước khi chuyển sang các chiến lược khác thực sự phân phối cho nhiều thiết bị / máy. Điều này sẽ thực hiện bộ máy chiến lược phân phối nhiều hơn một chút so với chiến lược mặc định, nhưng không ở mức độ đầy đủ như sử dụng MirroredStrategy hoặc TPUStrategy v.v. Nếu bạn muốn mã hoạt động như thể không có chiến lược, thì hãy sử dụng chiến lược mặc định.

Cho đến nay, bạn đã thấy các chiến lược khác nhau có sẵn và cách bạn có thể thực hiện chúng. Một số phần tiếp theo cho thấy các cách khác nhau mà bạn có thể sử dụng chúng để phân phối chương trình đào tạo của mình.

Sử dụng tf.distribute.Strategy với tf.keras.Model.fit

tf.distribute.Strategy được tích hợp vào tf.keras , đây là cách TensorFlow triển khai đặc tả API Keras . tf.keras là một API cấp cao để xây dựng và đào tạo các mô hình. Bằng cách tích hợp vào tf.keras phụ trợ của tf.keras , chúng tôi đã giúp bạn phân phối chương trình đào tạo của mình được viết trong khung đào tạo model.fit bằng model.fit cách model.fit .

Đây là những gì bạn cần thay đổi trong mã của mình:

  1. Tạo một phiên bản của tf.distribute.Strategy thích hợp.
  2. Di chuyển việc tạo ra các mô hình Keras, ưu và số liệu bên trong strategy.scope .

Chúng tôi hỗ trợ tất cả các loại mô hình Keras - tuần tự, chức năng và phân lớp.

Đây là một đoạn mã để thực hiện việc này cho một mô hình Keras rất đơn giản với một lớp dày đặc:

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

Ví dụ này sử dụng MirroredStrategy để bạn có thể chạy điều này trên máy có nhiều GPU. strategy.scope() lược.scope strategy.scope() chỉ ra cho Keras sử dụng chiến lược nào để phân phối khóa đào tạo. Tạo mô hình / trình tối ưu hóa / số liệu trong phạm vi này cho phép chúng tôi tạo các biến phân tán thay vì các biến thông thường. Khi điều này được thiết lập, bạn có thể điều chỉnh mô hình của mình như bình thường. MirroredStrategy quan tâm đến việc tái tạo quá trình đào tạo của mô hình trên các GPU có sẵn, tổng hợp các gradient và hơn thế nữa.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)
Epoch 1/2
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
10/10 [==============================] - 3s 2ms/step - loss: 2.3180
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 1.0246
10/10 [==============================] - 1s 2ms/step - loss: 0.5432
0.5432368516921997

Ở đây,tf.data.Dataset cung cấp đầu vào huấn luyện và đánh giá. Bạn cũng có thể sử dụng mảng numpy:

import numpy as np
inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)
Epoch 1/2
10/10 [==============================] - 1s 2ms/step - loss: 0.3866
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.1709
<tensorflow.python.keras.callbacks.History at 0x7f36bfa9fc18>

Trong cả hai trường hợp (tập dữ liệu hoặc số liệu), mỗi lô đầu vào đã cho được chia đều cho nhiều bản sao. Ví dụ: nếu sử dụng MirroredStrategy với 2 GPU, mỗi lô kích thước 10 sẽ được chia cho 2 GPU, với mỗi bộ nhận được 5 ví dụ đầu vào trong mỗi bước. Sau đó, mỗi kỷ nguyên sẽ đào tạo nhanh hơn khi bạn thêm nhiều GPU hơn. Thông thường, bạn sẽ muốn tăng kích thước lô của mình khi bạn thêm nhiều bộ tăng tốc để sử dụng hiệu quả sức mạnh tính toán bổ sung. Bạn cũng sẽ cần điều chỉnh lại tốc độ học của mình, tùy thuộc vào kiểu máy. Bạn có thể sử dụng strategy.num_replicas_in_sync để lấy số lượng bản sao.

# Compute global batch size using number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

Những gì được hỗ trợ bây giờ?

API đào tạo MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
API Keras Được hỗ trợ Được hỗ trợ Hỗ trợ thử nghiệm Hỗ trợ thử nghiệm Hỗ trợ thử nghiệm

Ví dụ và Hướng dẫn

Dưới đây là danh sách các hướng dẫn và ví dụ minh họa quá trình tích hợp ở trên kết thúc với Keras:

  1. Hướng dẫn đào tạo MNIST với MirroredStrategy .
  2. Hướng dẫn đào tạo MNIST bằng MultiWorkerMirroredStrategy .
  3. Hướng dẫn đào tạo MNIST sử dụng TPUStrategy .
  4. Hướng dẫn đào tạo máy chủ tham số trong TensorFlow 2 với ParameterServerStrategy .
  5. Kho lưu trữ TensorFlow Model Garden chứa các bộ sưu tập các mô hình hiện đại được thực hiện bằng cách sử dụng các chiến lược khác nhau.

Sử dụng tf.distribute.Strategy với các vòng huấn luyện tùy chỉnh

Như bạn đã thấy, việc sử dụng tf.distribute.Strategy với Keras model.fit yêu cầu thay đổi một vài dòng mã của bạn. Với một chút nỗ lực hơn, bạn cũng có thể sử dụng tf.distribute.Strategy với các vòng huấn luyện tùy chỉnh.

Nếu bạn cần sự linh hoạt và kiểm soát các vòng huấn luyện của mình nhiều hơn mức có thể với Công cụ ước tính hoặc Keras, bạn có thể viết các vòng huấn luyện tùy chỉnh. Ví dụ: khi sử dụng GAN, bạn có thể muốn thực hiện một số bước khác nhau của trình tạo hoặc bộ phân biệt trong mỗi vòng. Tương tự như vậy, các khung cấp cao không phù hợp lắm cho việc đào tạo Học tăng cường.

Các lớp tf.distribute.Strategy cung cấp một tập hợp các phương thức cốt lõi để hỗ trợ các vòng huấn luyện tùy chỉnh. Việc sử dụng những thứ này có thể yêu cầu cơ cấu lại mã nhỏ ban đầu, nhưng khi điều đó được thực hiện, bạn sẽ có thể chuyển đổi giữa GPU, TPU và nhiều máy chỉ bằng cách thay đổi phiên bản chiến lược.

Ở đây chúng tôi sẽ hiển thị một đoạn ngắn minh họa trường hợp sử dụng này cho một ví dụ đào tạo đơn giản sử dụng cùng một mô hình Keras như trước đây.

Đầu tiên, tạo mô hình và trình tối ưu hóa bên trong phạm vi của chiến lược. Điều này đảm bảo rằng bất kỳ biến nào được tạo bằng mô hình và trình tối ưu hóa đều là các biến được sao chép.

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  optimizer = tf.keras.optimizers.SGD()

Tiếp theo, tạo tập dữ liệu đầu vào và gọi tf.distribute.Strategy.experimental_distribute_dataset để phân phối tập dữ liệu dựa trên chiến lược.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(
    global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

Sau đó, xác định một bước của quá trình đào tạo. Sử dụngtf.GradientTape để tính toán các độ dốc và trình tối ưu hóa để áp dụng các độ dốc đó nhằm cập nhật các biến của mô hình của chúng tôi. Để phân phối bước đào tạo này, hãy đặt nó trong một hàm train_step và chuyển nó tới tf.distrbute.Strategy.run cùng với các đầu vào tập dữ liệu mà bạn nhận được từ dist_dataset đã tạo trước đó:

loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def compute_loss(labels, predictions):
  per_example_loss = loss_object(labels, predictions)
  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    loss = compute_loss(labels, predictions)

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

@tf.function
def distributed_train_step(dist_inputs):
  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

Một số điều khác cần lưu ý trong đoạn mã trên:

  1. Nó đã sử dụng tf.nn.compute_average_loss để tính toán khoản lỗ. tf.nn.compute_average_loss tổng số lỗ trên mỗi ví dụ và chia tổng cho kích thước global_batch_size. Điều này rất quan trọng vì sau này sau khi các độ dốc được tính toán trên mỗi bản sao, chúng sẽ được tổng hợp trên các bản sao bằng cách tính tổng chúng.
  2. Nó sử dụng API tf.distribute.Strategy.reduce để tổng hợp các kết quả được trả về bởi tf.distribute.Strategy.run . tf.distribute.Strategy.run trả về kết quả từ mỗi bản sao cục bộ trong chiến lược và có nhiều cách để sử dụng kết quả này. Bạn có thể reduce chúng để nhận được giá trị tổng hợp. Bạn cũng có thể thực hiện tf.distribute.Strategy.experimental_local_results để lấy danh sách các giá trị có trong kết quả, một giá trị cho mỗi bản sao cục bộ.
  3. Khi apply_gradients được gọi trong phạm vi chiến lược phân phối, hành vi của nó được sửa đổi. Cụ thể, trước khi áp dụng các gradient trên mỗi phiên bản song song trong quá trình đào tạo đồng bộ, nó thực hiện tính tổng trên tất cả các bản sao của các gradient.

Cuối cùng, khi bạn đã xác định bước huấn luyện, chúng ta có thể lặp qua dist_dataset và chạy huấn luyện trong một vòng lặp:

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))
tf.Tensor(0.7167484, shape=(), dtype=float32)
tf.Tensor(0.71152556, shape=(), dtype=float32)
tf.Tensor(0.70635474, shape=(), dtype=float32)
tf.Tensor(0.7012355, shape=(), dtype=float32)
tf.Tensor(0.6961673, shape=(), dtype=float32)
tf.Tensor(0.6911497, shape=(), dtype=float32)
tf.Tensor(0.6861821, shape=(), dtype=float32)
tf.Tensor(0.6812641, shape=(), dtype=float32)
tf.Tensor(0.6763952, shape=(), dtype=float32)
tf.Tensor(0.6715747, shape=(), dtype=float32)
tf.Tensor(0.66680235, shape=(), dtype=float32)
tf.Tensor(0.6620776, shape=(), dtype=float32)
tf.Tensor(0.6573999, shape=(), dtype=float32)
tf.Tensor(0.65276873, shape=(), dtype=float32)
tf.Tensor(0.6481838, shape=(), dtype=float32)
tf.Tensor(0.6436445, shape=(), dtype=float32)
tf.Tensor(0.6391503, shape=(), dtype=float32)
tf.Tensor(0.63470083, shape=(), dtype=float32)
tf.Tensor(0.6302957, shape=(), dtype=float32)
tf.Tensor(0.6259341, shape=(), dtype=float32)

Trong ví dụ trên, bạn đã lặp lại dist_dataset để cung cấp đầu vào cho quá trình đào tạo của mình. Chúng tôi cũng cung cấp tf.distribute.Strategy.make_experimental_numpy_dataset để hỗ trợ đầu vào numpy. Bạn có thể sử dụng API này để tạo tập dữ liệu trước khi gọi tf.distribute.Strategy.experimental_distribute_dataset .

Một cách khác để lặp lại dữ liệu của bạn là sử dụng các trình vòng lặp một cách rõ ràng. Bạn có thể muốn làm điều này khi bạn muốn chạy một số bước nhất định thay vì lặp lại trên toàn bộ tập dữ liệu. Lần lặp ở trên bây giờ sẽ được sửa đổi để đầu tiên tạo một trình lặp và sau đó gọi một cách rõ ràng next trên nó để lấy dữ liệu đầu vào.

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))
tf.Tensor(0.62161595, shape=(), dtype=float32)
tf.Tensor(0.6173406, shape=(), dtype=float32)
tf.Tensor(0.6131077, shape=(), dtype=float32)
tf.Tensor(0.60891676, shape=(), dtype=float32)
tf.Tensor(0.60476726, shape=(), dtype=float32)
tf.Tensor(0.60065883, shape=(), dtype=float32)
tf.Tensor(0.59659094, shape=(), dtype=float32)
tf.Tensor(0.5925633, shape=(), dtype=float32)
tf.Tensor(0.5885753, shape=(), dtype=float32)
tf.Tensor(0.5846267, shape=(), dtype=float32)

Điều này bao gồm trường hợp đơn giản nhất của việc sử dụng API tf.distribute.Strategy để phân phối các vòng đào tạo tùy chỉnh.

Những gì được hỗ trợ bây giờ?

API đào tạo MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
Vòng lặp đào tạo tùy chỉnh Được hỗ trợ Được hỗ trợ Hỗ trợ thử nghiệm Hỗ trợ thử nghiệm Hỗ trợ thử nghiệm

Ví dụ và Hướng dẫn

Dưới đây là một số ví dụ để sử dụng chiến lược phân phối với các vòng đào tạo tùy chỉnh:

  1. Hướng dẫn đào tạo MNIST bằng MirroredStrategy .
  2. Hướng dẫn đào tạo MNIST sử dụng TPUStrategy .
  3. Kho lưu trữ TensorFlow Model Garden chứa các bộ sưu tập các mô hình hiện đại được thực hiện bằng cách sử dụng các chiến lược khác nhau.

Các chủ đề khác

Phần này bao gồm một số chủ đề có liên quan đến nhiều trường hợp sử dụng.

Thiết lập biến môi trường TF_CONFIG

Đối với đào tạo nhiều nhân viên, như đã đề cập trước đây, bạn cần đặt biến môi trường TF_CONFIG cho mỗi tệp nhị phân đang chạy trong cụm của bạn. TF_CONFIG môi trường TF_CONFIG là một chuỗi JSON chỉ định những tác vụ nào tạo thành một cụm, địa chỉ của chúng và vai trò của từng tác vụ trong cụm. Kho lưu trữ tensorflow / hệ sinh thái cung cấp một mẫu Kubernetes trong đó đặt TF_CONFIG cho các nhiệm vụ đào tạo của bạn.

Có hai thành phần của TF_CONFIG: cụm và tác vụ. cluster cung cấp thông tin về cụm đào tạo, là một mệnh lệnh bao gồm các loại công việc khác nhau, chẳng hạn như worker. Trong đào tạo nhiều nhân viên, thường có một nhân viên đảm nhận nhiều trách nhiệm hơn một chút như lưu điểm kiểm tra và viết tệp tóm tắt cho TensorBoard ngoài những việc mà một nhân viên bình thường làm. Công nhân đó được gọi là công nhân 'trưởng' và theo thông lệ, công nhân có chỉ số 0 được chỉ định làm công nhân chính (trên thực tế đây là cách thực hiện tf.distribute.Strategy). mặt khác cung cấp thông tin của nhiệm vụ hiện tại. Cụm thành phần đầu tiên giống nhau đối với tất cả các công nhân và nhiệm vụ thành phần thứ hai là khác nhau đối với mỗi công nhân và chỉ định loại và chỉ số của công nhân đó.

Một ví dụ về TF_CONFIG là:

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})

TF_CONFIG này chỉ định rằng có ba công nhân và hai tác vụ ps trong cụm cùng với máy chủ và cổng của chúng. Phần "tác vụ" chỉ định rằng vai trò của tác vụ hiện tại trong cụm, công nhân 1 (công nhân thứ hai). Các vai trò hợp lệ trong một cụm là "trưởng", "công nhân", "ps" và "người đánh giá". Không nên có lệnh "ps" ngoại trừ khi sử dụng tf.distribute.experimental.ParameterServerStrategy .

Cái gì tiếp theo?

tf.distribute.Strategy đang được phát triển tích cực. Hãy dùng thử và cung cấp và phản hồi của bạn bằng cách sử dụng các sự cố GitHub .