Đà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ư học máy, 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.

Bạn có thể phân phối đào tạo bằng tf.distribute.Strategy với API cấp cao như Keras Model.fit , cũng như các vòng đào tạo 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 hiện 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 . 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 tf.distribute.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, vì các thành phần cơ bản của TensorFlow đã được thay đổi để trở nên có 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, bạn sẽ tìm hiểu về 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 .

Thiết lập 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ả công nhâ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ả công nhâ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 all-Reduce và async thông qua cấu 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, TensorFlow có sẵn MirroredStrategy , TPUStrategy , MultiWorkerMirroredStrategy , ParameterServerStrategy , CentralStorageStrategy , cũng như các chiến lược khác. Phần tiếp theo giải thích cái nào trong số này được hỗ trợ trong các kịch bản nào trong TensorFlow. Dưới đây là tổng quan nhanh:

API đào tạo MirroredStrategy TPUStrategy MultiWorkerMirroredStrategy CentralStorageStrategy ParameterServerStrategy
Keras Model.fit Được hỗ trợ Được hỗ trợ Được hỗ trợ Hỗ trợ thử nghiệm Hỗ trợ thử nghiệm
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ó 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 Thư viện Liên lạc Tập thể NVIDIA ( NCCL ) làm phương thức triển khai tất cả. 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à 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:1,/job:localhost/replica:0/task:0/device:GPU:0
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 TPUĐám mây TPU .

Về kiến ​​trúc đào tạo phân tán, TPUStrategy cũng giống như MirroredStrategy — nó triển khai đào tạo phân tán đồng bộ. TPU cung cấp triển khai riêng của họ về tất cả các hoạt động giảm thiểu hiệu quả 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ể 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 cái này cho Cloud TPU:

  • 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 worker.

Đâ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: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 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: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. Biến 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 trong phần thiết lập TF_CONFIG của tài liệu này.

Để biết thêm chi tiết về MultiWorkerMirroredStrategy , hãy xem các hướng dẫn sau:

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. Xem hướng dẫn đào tạo máy chủ tham số để biết chi tiết.

Trong TensorFlow 2, đào tạo máy chủ tham số 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 triển khai này, các 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)

Để tìm hiểu thêm về ParameterServerStrategy , hãy xem đào tạo máy chủ Tham số với Keras Model.fit và hướng dẫn vòng lặp đào tạo tùy chỉnh .

Trong TensorFlow 1, ParameterServerStrategy chỉ khả dụng với Estimator thông qua biểu tượng 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à một 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 nào 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ẽ gọi đơn giản là 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 bản sao 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, bạn 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à bạn có thể gọi API Strategy.reduce .
# 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. Đây là đ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 the 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 đủ của việc sử dụng, chẳng hạn như MirroredStrategy hoặc TPUStrategy . 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 đã học về các chiến lược khác nhau và cách bạn có thể thực hiện chúng. Một số phần tiếp theo chỉ ra 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 Keras Model.fit

tf.distribute.Strategy được tích hợp vào tf.keras , là phần triển khai đặc tả Keras API của TensorFlow. 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 chương trình phụ trợ tf.keras , bạn có thể phân phối chương trình đào tạo của mình được viết trong khung đào tạo Keras bằng Model.fit một cách liền mạch .

Đâ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 mô hình Keras, trình tối ưu strategy.scope và số liệu bên trong Strategy.scope.

Các chiến lược phân phối TensorFlow hỗ trợ tất cả các loại mô hình Keras— Tuần tự , Chức năngphâ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 Dense :

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 , vì vậy bạn có thể chạy điều này trên máy có nhiều GPU. 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 bên trong phạm vi này cho phép bạn 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ể lắp 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.v.

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',).
2021-10-26 01:27:56.527729: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
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.2552
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: 0.9968
2021-10-26 01:27:59.372113: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
10/10 [==============================] - 1s 2ms/step - loss: 0.6190
0.6190494298934937

Ở đâ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
2021-10-26 01:28:00.609977: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Did not find a shardable source, walked to a node which is not a dataset: name: "FlatMapDataset/_9"
op: "FlatMapDataset"
input: "PrefetchDataset/_8"
attr {
  key: "Targuments"
  value {
    list {
    }
  }
}
attr {
  key: "f"
  value {
    func {
      name: "__inference_Dataset_flat_map_slice_batch_indices_997"
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 10
        }
      }
    }
  }
}
attr {
  key: "output_types"
  value {
    list {
      type: DT_INT64
    }
  }
}
. Consider either turning off auto-sharding or switching the auto_shard_policy to DATA to shard this dataset. You can do this by creating a new `tf.data.Options()` object then setting `options.experimental_distribute.auto_shard_policy = AutoShardPolicy.DATA` before applying the options object to the dataset via `dataset.with_options(options)`.
10/10 [==============================] - 1s 2ms/step - loss: 0.4406
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.1947
<keras.callbacks.History at 0x7fb81813d2d0>

