Sử dụng GPU

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

Mã TensorFlow và các mô hình tf.keras sẽ chạy rõ ràng trên một GPU duy nhất mà không cần thay đổi mã.

Cách đơn giản nhất để chạy trên nhiều GPU, trên một hoặc nhiều máy, là sử dụng Chiến lược phân phối .

Hướng dẫn này dành cho những người dùng đã thử các phương pháp này và nhận thấy rằng họ cần kiểm soát chi tiết cách TensorFlow sử dụng GPU. Để tìm hiểu cách gỡ lỗi các vấn đề về hiệu suất cho các tình huống đơn và đa GPU, hãy xem hướng dẫn Tối ưu hóa hiệu suất GPU TensorFlow .

Thành lập

Đảm bảo bạn đã cài đặt bản phát hành gpu TensorFlow mới nhất.

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
Num GPUs Available:  1

Tổng quat

TensorFlow hỗ trợ chạy tính toán trên nhiều loại thiết bị, bao gồm cả CPU và GPU. Chúng được biểu diễn bằng các số nhận dạng chuỗi, ví dụ:

  • "/device:CPU:0" : CPU của máy tính của bạn.
  • "/GPU:0" : Ký hiệu ngắn gọn cho GPU đầu tiên trên máy của bạn hiển thị với TensorFlow.
  • "/job:localhost/replica:0/task:0/device:GPU:1" : Tên đủ điều kiện của GPU thứ hai trên máy của bạn hiển thị với TensorFlow.

Nếu một hoạt động TensorFlow có triển khai cả CPU và GPU, theo mặc định, thiết bị GPU sẽ được ưu tiên khi hoạt động được chỉ định. Ví dụ: tf.matmul có cả nhân CPU và GPU và trên một hệ thống có thiết bị CPU:0GPU:0 , thiết bị GPU:0 được chọn để chạy tf.matmul trừ khi bạn yêu cầu rõ ràng chạy nó trên thiết bị khác.

Nếu một hoạt động TensorFlow không có triển khai GPU tương ứng, thì hoạt động sẽ trở lại thiết bị CPU. Ví dụ: vì tf.cast chỉ có nhân CPU nên trên hệ thống có thiết bị CPU:0GPU:0 , thiết bị CPU:0 được chọn để chạy tf.cast , ngay cả khi được yêu cầu chạy trên thiết bị GPU:0 .

Ghi nhật ký vị trí thiết bị

Để biết các thiết bị hoạt động và tensors của bạn được chỉ định, hãy đặt tf.debugging.set_log_device_placement(True) làm câu lệnh đầu tiên của chương trình của bạn. Bật ghi nhật ký vị trí thiết bị sẽ in bất kỳ phân bổ hoặc hoạt động Tensor nào.

tf.debugging.set_log_device_placement(True)

# Create some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Đoạn mã trên sẽ in ra một chỉ báo op MatMul đã được thực thi trên GPU:0 .

Vị trí thiết bị thủ công

Nếu bạn muốn một hoạt động cụ thể chạy trên một thiết bị bạn chọn thay vì những gì được chọn tự động cho bạn, bạn có thể sử dụng with tf.device để tạo ngữ cảnh thiết bị và tất cả các hoạt động trong ngữ cảnh đó sẽ chạy trên cùng một thiết bị được chỉ định .

tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Bạn sẽ thấy rằng bây giờ ab được gán cho CPU:0 . Vì một thiết bị không được chỉ định rõ ràng cho hoạt động MatMul , thời gian chạy TensorFlow sẽ chọn một thiết bị dựa trên hoạt động và các thiết bị có sẵn (trong ví dụ này là GPU:0 ) và tự động sao chép các bộ căng giữa các thiết bị nếu được yêu cầu.

Hạn chế tăng trưởng bộ nhớ GPU

Theo mặc định, TensorFlow ánh xạ gần như tất cả bộ nhớ GPU của tất cả các GPU (tùy thuộc vào CUDA_VISIBLE_DEVICES ) hiển thị cho quá trình. Điều này được thực hiện để sử dụng hiệu quả hơn tài nguyên bộ nhớ GPU tương đối quý giá trên thiết bị bằng cách giảm phân mảnh bộ nhớ. Để giới hạn TensorFlow trong một tập hợp GPU cụ thể, hãy sử dụng phương pháp tf.config.set_visible_devices .

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
  except RuntimeError as e:
    # Visible devices must be set before GPUs have been initialized
    print(e)
1 Physical GPUs, 1 Logical GPU

Trong một số trường hợp, quá trình mong muốn chỉ cấp phát một tập hợp con của bộ nhớ khả dụng hoặc chỉ tăng mức sử dụng bộ nhớ khi tiến trình cần. TensorFlow cung cấp hai phương pháp để kiểm soát điều này.

