Tối ưu hóa đồ thị TensorFlow với Grappler

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

TensorFlow sử dụng cả biểu đồ và thực thi háo hức để thực thi các phép tính. Một tf.Graph chứa một tập hợp các đối tượng tf.Operation (ops) đại diện cho các đơn vị tính toán và các đối tượng tf.Tensor đại diện cho các đơn vị dữ liệu lưu chuyển giữa các ops.

Grappler là hệ thống tối ưu hóa đồ thị mặc định trong thời gian chạy TensorFlow. Grappler áp dụng tối ưu hóa trong chế độ đồ thị (trong tf.function . function) để cải thiện hiệu suất của các tính toán TensorFlow của bạn thông qua đơn giản hóa đồ thị và các tối ưu hóa cấp cao khác như nội tuyến các cơ quan chức năng để cho phép tối ưu hóa liên thủ tục. Tối ưu hóa tf.Graph cũng làm giảm mức sử dụng bộ nhớ tối đa của thiết bị và cải thiện việc sử dụng phần cứng bằng cách tối ưu hóa ánh xạ các nút đồ thị để tính toán tài nguyên.

Sử dụng tf.config.optimizer.set_experimental_options() để kiểm soát tốt hơn các tối ưu hóa tf.Graph của bạn.

Trình tối ưu hóa đồ thị có sẵn

Grappler thực hiện tối ưu hóa đồ thị thông qua trình điều khiển cấp cao nhất được gọi là MetaOptimizer . Các trình tối ưu hóa đồ thị sau có sẵn với TensorFlow:

  • Trình tối ưu hóa gấp liên tục - Suy ra giá trị của tensors khi có thể bằng cách gấp các nút không đổi trong biểu đồ và hiện thực hóa kết quả bằng cách sử dụng hằng số.
  • Trình tối ưu hóa số học - Đơn giản hóa các phép toán số học bằng cách loại bỏ các biểu thức con phổ biến và đơn giản hóa các câu lệnh số học.
  • Trình tối ưu hóa bố cục - Tối ưu hóa bố cục tensor để thực hiện các hoạt động phụ thuộc vào định dạng dữ liệu chẳng hạn như kết hợp hiệu quả hơn.
  • Trình tối ưu hóa Remapper - Bản đồ lại các đồ thị con thành các triển khai hiệu quả hơn bằng cách thay thế các đồ thị con thường xảy ra bằng các hạt nhân nguyên khối hợp nhất được tối ưu hóa.
  • Trình tối ưu hóa bộ nhớ - Phân tích biểu đồ để kiểm tra mức sử dụng bộ nhớ cao nhất cho từng hoạt động và chèn các hoạt động sao chép bộ nhớ CPU-GPU để hoán đổi bộ nhớ GPU sang CPU nhằm giảm mức sử dụng bộ nhớ tối đa.
  • Trình tối ưu hóa phụ thuộc - Loại bỏ hoặc sắp xếp lại các phụ thuộc điều khiển để rút ngắn đường dẫn quan trọng cho một bước mô hình hoặc cho phép các tối ưu hóa khác. Cũng loại bỏ các nút không hoạt động hiệu quả như Identity.
  • Trình tối ưu hóa cắt tỉa - Cắt tỉa các nút không ảnh hưởng đến kết quả đầu ra từ biểu đồ. Nó thường được chạy đầu tiên để giảm kích thước của đồ thị và tăng tốc độ xử lý trong các đường chuyền Grappler khác.
  • Trình tối ưu hóa chức năng - Tối ưu hóa thư viện chức năng của chương trình TensorFlow và nội tuyến các thân chức năng để cho phép tối ưu hóa liên thủ tục khác.
  • Trình tối ưu hóa hình dạng - Tối ưu hóa các đồ thị con hoạt động dựa trên thông tin liên quan đến hình dạng và hình dạng.
  • Trình tối ưu hóa tự động song song - Tự động song song hóa các biểu đồ bằng cách tách dọc theo thứ nguyên hàng loạt. Trình tối ưu hóa này được TẮT theo mặc định.
  • Trình tối ưu hóa vòng lặp - Tối ưu hóa luồng điều khiển đồ thị bằng cách nâng các đồ thị con bất biến của vòng lặp ra khỏi vòng lặp và bằng cách loại bỏ các hoạt động ngăn xếp dư thừa trong các vòng lặp. Cũng tối ưu hóa các vòng lặp với số lượng chuyến đi tĩnh đã biết và loại bỏ các nhánh chết đã biết tĩnh trong các điều kiện.
  • Trình tối ưu hóa trình phân bổ theo phạm vi - Giới thiệu trình phân bổ theo phạm vi để giảm di chuyển dữ liệu và hợp nhất một số hoạt động.
  • Ghim vào trình tối ưu hóa máy chủ - Hoán đổi các hoạt động nhỏ trên CPU. Trình tối ưu hóa này được TẮT theo mặc định.
  • Tự động tối ưu hóa độ chính xác hỗn hợp - Chuyển đổi kiểu dữ liệu thành float16 nếu có thể để cải thiện hiệu suất. Hiện chỉ áp dụng cho GPU.
  • Gỡ lỗi thoát y - Tách các nút liên quan đến hoạt động gỡ lỗi như tf.debugging.Assert , tf.debugging.check_numericstf.print khỏi biểu đồ. Trình tối ưu hóa này được TẮT theo mặc định.