Trong cả hai trường hợp — với Dataset hoặc NumPy — mỗi lô đầu vào đã cho được chia đều cho nhiều bản sao. Ví dụ: nếu bạn đang sử dụng MirroredStrategy với 2 GPU, mỗi lô có 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ộ gia tốc hơn, để 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 a global batch size using a 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
Keras Model.fit Được hỗ trợ Được hỗ trợ Được hỗ trợ 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 việc tích hợp end-to-end ở trên với Keras Model.fit :

  1. Hướng dẫn : Đào tạo với Model.fitMirroredStrategy .
  2. Hướng dẫn : Đào tạo với Model.fitMultiWorkerMirroredStrategy .
  3. Hướng dẫn : Chứa ví dụ về cách sử dụng Model.fitTPUStrategy .
  4. Hướng dẫn : Đào tạo máy chủ tham số với Model.fitParameterServerStrategy .
  5. Hướng dẫn : Tinh chỉnh BERT cho nhiều tác vụ từ tiêu chuẩn GLUE với Model.fitTPUStrategy .
  6. 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ư đã trình bày ở trên, việc sử dụng tf.distribute.Strategy với Keras Model.fit chỉ 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ự, 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 bộ phương pháp 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 ban đầu có thể yêu cầu cơ cấu lại mã nhỏ, nhưng sau khi hoàn tất, bạn sẽ có thể chuyển đổi giữa các GPU, TPU và nhiều máy chỉ bằng cách thay đổi phiên bản chiến lược.

Dưới đây là 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 bằng cách 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)
2021-10-26 01:28:01.831942: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:695] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}

Sau đó, xác định một bước của quá trình đào tạo. Sử dụng tf.GradientTape để tính toán độ 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 số của mô hình của bạn. Để 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.distribute.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. Bạn đã sử dụng tf.nn.compute_average_loss để tính toán khoản lỗ. tf.nn.compute_average_loss tổng tổn thất trên mỗi ví dụ và chia tổng cho global_batch_size . Điều này rất quan trọng bởi vì sau khi các độ dốc được tính toán trên mỗi bản sao, chúng được tổng hợp trên các bản sao bằng cách tổng hợp chúng.
  2. Bạn cũng đã sử dụng API tf.distribute.Strategy.reduce để tổng hợp 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 bạn gọi apply_gradients trong phạm vi chiến lược phân phối, hành vi của nó sẽ đượ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 đào tạo, bạn có thể lặp lại qua dist_dataset và chạy đào tạo trong một vòng lặp:

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))
tf.Tensor(0.18686396, shape=(), dtype=float32)
tf.Tensor(0.18628375, shape=(), dtype=float32)
tf.Tensor(0.18570684, shape=(), dtype=float32)
tf.Tensor(0.18513316, shape=(), dtype=float32)
tf.Tensor(0.1845627, shape=(), dtype=float32)
tf.Tensor(0.18399543, shape=(), dtype=float32)
tf.Tensor(0.18343134, shape=(), dtype=float32)
tf.Tensor(0.18287037, shape=(), dtype=float32)
tf.Tensor(0.18231256, shape=(), dtype=float32)
tf.Tensor(0.18175781, shape=(), dtype=float32)
tf.Tensor(0.18120615, shape=(), dtype=float32)
tf.Tensor(0.18065754, shape=(), dtype=float32)
tf.Tensor(0.18011193, shape=(), dtype=float32)
tf.Tensor(0.17956935, shape=(), dtype=float32)
tf.Tensor(0.17902976, shape=(), dtype=float32)
tf.Tensor(0.17849308, shape=(), dtype=float32)
tf.Tensor(0.17795937, shape=(), dtype=float32)
tf.Tensor(0.17742859, shape=(), dtype=float32)
tf.Tensor(0.17690066, shape=(), dtype=float32)
tf.Tensor(0.17637561, 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. Bạn cũng được 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 để tạo ra một trình lặp đầu tiên 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.17585339, shape=(), dtype=float32)
tf.Tensor(0.17533402, shape=(), dtype=float32)
tf.Tensor(0.17481743, shape=(), dtype=float32)
tf.Tensor(0.17430364, shape=(), dtype=float32)
tf.Tensor(0.17379259, shape=(), dtype=float32)
tf.Tensor(0.17328428, shape=(), dtype=float32)
tf.Tensor(0.17277871, shape=(), dtype=float32)
tf.Tensor(0.17227581, shape=(), dtype=float32)
tf.Tensor(0.17177561, shape=(), dtype=float32)
tf.Tensor(0.17127804, 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 đà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

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 huấn luyện tùy chỉnh:

  1. Hướng dẫn : Đào tạo với vòng lặp đào tạo tùy chỉnh và MirroredStrategy .
  2. Hướng dẫn : Đào tạo với vòng lặp đào tạo tùy chỉnh và MultiWorkerMirroredStrategy .
  3. Hướng dẫn : Chứa ví dụ về vòng lặp đào tạo tùy chỉnh với TPUStrategy .
  4. Hướng dẫn : Đào tạo máy chủ tham số với vòng lặp đào tạo tùy chỉnh và 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.

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 thiết lập 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. Biến 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. tensorflow/ecosystem cung cấp mẫu Kubernetes, mẫu này thiết lập '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' : một cụm và một tác vụ.

  • Một cụm 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ư công nhân. 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 tf.distribute.Strategy được triển khai).
  • Mặt khác, một nhiệm vụ cung cấp thông tin về 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 "cluster" cùng với máy chủ và cổng của chúng. Phần "task" chỉ định vai trò của nhiệm vụ hiện tại trong "cluster" —worker 1 (worker thứ hai). Các vai trò hợp lệ trong một cụm là "chief" , "worker" , "ps""evaluator" . 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 phản hồi của bạn bằng các sự cố GitHub .