Tổng quan
Trang này mô tả thiết kế và các bước cần thiết để chuyển đổi phép toán kết hợp trong TensorFlow sang các phép toán hợp nhất trong LiteRT. Cơ sở hạ tầng này cho mục đích chung và hỗ trợ chuyển đổi bất kỳ phép toán tổng hợp nào trong TensorFlow thành toán tử hợp nhất tương ứng trong LiteRT.
Một ví dụ về cách sử dụng cơ sở hạ tầng này là phương pháp kết hợp hoạt động TensorFlow RNN cho TensorFlow LiteRT, xem thông tin chi tiết tại đây.
Toán tử kết hợp là gì
Các hoạt động TensorFlow có thể là hoạt động gốc, ví dụ: tf.add hoặc chúng có thể bao gồm các thao tác gốc khác, ví dụ: tf.einsum. Nguyên gốc phép toán tử xuất hiện dưới dạng một nút duy nhất trong biểu đồ TensorFlow trong khi một hợp chất phép toán là một tập hợp các nút trong biểu đồ TensorFlow. Thực thi một toán tử kết hợp tương đương với việc thực thi từng phần tử gốc cấu thành của nó các toán tử.
Toán tử kết hợp tương ứng với một toán tử duy nhất chứa tất cả phép tính được thực hiện bởi từng phép toán gốc trong phép tính tương ứng phép toán tổng hợp.
Lợi ích của phép toán kết hợp
Có các toán tử kết hợp để tối đa hoá hiệu suất của nhân cơ bản bằng cách tối ưu hoá hoạt động tính toán tổng thể và giảm bộ nhớ dấu chân. Điều này rất có giá trị, đặc biệt đối với khối lượng công việc suy luận có độ trễ thấp và các nền tảng di động hạn chế tài nguyên.
Thao tác kết hợp cũng cung cấp giao diện cấp cao hơn để xác định các như lượng tử hoá, vốn không khả thi hoặc rất khó thực hiện khó thực hiện ở mức độ chi tiết hơn.
LiteRT có nhiều thực thể thao tác hợp nhất vì các lý do được trình bày ở trên. Các toán tử kết hợp này thường tương ứng với toán tử tổng hợp các hoạt động trong chương trình nguồn TensorFlow. Ví dụ về phép toán tổng hợp trong TensorFlow được triển khai dưới dạng một phép toán hợp nhất trong LiteRT bao gồm nhiều hoạt động RNN như trình tự một chiều và hai chiều LSTM, tích chập (Conv2d, bias add, relu), kết nối hoàn toàn (matmul, bias add, relu) và các lệnh khác. Trong LiteRT, việc định lượng LSTM hiện chỉ được triển khai trong các hoạt động LSTM kết hợp.
Những thách thức khi thao tác kết hợp
Chuyển đổi phép toán tổng hợp từ TensorFlow thành phép toán hợp nhất trong LiteRT là một bài toán khó. Điều này là do:
Toán tử tổng hợp được biểu thị trong biểu đồ TensorFlow dưới dạng tập hợp các phép toán nguyên gốc mà không có ranh giới được xác định rõ ràng. Có thể rất khó xác định (ví dụ: qua quá trình so khớp mẫu) biểu đồ phụ tương ứng với phép toán kết hợp đó.
Có thể có nhiều cách triển khai TensorFlow nhắm mục tiêu đến một công cụ kết hợp Hoạt động LiteRT. Ví dụ: có nhiều phương pháp triển khai LSTM trong TensorFlow (Keras, Breakfish/lingvo, v.v.) và mỗi TensorFlow bao gồm các thao tác gốc khác nhau nhưng tất cả vẫn có thể được chuyển đổi thành cùng một hoạt động LSTM kết hợp trong LiteRT.
Do đó, việc chuyển đổi hoạt động kết hợp đã được chứng minh là khá khó khăn.
Chuyển đổi từ hoạt động tổng hợp sang thao tác tuỳ chỉnh TFLite (nên dùng)
Gói toán tử tổng hợp trong tf.function
Trong nhiều trường hợp, một số phần của mô hình có thể được ánh xạ tới một toán tử trong
TFLite. Điều này có thể giúp cải thiện hiệu suất khi viết quy trình triển khai được tối ưu hoá
cho các hoạt động cụ thể. Để có thể tạo một thao tác kết hợp trong TFLite,
xác định một phần của biểu đồ đại diện cho phép kết hợp và gói phần đó vào
một tf.function với
"thử nghiệm_triển khai" cho tf.function, trong đó có thuộc tính này
giá trị tfl_fusable_op với giá trị true. Nếu thao tác tuỳ chỉnh mất
sau đó chuyển chúng dưới dạng một phần của cùng "experimental_implementations".
Ví dụ:
def get_implements_signature():
implements_signature = [
# 'name' will be used as a name for the operation.
'name: "my_custom_fused_op"',
# attr "tfl_fusable_op" is required to be set with true value.
'attr {key: "tfl_fusable_op" value { b: true } }',
# Example attribute "example_option" that the op accepts.
'attr {key: "example_option" value { i: %d } }' % 10
]
return ' '.join(implements_signature)
@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
# An empty function that represents pre/post processing example that
# is not represented as part of the Tensorflow graph.
output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
return output_1, output_2
class TestModel(tf.Module):
def __init__(self):
super(TestModel, self).__init__()
self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
@tf.function(input_signature=[
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
])
def simple_eval(self, input_a, input_b):
return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))
Lưu ý rằng bạn không cần đặt allow_custom_ops trên trình chuyển đổi dưới dạng
Thuộc tính tfl_fusable_op ngụ ý điều này đã.
Triển khai hoạt động tuỳ chỉnh và đăng ký bằng Trình thông dịch TFLite
Triển khai thao tác kết hợp dưới dạng thao tác TFLite Tuỳ chỉnh – xem hướng dẫn.
Lưu ý rằng tên để đăng ký op phải giống với tên
được chỉ định trong thuộc tính name trong chữ ký triển khai.
Ví dụ về hoạt động trong ví dụ này là
TfLiteRegistration reg = {};
// This name must match the name specified in the implements signature.
static constexpr char kOpName[] = "my_custom_fused_op";
reg.custom_name = kOpName;
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, ®);
Chuyển đổi từ thao tác kết hợp sang thao tác kết hợp (Nâng cao)
Kiến trúc tổng thể để chuyển đổi các phép toán tổng hợp TensorFlow thành Dưới đây là thao tác kết hợp LiteRT:
Gói toán tử tổng hợp trong tf.function
Trong mã nguồn mô hình TensorFlow, hãy xác định và trừu tượng hoá tổng hợp vào một tf.function bằng hàm experimental_implements chú thích hàm. Xem ví dụ về tra cứu nhúng. Chiến lược phát hành đĩa đơn hàm xác định giao diện và các đối số của nó sẽ được dùng để triển khai logic chuyển đổi.
Viết mã chuyển đổi
Mã chuyển đổi được viết trên giao diện của hàm bằng
Chú giải implements. Xem ví dụ về kiểu kết hợp khi nhúng
tìm kiếm. Về mặt lý thuyết, mã chuyển đổi thay thế chế độ kết hợp
phương thức triển khai giao diện này với giao diện kết hợp.
Trong lượt chuyển đổi các hàm chuẩn bị-composite, trình bổ trợ trong lượt chuyển đổi .
Trong các trường hợp sử dụng nâng cao hơn, có thể triển khai các phép biến đổi phức tạp của toán hạng của toán tử hợp nhất để lấy toán hạng của hợp nhất hoạt động. Xem Keras LSTM. mã chuyển đổi làm ví dụ.
Chuyển đổi sang LiteRT
Sử dụng TFLiteConverter.from_saved_model API để chuyển đổi thành LiteRT.
Tìm hiểu sâu
Hiện tại, chúng tôi mô tả chi tiết cấp cao về thiết kế tổng thể trong quá trình chuyển đổi sang kết hợp trong LiteRT.
Toán tử soạn thảo trong TensorFlow
Cách sử dụng tf.function với experimental_implements cho phép người dùng kết hợp rõ ràng các toán tử mới bằng cách sử dụng Thao tác gốc TensorFlow và chỉ định giao diện mà kết quả triển khai phép toán kết hợp. Điều này rất hữu ích vì nó cung cấp:
- Ranh giới được xác định rõ ràng cho phép toán kết hợp trong cơ sở Biểu đồ TensorFlow.
- Chỉ định rõ ràng giao diện mà thao tác này triển khai. Chiến lược phát hành đĩa đơn đối số của tf.function tương ứng với các đối số của giao diện này.
Ví dụ: hãy xem xét một phép toán kết hợp được xác định để triển khai tra cứu nhúng. Hàm này liên kết với một thao tác hợp nhất trong LiteRT.
@tf.function(
experimental_implements="embedding_lookup")
def EmbFprop(embs, ids_vec):
"""Embedding forward prop.
Effectively, it computes:
num = size of ids_vec
rets = zeros([num, embedding dim])
for i in range(num):
rets[i, :] = embs[ids_vec[i], :]
return rets
Args:
embs: The embedding matrix.
ids_vec: A vector of int32 embedding ids.
Returns:
The result of embedding lookups. A matrix of shape
[num ids in ids_vec, embedding dims].
"""
num = tf.shape(ids_vec)[0]
rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))
def EmbFpropLoop(i, embs, ids_vec, rets):
# row_id = ids_vec[i]
row_id = tf.gather(ids_vec, i)
# row = embs[row_id]
row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
# rets[i] = row
rets = inplace_ops.alias_inplace_update(rets, [i], row)
return embs, ids_vec, rets
_, _, rets = functional_ops.For(
start=0,
limit=num,
delta=1,
inputs=[embs, ids_vec, rets],
body=EmbFpropLoop,
rewrite_with_while=compiled)
if len(weight_shape) > 2:
rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
return rets
Bằng cách làm cho các mô hình sử dụng phép toán tổng hợp qua tf.function dưới dạng minh hoạ ở trên, có thể xây dựng một cơ sở hạ tầng chung để xác định và chuyển đổi các thao tác đó thành thao tác LiteRT kết hợp.
Mở rộng bộ chuyển đổi LiteRT
Công cụ chuyển đổi LiteRT được phát hành đầu năm nay chỉ được hỗ trợ nhập các mô hình TensorFlow dưới dạng biểu đồ, trong đó tất cả các biến được thay thế bằng các giá trị hằng số tương ứng. Điều này không phù hợp với hợp nhất hoạt động vì các đồ thị như vậy có tất cả các hàm cùng dòng để có thể biến các biến thành hằng số.
Để tận dụng
tf.function bằng hàm
experimental_implements trong quá trình chuyển đổi, các hàm
cần được lưu giữ cho đến sau này trong quá trình chuyển đổi.
Do đó, chúng tôi đã triển khai một quy trình mới để nhập và chuyển đổi TensorFlow trong bộ chuyển đổi để hỗ trợ trường hợp sử dụng hợp nhất toán tử tổng hợp. Cụ thể, chúng tôi đã bổ sung các tính năng mới:
- Nhập các mô hình đã lưu vào MLIR của TensorFlow
- hợp nhất phép toán
- phân tích khả năng biến đổi của biến
- cố định tất cả biến chỉ có thể đọc
Điều này cho phép chúng ta thực hiện hợp nhất thao tác bằng cách sử dụng các hàm biểu thị các phép toán kết hợp trước khi hàm cùng dòng và cố định biến.
Triển khai hoạt động kết hợp
Hãy cùng tìm hiểu cụ thể hơn về kỹ thuật truyền dữ liệu kết hợp hoạt động. Thẻ/vé này sau:
- Lặp lại qua tất cả các hàm trong mô-đun MLIR.
- Nếu một hàm có thuộc tính tf._implementations, dựa trên thuộc tính đó gọi là phần mềm tiện ích tổng hợp toán tử thích hợp.
- Phần mềm tiện ích hợp nhất toán tử hoạt động dựa trên toán hạng của hàm và (đóng vai trò là giao diện cho lượt chuyển đổi) và thay thế phần nội dung của hàm có phần nội dung hàm tương đương chứa thao tác kết hợp.
- Trong nhiều trường hợp, phần nội dung được thay thế sẽ chứa các toán tử khác với thao tác kết hợp. Các giá trị này tương ứng với một số biến đổi tĩnh trên toán hạng của hàm để lấy toán hạng của toán tử hợp nhất. Do tất cả các phép tính này có thể liên tục gấp lại, nên chúng sẽ không có trong vùng đệm phẳng đã xuất, trong đó chỉ có thao tác kết hợp đã tồn tại.
Dưới đây là đoạn mã lấy từ thẻ cho thấy quy trình làm việc chính:
void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
StringAttr attr) {
if (attr.getValue() == "embedding_lookup") {
func.eraseBody();
func.addEntryBlock();
// Convert the composite embedding_lookup function body to a
// TFLite fused embedding_lookup op.
ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
if (failed(convert_embedded_lookup.VerifySignature())) {
return signalPassFailure();
}
convert_embedded_lookup.RewriteFunc();
} else if (attr.getValue() == mlir::TFL::kKerasLstm) {
func.eraseBody();
func.addEntryBlock();
OpBuilder builder(func.getBody());
if (failed(ConvertKerasLSTMLayer(func, &builder))) {
return signalPassFailure();
}
} else if (.....) /* Other fusions can plug in here */
}
Dưới đây là đoạn mã cho thấy việc ánh xạ toán tử kết hợp này đến hàm hợp nhất trong LiteRT tận dụng hàm làm giao diện chuyển đổi.
void RewriteFunc() {
Value lookup = func_.getArgument(1);
Value value = func_.getArgument(0);
auto output_type = func_.getType().getResult(0);
OpBuilder builder(func_.getBody());
auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
func_.getLoc(), output_type, lookup, value);
builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
}