Tùy chọn đầu tiên là bật tăng trưởng bộ nhớ bằng cách gọi tf.config.experimental.set_memory_growth , cố gắng chỉ phân bổ nhiều bộ nhớ GPU nhất cần thiết cho phân bổ thời gian chạy: nó bắt đầu phân bổ rất ít bộ nhớ và khi chương trình được chạy và cần thêm bộ nhớ GPU, vùng bộ nhớ GPU được mở rộng cho quá trình TensorFlow. Bộ nhớ không được giải phóng vì nó có thể dẫn đến phân mảnh bộ nhớ. Để bật tăng trưởng bộ nhớ cho một GPU cụ thể, hãy sử dụng mã sau trước khi phân bổ bất kỳ bộ căng nào hoặc thực hiện bất kỳ hoạt động nào.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)
Physical devices cannot be modified after being initialized

Một cách khác để bật tùy chọn này là đặt biến môi trường TF_FORCE_GPU_ALLOW_GROWTH thành true . Cấu hình này dành riêng cho nền tảng.

Phương pháp thứ hai là định cấu hình thiết bị GPU ảo với tf.config.set_logical_device_configuration và đặt giới hạn cứng trên tổng bộ nhớ để phân bổ trên GPU.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
Virtual devices cannot be modified after being initialized

Điều này rất hữu ích nếu bạn muốn thực sự ràng buộc dung lượng bộ nhớ GPU có sẵn với quy trình TensorFlow. Đây là thực tế phổ biến đối với phát triển cục bộ khi GPU được chia sẻ với các ứng dụng khác như GUI máy trạm.

Sử dụng một GPU duy nhất trên hệ thống đa GPU

Nếu bạn có nhiều GPU trong hệ thống của mình, GPU có ID thấp nhất sẽ được chọn theo mặc định. Nếu bạn muốn chạy trên một GPU khác, bạn sẽ cần chỉ định tùy chọn một cách rõ ràng:

tf.debugging.set_log_device_placement(True)

try:
  # Specify an invalid GPU device
  with tf.device('/device:GPU:2'):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    c = tf.matmul(a, b)
except RuntimeError as e:
  print(e)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0

Nếu thiết bị bạn đã chỉ định không tồn tại, bạn sẽ nhận được RuntimeError : .../device:GPU:2 unknown device .

Nếu bạn muốn TensorFlow tự động chọn một thiết bị hiện có và được hỗ trợ để chạy các hoạt động trong trường hợp thiết bị được chỉ định không tồn tại, bạn có thể gọi tf.config.set_soft_device_placement(True) .

tf.config.set_soft_device_placement(True)
tf.debugging.set_log_device_placement(True)

# Creates some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Sử dụng nhiều GPU

Việc phát triển cho nhiều GPU sẽ cho phép một mô hình mở rộng quy mô với các tài nguyên bổ sung. Nếu phát triển trên hệ thống có một GPU, bạn có thể mô phỏng nhiều GPU bằng các thiết bị ảo. Điều này cho phép dễ dàng kiểm tra các thiết lập đa GPU mà không yêu cầu thêm tài nguyên.

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Create 2 virtual GPUs with 1GB memory each
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024),
         tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)
Virtual devices cannot be modified after being initialized

Khi có nhiều GPU logic có sẵn cho thời gian chạy, bạn có thể sử dụng nhiều GPU với tf.distribute.Strategy hoặc với vị trí thủ công.

Với tf.distribute.Strategy

Cách tốt nhất để sử dụng nhiều GPU là sử dụng tf.distribute.Strategy . Đây là một ví dụ đơn giản:

tf.debugging.set_log_device_placement(True)
gpus = tf.config.list_logical_devices('GPU')
strategy = tf.distribute.MirroredStrategy(gpus)
with strategy.scope():
  inputs = tf.keras.layers.Input(shape=(1,))
  predictions = tf.keras.layers.Dense(1)(inputs)
  model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
  model.compile(loss='mse',
                optimizer=tf.keras.optimizers.SGD(learning_rate=0.2))
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op RandomUniform in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Sub in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /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',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /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',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /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',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /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',).
Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op Identity in device /job:localhost/replica:0/task:0/device:CPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op NoOp in device /job:localhost/replica:0/task:0/device:GPU:0

Chương trình này sẽ chạy một bản sao mô hình của bạn trên mỗi GPU, tách dữ liệu đầu vào giữa chúng, còn được gọi là " dữ liệu song song ".

Để biết thêm thông tin về các chiến lược phân phối, hãy xem hướng dẫn tại đây .

Vị trí thủ công

tf.distribute.Strategy hoạt động ngầm bằng cách sao chép tính toán trên các thiết bị. Bạn có thể triển khai nhân rộng theo cách thủ công bằng cách xây dựng mô hình của mình trên mỗi GPU. Ví dụ:

tf.debugging.set_log_device_placement(True)

gpus = tf.config.list_logical_devices('GPU')
if gpus:
  # Replicate your computation on multiple GPUs
  c = []
  for gpu in gpus:
    with tf.device(gpu.name):
      a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
      c.append(tf.matmul(a, b))

  with tf.device('/CPU:0'):
    matmul_sum = tf.add_n(c)

  print(matmul_sum)
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)