Thành lập

import numpy as np
import timeit
import traceback
import contextlib


import tensorflow as tf

Tạo trình quản lý ngữ cảnh để dễ dàng chuyển đổi trạng thái của trình tối ưu hóa.

@contextlib.contextmanager
def options(options):
  old_opts = tf.config.optimizer.get_experimental_options()
  tf.config.optimizer.set_experimental_options(options)
  try:
    yield
  finally:
    tf.config.optimizer.set_experimental_options(old_opts)

So sánh hiệu suất thực thi có và không có Grappler

TensorFlow 2 trở lên thực thi theo mặc định. Sử dụng tf.function để chuyển việc thực thi mặc định sang chế độ Đồ thị. Grappler chạy tự động trong nền để áp dụng các tối ưu hóa đồ thị ở trên và cải thiện hiệu suất thực thi.

Trình tối ưu hóa gấp liên tục

Ví dụ sơ bộ, hãy xem xét một hàm thực hiện các hoạt động trên hằng số và trả về một đầu ra.

def test_function_1():
  @tf.function
  def simple_function(input_arg):
    print('Tracing!')
    a = tf.constant(np.random.randn(2000,2000), dtype = tf.float32)
    c = a
    for n in range(50):
      c = c@a
    return tf.reduce_mean(c+input_arg)

  return simple_function

Tắt trình tối ưu hóa gấp liên tục và thực hiện chức năng:

with options({'constant_folding': False}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Vanilla execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")
{'constant_folding': False, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Vanilla execution: 0.0018392090000816097 s

Bật trình tối ưu hóa gấp liên tục và thực thi lại chức năng để quan sát tốc độ thực thi chức năng.

with options({'constant_folding': True}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Constant folded execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")
{'constant_folding': True, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Constant folded execution: 0.0006749789999958011 s

Gỡ lỗi trình tối ưu hóa vũ nữ thoát y

Hãy xem xét một hàm đơn giản kiểm tra giá trị số của đối số đầu vào của nó và trả về nó.

def test_function_2():
  @tf.function
  def simple_func(input_arg):
    output = input_arg
    tf.debugging.check_numerics(output, "Bad!")
    return output
  return simple_func

Đầu tiên, thực thi chức năng với trình tối ưu hóa bộ thoát lỗi đã tắt.

test_func = test_function_2()
p1 = tf.constant(float('inf'))
try:
  test_func(p1)
except tf.errors.InvalidArgumentError as e:
  traceback.print_exc(limit=2)
2021-09-22 20:34:55.871238: E tensorflow/core/kernels/check_numerics_op.cc:292] abnormal_detected_host @0x7f4878e00100 = {0, 1} Bad!
Traceback (most recent call last):
  File "/tmp/ipykernel_22954/3616845043.py", line 4, in <module>
    test_func(p1)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 885, in __call__
    result = self._call(*args, **kwds)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  Bad! : Tensor had Inf values
     [[node CheckNumerics (defined at tmp/ipykernel_22954/2241890286.py:5) ]] [Op:__inference_simple_func_131]

Errors may have originated from an input operation.
Input Source operations connected to node CheckNumerics:
 input_arg (defined at tmp/ipykernel_22954/3616845043.py:4)

Function call stack:
simple_func

tf.debugging.check_numerics gây ra lỗi đối số không hợp lệ vì đối số Inf đối với test_func .

Bật trình tối ưu hóa vũ nữ thoát y gỡ lỗi và thực thi lại chức năng.

with options({'debug_stripper': True}):
  test_func2 = test_function_2()
  p1 = tf.constant(float('inf'))
  try:
    test_func2(p1)
  except tf.errors.InvalidArgumentError as e:
    traceback.print_exc(limit=2)

Trình tối ưu hóa trình gỡ lỗi thoát y tách nút tf.debug.check_numerics khỏi biểu đồ và thực thi chức năng mà không gây ra bất kỳ lỗi nào.

Bản tóm tắt

Thời gian chạy TensorFlow sử dụng Grappler để tự động tối ưu hóa đồ thị trước khi thực thi. Sử dụng tf.config.optimizer.set_experimental_options để bật hoặc tắt các trình tối ưu hóa đồ thị khác nhau.

Để biết thêm thông tin về Grappler, hãy xem Tối ưu hóa đồ thị TensorFlow .