ตัวดำเนินการที่กำหนดเอง

เนื่องจากไลบรารีตัวดำเนินการในตัว TensorFlow Lite รองรับตัวดำเนินการ TensorFlow ในจำนวนจำกัดเท่านั้น จึงไม่ใช่ทุกรุ่นที่จะแปลงได้ สำหรับรายละเอียด โปรดดูที่ ความเข้ากันได้ของผู้ให้บริการ

หากต้องการอนุญาตการแปลง ผู้ใช้สามารถจัดเตรียมการใช้งานตัวดำเนินการ TensorFlow ที่ไม่รองรับใน TensorFlow Lite หรือที่เรียกว่าตัวดำเนินการแบบกำหนดเองแบบกำหนดเองได้ หากคุณต้องการรวมชุดตัวดำเนินการ TensorFlow ที่ไม่รองรับ (หรือรองรับ) เข้ากับตัวดำเนินการแบบกำหนดเองที่ได้รับการปรับให้เหมาะสมที่สุดเพียงตัวเดียว โปรดดูที่ การหลอมรวมตัวดำเนินการ

การใช้ตัวดำเนินการแบบกำหนดเองประกอบด้วยสี่ขั้นตอน

  • สร้างโมเดล TensorFlow ตรวจสอบให้แน่ใจว่าโมเดลที่บันทึกไว้ (หรือ Graph Def) อ้างอิงถึงตัวดำเนินการ TensorFlow Lite ที่มีชื่ออย่างถูกต้อง

  • แปลงเป็นโมเดล TensorFlow Lite ตรวจสอบให้แน่ใจว่าคุณตั้งค่าแอตทริบิวต์ตัวแปลง TensorFlow Lite ที่ถูกต้องเพื่อให้สามารถแปลงโมเดลได้สำเร็จ

  • สร้างและลงทะเบียนตัวดำเนินการ ทั้งนี้เพื่อให้รันไทม์ TensorFlow Lite รู้วิธีแมปตัวดำเนินการและพารามิเตอร์ในกราฟของคุณกับโค้ด C/C++ ที่รันได้

  • ทดสอบและโปรไฟล์ผู้ให้บริการของคุณ หากคุณต้องการทดสอบเฉพาะโอเปอเรเตอร์ที่คุณกำหนดเอง วิธีที่ดีที่สุดคือสร้างโมเดลด้วยโอเปอเรเตอร์ที่คุณกำหนดเองและใช้โปรแกรม benchmark_model

มาดูตัวอย่างแบบครบวงจรของการรันโมเดลด้วยโอเปอเรเตอร์ที่กำหนดเอง tf.atan (ชื่อ Atan อ้างอิงถึง #create_a_tensorflow_model) ซึ่งได้รับการรองรับใน TensorFlow แต่ไม่รองรับใน TensorFlow Lite

ตัวดำเนินการข้อความ TensorFlow คือตัวอย่างของตัวดำเนินการแบบกำหนดเอง ดูบทช่วยสอน การแปลงข้อความ TF เป็น TF Lite สำหรับตัวอย่างโค้ด

ตัวอย่าง: ตัวดำเนินการ Atan แบบกำหนดเอง

มาดูตัวอย่างการสนับสนุนตัวดำเนินการ TensorFlow ที่ TensorFlow Lite ไม่มี สมมติว่าเรากำลังใช้ตัวดำเนินการ Atan และเรากำลังสร้างโมเดลที่เรียบง่ายมากสำหรับฟังก์ชัน y = atan(x + offset) โดยที่ offset สามารถฝึกได้

สร้างโมเดล TensorFlow

ข้อมูลโค้ดต่อไปนี้ฝึกโมเดล TensorFlow อย่างง่าย โมเดลนี้มีเพียงโอเปอเรเตอร์ที่กำหนดเองชื่อ Atan ซึ่งเป็นฟังก์ชัน y = atan(x + offset) โดยที่ offset สามารถฝึกได้

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-1.4288993, 0.98279375, 1.2490457, 1.2679114, 1.5658458]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Atan`
@tf.function(input_signature=[tf.TensorSpec.from_tensor(tf.constant(x))])
def atan(x):
  return tf.atan(x + offset, name="Atan")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = atan(x)
      loss = tf.reduce_sum(tf.square(predicted_y - y))
    grads = t.gradient(loss, [offset])
    optimizer.apply_gradients(zip(grads, [offset]))

for i in range(1000):
    train(x, y)

print("The actual offset is: 1.0")
print("The predicted offset is:", offset.numpy())
The actual offset is: 1.0
The predicted offset is: 0.99999905

ณ จุดนี้ หากคุณพยายามสร้างโมเดล TensorFlow Lite ด้วยแฟล็กตัวแปลงเริ่มต้น คุณจะได้รับข้อความแสดงข้อผิดพลาดต่อไปนี้:

Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.

แปลงเป็นโมเดล TensorFlow Lite

สร้างโมเดล TensorFlow Lite ด้วยตัวดำเนินการที่กำหนดเอง โดยการตั้งค่าแอตทริบิวต์ตัวแปลง allow_custom_ops ดังที่แสดงด้านล่าง:

converter = tf.lite.TFLiteConverter.from_concrete_functions([atan.get_concrete_function()], atan)
converter.allow_custom_ops = True
tflite_model = converter.convert()

ณ จุดนี้ หากคุณรันด้วยล่ามเริ่มต้นโดยใช้คำสั่งดังต่อไปนี้:

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

คุณจะยังคงได้รับข้อผิดพลาด:

Encountered unresolved custom op: Atan.

สร้างและลงทะเบียนตัวดำเนินการ

ตัวดำเนินการ TensorFlow Lite ทั้งหมด (ทั้งแบบกำหนดเองและในตัว) ได้รับการกำหนดโดยใช้อินเทอร์เฟซ pure-C แบบธรรมดาที่ประกอบด้วยสี่ฟังก์ชัน:

typedef struct {
  void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
  void (*free)(TfLiteContext* context, void* buffer);
  TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
} TfLiteRegistration;

อ้างถึง common.h สำหรับรายละเอียดเกี่ยวกับ TfLiteContext และ TfLiteNode แบบแรกอำนวยความสะดวกในการรายงานข้อผิดพลาดและการเข้าถึงอ็อบเจ็กต์ระดับโลก รวมถึงเทนเซอร์ทั้งหมด ส่วนหลังอนุญาตให้นำไปใช้งานในการเข้าถึงอินพุตและเอาต์พุต

เมื่อล่ามโหลดโมเดล มันจะเรียก init() หนึ่งครั้งสำหรับแต่ละโหนดในกราฟ init() ที่กำหนดจะถูกเรียกมากกว่าหนึ่งครั้งหากใช้ op หลายครั้งในกราฟ สำหรับการดำเนินการแบบกำหนดเอง จะมีการจัดเตรียมบัฟเฟอร์การกำหนดค่า ซึ่งมี flexbuffer ที่แมปชื่อพารามิเตอร์กับค่าของมัน บัฟเฟอร์ว่างเปล่าสำหรับ ops ในตัวเนื่องจากล่ามได้แยกวิเคราะห์พารามิเตอร์ op แล้ว การใช้งานเคอร์เนลที่จำเป็นต้องมีสถานะควรเริ่มต้นที่นี่และโอนความเป็นเจ้าของให้กับผู้โทร สำหรับการเรียก init() แต่ละครั้ง จะมีการเรียกที่สอดคล้องกันไปยัง free() ซึ่งช่วยให้การใช้งานสามารถกำจัดบัฟเฟอร์ที่พวกเขาอาจจัดสรรไว้ใน init()

เมื่อใดก็ตามที่เทนเซอร์อินพุตถูกปรับขนาด ล่ามจะดูกราฟเพื่อแจ้งการดำเนินการเปลี่ยนแปลง สิ่งนี้ทำให้พวกเขามีโอกาสที่จะปรับขนาดบัฟเฟอร์ภายใน ตรวจสอบความถูกต้องของรูปร่างและประเภทอินพุต และคำนวณรูปร่างเอาต์พุตใหม่ ทั้งหมดนี้ทำได้ผ่าน prepare() และการใช้งานสามารถเข้าถึงสถานะของตนได้โดยใช้ node->user_data

ในที่สุด แต่ละครั้งที่การอนุมานทำงาน ล่ามจะสำรวจกราฟที่เรียก invoke() และที่นี่สถานะก็พร้อมใช้งานเป็น node->user_data กัน

Custom ops สามารถนำไปใช้ในลักษณะเดียวกับ ops ในตัว โดยการกำหนดฟังก์ชันทั้งสี่ดังกล่าวและฟังก์ชันการลงทะเบียนส่วนกลางที่มักจะมีลักษณะดังนี้:

namespace my_namespace {
  const TfLiteRegistration* Register_MY_CUSTOM_OP() {
    static const TfLiteRegistration r = {my_custom_op::Init,
                                         my_custom_op::Free,
                                         my_custom_op::Prepare,
                                         my_custom_op::Eval};
    return &r;
  }
}  // namespace my_namespace

โปรดทราบว่าการลงทะเบียนจะไม่เกิดขึ้นโดยอัตโนมัติ และควรทำการเรียก Register_MY_CUSTOM_OP อย่างชัดเจน แม้ว่า BuiltinOpResolver มาตรฐาน (พร้อมใช้งานจาก :builtin_ops เป้าหมาย) จะดูแลการลงทะเบียนบิวด์อิน แต่ ops ที่กำหนดเองจะต้องถูกรวบรวมในไลบรารีที่กำหนดเองแยกต่างหาก

การกำหนดเคอร์เนลในรันไทม์ TensorFlow Lite

สิ่งที่เราต้องทำเพื่อใช้ op ใน TensorFlow Lite คือการกำหนดสองฟังก์ชัน ( Prepare และ Eval ) และสร้าง TfLiteRegistration :

TfLiteStatus AtanPrepare(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  int num_dims = NumDimensions(input);

  TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
  for (int i=0; i<num_dims; ++i) {
    output_size->data[i] = input->dims->data[i];
  }

  return context->ResizeTensor(context, output, output_size);
}

TfLiteStatus AtanEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  float* input_data = GetTensorData<float>(input);
  float* output_data = GetTensorData<float>(output);

  size_t count = 1;
  int num_dims = NumDimensions(input);
  for (int i = 0; i < num_dims; ++i) {
    count *= input->dims->data[i];
  }

  for (size_t i=0; i<count; ++i) {
    output_data[i] = atan(input_data[i]);
  }
  return kTfLiteOk;
}

const TfLiteRegistration* Register_ATAN() {
  static const TfLiteRegistration r = {nullptr, nullptr, AtanPrepare, AtanEval};
  return &r;
}

เมื่อเริ่มต้น OpResolver ให้เพิ่ม op ที่กำหนดเองลงในตัวแก้ไข (ดูตัวอย่างด้านล่าง) การดำเนินการนี้จะลงทะเบียนโอเปอเรเตอร์กับ Tensorflow Lite เพื่อให้ TensorFlow Lite สามารถใช้การใช้งานใหม่ได้ โปรดทราบว่าอาร์กิวเมนต์สองรายการสุดท้ายใน TfLiteRegistration สอดคล้องกับฟังก์ชัน AtanPrepare และ AtanEval ที่คุณกำหนดไว้สำหรับ op ที่กำหนดเอง หากคุณใช้ฟังก์ชัน AtanInit และ AtanFree เพื่อเตรียมใช้งานตัวแปรที่ใช้ใน op และเพื่อเพิ่มพื้นที่ว่าง ตามลำดับ ฟังก์ชันเหล่านั้นจะถูกเพิ่มเข้าไปในอาร์กิวเมนต์สองตัวแรกของ TfLiteRegistration อาร์กิวเมนต์เหล่านั้นถูกตั้งค่าเป็น nullptr ในตัวอย่างนี้

ลงทะเบียนโอเปอเรเตอร์กับไลบรารีเคอร์เนล

ตอนนี้เราจำเป็นต้องลงทะเบียนโอเปอเรเตอร์กับไลบรารีเคอร์เนล นี้จะกระทำด้วย OpResolver เบื้องหลัง ล่ามจะโหลดไลบรารีของเคอร์เนลซึ่งจะถูกกำหนดให้ดำเนินการกับโอเปอเรเตอร์แต่ละตัวในโมเดล แม้ว่าไลบรารีเริ่มต้นจะมีเคอร์เนลในตัวเท่านั้น แต่ก็สามารถแทนที่/เสริมด้วยตัวดำเนินการไลบรารีแบบกำหนดเองได้

คลาส OpResolver ซึ่งแปลรหัสโอเปอเรเตอร์และชื่อเป็นโค้ดจริง ถูกกำหนดไว้ดังนี้:

class OpResolver {
 public:
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  ...
};

คลาส MutableOpResolver และ BuiltinOpResolver มาจาก OpResolver :

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddBuiltin(tflite::BuiltinOperator op, const TfLiteRegistration* registration) = 0;
  void AddCustom(const char* op, const TfLiteRegistration* registration) = 0;
  void AddAll(const MutableOpResolver& other);
  ...
};

class BuiltinOpResolver : public MutableOpResolver {
 public:
  BuiltinOpResolver();  // Constructs an op resolver with all the builtin ops.
};

การใช้งานปกติกำหนดให้คุณใช้ BuiltinOpResolver และเขียน:

tflite::ops::builtin::BuiltinOpResolver resolver;

หากต้องการเพิ่ม op แบบกำหนดเองที่สร้างขึ้นด้านบน คุณสามารถใช้ MutableOpResolver แทน และเรียก AddCustom (ก่อนที่คุณจะส่งตัวแก้ไขไปที่ InterpreterBuilder ):

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
resolver.AddCustom("Atan", Register_ATAN());

หากถือว่าชุดของ ops ในตัวมีขนาดใหญ่เกินไป OpResolver ใหม่สามารถสร้างโค้ดโดยอิงจากชุดย่อยของ ops ที่กำหนด อาจเป็นเฉพาะชุดที่มีอยู่ในโมเดลที่กำหนดเท่านั้น ซึ่งเทียบเท่ากับการลงทะเบียนแบบเลือกของ TensorFlow (และมีเวอร์ชันเรียบง่ายอยู่ในไดเรกทอรี tools )

หากคุณต้องการกำหนดโอเปอเรเตอร์ที่กำหนดเองใน Java คุณจะต้องสร้างเลเยอร์ JNI ที่กำหนดเองของคุณเองและคอมไพล์ AAR ของคุณเอง ในโค้ด jni นี้ ในทำนองเดียวกัน หากคุณต้องการกำหนดโอเปอเรเตอร์เหล่านี้ที่มีอยู่ใน Python คุณสามารถลงทะเบียนของคุณใน โค้ด Wrapper ของ Python

โปรดทราบว่าสามารถปฏิบัติตามกระบวนการที่คล้ายกันข้างต้นเพื่อสนับสนุนชุดการดำเนินการแทนที่จะเป็นตัวดำเนินการเพียงตัวเดียว เพียงเพิ่มตัวดำเนินการ AddCustom ได้มากเท่าที่คุณต้องการ นอกจากนี้ MutableOpResolver ยังช่วยให้คุณสามารถแทนที่การใช้งานบิวด์อินได้โดยใช้ AddBuiltin

ทดสอบและโปรไฟล์ผู้ให้บริการของคุณ

หากต้องการโปรไฟล์การดำเนินการของคุณด้วยเครื่องมือวัดประสิทธิภาพ TensorFlow Lite คุณสามารถใช้ เครื่องมือโมเดลวัดประสิทธิภาพ สำหรับ TensorFlow Lite เพื่อวัตถุประสงค์ในการทดสอบ คุณสามารถทำให้ TensorFlow Lite บิลด์ในเครื่องของคุณทราบถึง op ที่คุณกำหนดเองโดยเพิ่มการเรียก AddCustom ที่เหมาะสม (ดังที่แสดงด้านบน) เพื่อ register.cc

ปฏิบัติที่ดีที่สุด

  1. ปรับการจัดสรรหน่วยความจำและยกเลิกการจัดสรรหน่วยความจำให้เหมาะสมด้วยความระมัดระวัง การจัดสรรหน่วยความจำใน Prepare จะมีประสิทธิภาพมากกว่าใน Invoke และการจัดสรรหน่วยความจำก่อนการวนซ้ำจะดีกว่าในการวนซ้ำทุกครั้ง ใช้ข้อมูลเทนเซอร์ชั่วคราวแทนที่จะทำการแบ่งตัวเอง (ดูรายการที่ 2) ใช้พอยน์เตอร์/ข้อมูลอ้างอิงแทนการคัดลอกให้มากที่สุด

  2. หากโครงสร้างข้อมูลยังคงอยู่ระหว่างการดำเนินการทั้งหมด เราแนะนำให้จัดสรรหน่วยความจำล่วงหน้าโดยใช้เทนเซอร์ชั่วคราว คุณอาจจำเป็นต้องใช้โครงสร้าง OpData เพื่ออ้างอิงดัชนีเทนเซอร์ในฟังก์ชันอื่นๆ ดูตัวอย่างใน เคอร์เนลสำหรับการบิด ข้อมูลโค้ดตัวอย่างอยู่ด้านล่าง

    auto* op_data = reinterpret_cast<OpData*>(node->user_data);
    TfLiteIntArrayFree(node->temporaries);
    node->temporaries = TfLiteIntArrayCreate(1);
    node->temporaries->data[0] = op_data->temp_tensor_index;
    TfLiteTensor* temp_tensor = &context->tensors[op_data->temp_tensor_index];
    temp_tensor->type =  kTfLiteFloat32;
    temp_tensor->allocation_type = kTfLiteArenaRw;
    
  3. หากไม่เปลืองหน่วยความจำมากเกินไป ให้เลือกใช้อาร์เรย์ขนาดคงที่แบบคงที่ (หรือ std::vector ที่จัดสรรไว้ล่วงหน้าใน Resize ) แทนที่จะใช้ std::vector ที่จัดสรรแบบไดนามิกทุกๆ การวนซ้ำของการดำเนินการ

  4. หลีกเลี่ยงการสร้างอินสแตนซ์เทมเพลตคอนเทนเนอร์ไลบรารีมาตรฐานที่ไม่มีอยู่แล้ว เนื่องจากจะส่งผลต่อขนาดไบนารี ตัวอย่างเช่น หากคุณต้องการ std::map ในการดำเนินการของคุณที่ไม่มีอยู่ในเคอร์เนลอื่น การใช้ std::vector ที่มีการแมปการจัดทำดัชนีโดยตรงอาจทำงานได้ในขณะที่รักษาขนาดไบนารี่ให้เล็ก ดูว่าเคอร์เนลอื่นใช้อะไรเพื่อรับข้อมูลเชิงลึก (หรือถาม)

  5. ตรวจสอบตัวชี้ไปยังหน่วยความจำที่ส่งคืนโดย malloc หากตัวชี้นี้เป็น nullptr ไม่ควรดำเนินการใดๆ โดยใช้ตัวชี้นั้น หากคุณ malloc ในฟังก์ชันและมีข้อผิดพลาดในการออก ให้จัดสรรหน่วยความจำใหม่ก่อนที่คุณจะออก

  6. ใช้ TF_LITE_ENSURE(context, condition) เพื่อตรวจสอบเงื่อนไขเฉพาะ รหัสของคุณจะต้องไม่ปล่อยให้หน่วยความจำค้างเมื่อมีการใช้ TF_LITE_ENSURE กล่าวคือ ควรใช้แมโครเหล่านี้ก่อนที่จะจัดสรรทรัพยากรใดๆ ที่จะรั่วไหล