צור אופ

אם תרצה ליצור אופציה שאינה מכוסה על ידי ספריית TensorFlow הקיימת, אנו ממליצים לך תחילה לנסות לכתוב את האופציה ב-Python כהרכב של פעולות או פונקציות קיימות של Python. אם זה לא אפשרי, אתה יכול ליצור אופציה מותאמת אישית של C++. ישנן מספר סיבות מדוע ייתכן שתרצה ליצור אופציה מותאמת אישית של C++:

  • זה לא קל או אפשרי לבטא את הפעולה שלך כהרכב של מבצעים קיימים.
  • זה לא יעיל לבטא את הפעולה שלך כהרכב של פרימיטיבים קיימים.
  • אתה רוצה למזג ידנית קומפוזיציה של פרימיטיבים שלמהדר עתידי יהיה קשה להתמזג.

לדוגמה, תאר לעצמך שאתה רוצה ליישם משהו כמו "איגוד חציוני", בדומה לאופרטור "MaxPool", אבל מחשוב חציונים על חלונות הזזה במקום ערכים מקסימליים. ביצוע זה באמצעות קומפוזיציה של פעולות עשוי להיות אפשרי (למשל, שימוש ב-ExtractImagePatches ו-TopK), אך עשוי להיות לא יעיל בביצועים או בזיכרון כמו פעולה מקורית שבה אתה יכול לעשות משהו חכם יותר בפעולה אחת ומחוברת. כמו תמיד, בדרך כלל כדאי קודם כל לנסות להביע את מה שאתה רוצה באמצעות הרכב מפעיל, רק לבחור להוסיף פעולה חדשה אם זה מתגלה כקשה או לא יעיל.

כדי לשלב את האופציה המותאמת אישית שלך, תצטרך:

  1. רשום את האופציה החדשה בקובץ C++. רישום Op מגדיר ממשק (מפרט) לפונקציונליות של ה-Op, שאינו תלוי ביישום ה-Op. לדוגמה, רישום op מגדיר את שם ה-op ואת הכניסות והפלטים של ה-op. זה גם מגדיר את פונקציית הצורה המשמשת להסקת צורת טנזור.
  2. הטמיע את ה-op ב-C++. היישום של op ידוע בתור קרנל, והוא היישום הקונקרטי של המפרט שרשמתם בשלב 1. יכולים להיות מספר גרעינים עבור סוגי קלט/פלט או ארכיטקטורות שונות (לדוגמה, CPUs, GPUs).
  3. צור מעטפת Python (אופציונלי). מעטפת זו היא ה-API הציבורי המשמש ליצירת ה-op ב-Python. מעטפת ברירת מחדל נוצרת מהרישום המבצע, שניתן להשתמש בה ישירות או להוסיף אליה.
  4. כתוב פונקציה לחישוב מעברי צבע עבור ה-op (אופציונלי).
  5. בדוק את האופציה. בדרך כלל אנו עושים זאת ב-Python מטעמי נוחות, אבל אתה יכול גם לבדוק את ה-Op ב-C++. אם אתה מגדיר מעברי צבע, אתה יכול לאמת אותם עם Python tf.test.compute_gradient_error . ראה relu_op_test.py כדוגמה הבודקת את הפונקציות קדימה של אופרטורים דמויי Relu והשיפועים שלהם.

דרישות מוקדמות

הגדר את ממשק ההפעלה

אתה מגדיר את הממשק של הפעלה על ידי רישום שלו במערכת TensorFlow. ברישום, אתה מציין את שם המבצע שלך, הקלט שלו (סוגים ושמות) ופלטים (טיפוסים ושמות), כמו גם מחרוזות docstrings וכל כתובות שהאופ עשוי לדרוש.

כדי לראות איך זה עובד, נניח שברצונך ליצור אופ שלוקח טנסור של int32 s ומוציא עותק של הטנסור, כשהכל מלבד האלמנט הראשון מוגדר לאפס. לשם כך, צור קובץ בשם zero_out.cc . לאחר מכן הוסף קריאה למאקרו REGISTER_OP המגדיר את הממשק עבור הפעולה שלך:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

אופציה זו ZeroOut לוקחת טנסור אחד to_zero של מספרים שלמים של 32 סיביות כקלט, ומוציאה טנזור zeroed של מספרים שלמים של 32 סיביות. ה-op משתמש גם בפונקציית צורה כדי להבטיח שטנסור הפלט זהה לצורה של טנסור הקלט. לדוגמה, אם הקלט הוא טנסור של צורה [10, 20], אז פונקציית צורה זו מציינת שצורת הפלט היא גם [10, 20].

יישם את הליבה עבור האופ

לאחר הגדרת הממשק, ספק מימוש אחד או יותר של ה-op. כדי ליצור אחד מהקרנלים הללו, צור מחלקה שמרחיבה את OpKernel ועוברת את שיטת Compute . שיטת Compute מספקת ארגומנט context אחד מסוג OpKernelContext* , שממנו אתה יכול לגשת לדברים שימושיים כמו טנסורי הקלט והפלט.

הוסף את הקרנל שלך לקובץ שיצרת למעלה. הקרנל עשוי להיראות בערך כך:

#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output_flat(0) = input(0);
  }
};

לאחר הטמעת הליבה שלך, אתה רושם אותו במערכת TensorFlow. ברישום, אתה מציין אילוצים שונים שתחתם הליבה הזו תפעל. לדוגמה, ייתכן שיהיה לך ליבה אחת המיועדת למעבדים, ואחד נפרד למעבדי GPU.

כדי לעשות זאת עבור ZeroOut op, הוסף את הדברים הבאים ל- zero_out.cc :

REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

גרעיני מעבד מרובי הליכי

כדי לכתוב ליבת CPU מרובה הליכי, ניתן להשתמש בפונקציית Shard ב- work_sharder.h . פונקציה זו חותכת פונקציית חישוב על פני השרשורים המוגדרים לשימוש עבור השרשור תוך-אופ (ראה intra_op_parallelism_threads ב- config.proto ).

גרעיני GPU

ליבת GPU מיושמת בשני חלקים: OpKernel ו-CUDA וקוד ההשקה שלו.

לפעמים יישום OpKernel נפוץ בין גרעין מעבד ו-GPU, כגון סביב בדיקת תשומות והקצאת פלטים. במקרה זה, יישום מוצע הוא:

  1. הגדר את התבנית OpKernel במכשיר ואת הסוג הפרימיטיבי של הטנזור.
  2. כדי לבצע את החישוב בפועל של הפלט, הפונקציה Compute קוראת למבנה functor בתבנית.
  3. ההתמחות של אותו פונקטור עבור ה-CPUDevice מוגדרת באותו קובץ, אך ההתמחות עבור ה-GPUDevice מוגדרת בקובץ .cu.cc, מכיוון שהוא יודר עם מהדר CUDA.

הנה דוגמה ליישום.

// kernel_example.h
#ifndef KERNEL_EXAMPLE_H_
#define KERNEL_EXAMPLE_H_

#include <unsupported/Eigen/CXX11/Tensor>

template <typename Device, typename T>
struct ExampleFunctor {
  void operator()(const Device& d, int size, const T* in, T* out);
};

#if GOOGLE_CUDA
// Partially specialize functor for GpuDevice.
template <typename T>
struct ExampleFunctor<Eigen::GpuDevice, T> {
  void operator()(const Eigen::GpuDevice& d, int size, const T* in, T* out);
};
#endif

#endif KERNEL_EXAMPLE_H_
// kernel_example.cc
#include "kernel_example.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

REGISTER_OP("Example")
    .Attr("T: numbertype")
    .Input("input: T")
    .Output("input_times_two: T")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

// CPU specialization of actual computation.
template <typename T>
struct ExampleFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, int size, const T* in, T* out) {
    for (int i = 0; i < size; ++i) {
      out[i] = 2 * in[i];
    }
  }
};

// OpKernel definition.
// template parameter <T> is the datatype of the tensors.
template <typename Device, typename T>
class ExampleOp : public OpKernel {
 public:
  explicit ExampleOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));

    // Do the computation.
    OP_REQUIRES(context, input_tensor.NumElements() <= tensorflow::kint32max,
                errors::InvalidArgument("Too many elements in tensor"));
    ExampleFunctor<Device, T>()(
        context->eigen_device<Device>(),
        static_cast<int>(input_tensor.NumElements()),
        input_tensor.flat<T>().data(),
        output_tensor->flat<T>().data());
  }
};

// Register the CPU kernels.
#define REGISTER_CPU(T)                                          \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
      ExampleOp<CPUDevice, T>);
REGISTER_CPU(float);
REGISTER_CPU(int32);

// Register the GPU kernels.
#ifdef GOOGLE_CUDA
#define REGISTER_GPU(T)                                          \
  /* Declare explicit instantiations in kernel_example.cu.cc. */ \
  extern template class ExampleFunctor<GPUDevice, T>;            \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
      ExampleOp<GPUDevice, T>);
REGISTER_GPU(float);
REGISTER_GPU(int32);
#endif  // GOOGLE_CUDA
// kernel_example.cu.cc
#ifdef GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "kernel_example.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"

using namespace tensorflow;

using GPUDevice = Eigen::GpuDevice;

// Define the CUDA kernel.
template <typename T>
__global__ void ExampleCudaKernel(const int size, const T* in, T* out) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;
       i += blockDim.x * gridDim.x) {
    out[i] = 2 * __ldg(in + i);
  }
}

// Define the GPU implementation that launches the CUDA kernel.
template <typename T>
void ExampleFunctor<GPUDevice, T>::operator()(
    const GPUDevice& d, int size, const T* in, T* out) {
  // Launch the cuda kernel.
  //
  // See core/util/gpu_kernel_helper.h for example of computing
  // block count and thread_per_block count.
  int block_count = 1024;
  int thread_per_block = 20;
  ExampleCudaKernel<T>
      <<<block_count, thread_per_block, 0, d.stream()>>>(size, in, out);
}

// Explicitly instantiate functors for the types of OpKernels registered.
template struct ExampleFunctor<GPUDevice, float>;
template struct ExampleFunctor<GPUDevice, int32>;

#endif  // GOOGLE_CUDA

בנה את ספריית ה-op

קומפל את האופ באמצעות מהדר המערכת שלך (התקנה בינארית של TensorFlow)

אתה אמור להיות מסוגל להדר zero_out.cc עם מהדר C++ כגון g++ או clang הזמין במערכת שלך. חבילת ה-PIP הבינארית מתקינה את קבצי הכותרות ואת הספרייה שאתה צריך כדי להרכיב את ה-Op שלך במיקומים ספציפיים למערכת. עם זאת, ספריית TensorFlow python מספקת את הפונקציה get_include כדי לקבל את ספריית הכותרת, ולספריית get_lib יש אובייקט משותף לקשר אליו. להלן הפלטים של פונקציות אלה במכונת אובונטו.

$ python
>>> import tensorflow as tf
>>> tf.sysconfig.get_include()
'/usr/local/lib/python3.6/site-packages/tensorflow/include'
>>> tf.sysconfig.get_lib()
'/usr/local/lib/python3.6/site-packages/tensorflow'

בהנחה שהתקנת את g++ , הנה רצף הפקודות שבהן אתה יכול להשתמש כדי לקמפל את האופציה שלך לספרייה דינמית.

TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++14 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

ב-macOS, הדגל הנוסף "-undefined dynamic_lookup" נדרש בעת בניית קובץ ה- .so .

הערה לגבי גרסת gcc >=5 : gcc משתמש ב-C++ ABI החדש מאז גרסה 5 . TensorFlow 2.8 ואילך נבנו עם gcc4 שמשתמש ב-ABI הישן יותר. אם אתה משתמש בגרסאות אלה של TensorFlow ומנסה להרכיב את ספריית ה-op שלך עם gcc>=5 , הוסף -D_GLIBCXX_USE_CXX11_ABI=0 לשורת הפקודה כדי להפוך את הספרייה לתואמת ל-ABI הישן יותר. חבילות TensorFlow 2.9+ תואמות ל-ABI החדש יותר כברירת מחדל.

קומפל את האופ באמצעות bazel (התקנת מקור TensorFlow)

אם מותקנים לך מקורות TensorFlow, אתה יכול להשתמש במערכת הבנייה של TensorFlow כדי להרכיב את האופציה שלך. הצב קובץ BUILD עם כלל הבנייה הבא של Bazel בספריית tensorflow/core/user_ops .

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"],
)

הפעל את הפקודה הבאה כדי לבנות zero_out.so .

$ bazel build --config opt //tensorflow/core/user_ops:zero_out.so

כדי להרכיב את פעולת Example , עם ליבת CUDA, עליך להשתמש בפרמטר gpu_srcs של tf_custom_op_library . הצב קובץ BUILD עם כלל הבנייה הבא של Bazel בתיקייה חדשה בתוך ספריית tensorflow/core/user_ops (למשל "example_gpu").

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    # kernel_example.cc  kernel_example.cu.cc  kernel_example.h
    name = "kernel_example.so",
    srcs = ["kernel_example.h", "kernel_example.cc"],
    gpu_srcs = ["kernel_example.cu.cc", "kernel_example.h"],
)

הפעל את הפקודה הבאה כדי לבנות kernel_example.so .

$ bazel build --config opt //tensorflow/core/user_ops/example_gpu:kernel_example.so

השתמש ב-op ב-Python

TensorFlow Python API מספק את הפונקציה tf.load_op_library לטעינת הספרייה הדינמית ורישום ה-op עם המסגרת של TensorFlow. load_op_library מחזירה מודול Python המכיל את עטיפות Python עבור ה-op והקרנל. לפיכך, לאחר שבנית את ה-Op, אתה יכול לעשות את הפעולות הבאות כדי להפעיל אותו מ-Python:

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
print(zero_out_module.zero_out([[1, 2], [3, 4]]).numpy())

# Prints
array([[1, 0], [0, 0]], dtype=int32)

זכור, לפונקציה שנוצרת יינתן שם נחש_מקרה (כדי לעמוד ב- PEP8 ). לכן, אם האופציה שלך נקראת ZeroOut בקבצי C++, הפונקציה python תיקרא zero_out .

כדי להפוך את ה-op לזמין בתור פונקציה רגילה import ממודול Python, אולי שימושי לקיים את הקריאה load_op_library בקובץ מקור של Python באופן הבא:

import tensorflow as tf

zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_out

ודא שהאופציה עובדת

דרך טובה לוודא שיישמת בהצלחה את האופציה שלך היא לכתוב עבורו מבחן. צור את הקובץ zero_out_op_test.py עם התוכן:

import tensorflow as tf

class ZeroOutTest(tf.test.TestCase):
  def testZeroOut(self):
    zero_out_module = tf.load_op_library('./zero_out.so')
    with self.test_session():
      result = zero_out_module.zero_out([5, 4, 3, 2, 1])
      self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])

if __name__ == "__main__":
  tf.test.main()

לאחר מכן הפעל את הבדיקה שלך (בהנחה שהותקן לך tensorflow):

$ python zero_out_op_test.py

בנה תכונות מתקדמות לתוך האופציה שלך

עכשיו כשאתה יודע איך לבנות אופציה בסיסית (והגבלה במקצת) והטמעה, נסתכל על כמה מהדברים היותר מסובכים שבדרך כלל תצטרך לבנות באופ שלך. זה כולל:

בדיקות ותיקוף על תנאי

הדוגמה לעיל הניחה שהאופ חל על טנזור מכל צורה. מה אם זה חל רק על וקטורים? המשמעות היא הוספת המחאה למימוש OpKernel לעיל.

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),
                errors::InvalidArgument("ZeroOut expects a 1-D vector."));
    // ...
  }

זה טוען שהקלט הוא וקטור, וחוזר לאחר הגדרת המצב InvalidArgument אם לא. המאקרו OP_REQUIRES לוקח שלושה ארגומנטים:

לחלופין, אם ברצונך לבדוק אם אובייקט Status שהוחזר מפונקציה כלשהי הוא שגיאה, ואם כן החזר אותו, השתמש ב- OP_REQUIRES_OK . שתי פקודות המאקרו הללו חוזרות מהפונקציה על שגיאה.

אופ רישום

Attrs

Ops יכול להיות attrs, שהערכים שלהם נקבעים כאשר ה-op מתווסף לגרף. אלה משמשים כדי להגדיר את ה-op, וניתן לגשת לערכיהם הן ביישום הליבה והן בסוגי הקלט והפלטים ברישום ה-op. העדיפו להשתמש בקלט במקום ב-attr במידת האפשר, מכיוון שהקלטות גמישות יותר. הסיבה לכך היא ש-attrs הם קבועים ויש להגדיר אותם בזמן בניית הגרף. לעומת זאת, תשומות הן Tensors שהערכים שלהן יכולים להיות דינמיים; כלומר, כניסות יכולות לשנות כל שלב, להיות מוגדרות באמצעות פיד וכו'. Attrs משמשות לדברים שלא ניתן לעשות עם כניסות: כל תצורה שמשפיעה על החתימה (מספר או סוג כניסות או יציאות) או שיכולה' לא לשנות משלב לשלב.

אתה מגדיר attr כשאתה רושם את ה-op, על ידי ציון השם והסוג שלו בשיטת Attr , שמצפה למפרט של הטופס:

<name>: <attr-type-expr>

כאשר <name> מתחיל באות ויכול להיות מורכב מתווים אלפאנומריים וקווים תחתונים, ו- <attr-type-expr> הוא ביטוי מסוג של הצורה המתוארת להלן .

לדוגמה, אם תרצה שהאופ ZeroOut ישמר אינדקס שצוין על ידי המשתמש, במקום רק האלמנט ה-0, תוכל לרשום את האופ כך:

REGISTER_OP("ZeroOut")
    .Attr("preserve_index: int")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

(שים לב שקבוצת סוגי התכונות שונה מה- tf.DType המשמש לכניסות ויציאות.)

הגרעין שלך יכול לגשת ל-attr הזה בבנאי שלו באמצעות פרמטר context :

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
    // Get the index of the value to preserve
    OP_REQUIRES_OK(context,
                   context->GetAttr("preserve_index", &preserve_index_));
    // Check that preserve_index is positive
    OP_REQUIRES(context, preserve_index_ >= 0,
                errors::InvalidArgument("Need preserve_index >= 0, got ",
                                        preserve_index_));
  }
  void Compute(OpKernelContext* context) override {
    // ...
  }
 private:
  int preserve_index_;
};

אשר לאחר מכן ניתן להשתמש בשיטת Compute :

  void Compute(OpKernelContext* context) override {
    // ...

    // We're using saved attr to validate potentially dynamic input
    // So we check that preserve_index is in range
    OP_REQUIRES(context, preserve_index_ < input.dimension(0),
                errors::InvalidArgument("preserve_index out of range"));

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the requested input value
    output_flat(preserve_index_) = input(preserve_index_);
  }

סוגים של Attr

הסוגים הבאים נתמכים ב-attr:

  • string : כל רצף של בתים (לא חייב להיות UTF8).
  • int : מספר שלם חתום.
  • float : מספר נקודה צפה.
  • bool : נכון או לא נכון.
  • type : אחד מהערכים (שאינם ר') של DataType .
  • shape : TensorShapeProto .
  • list(<type>) : רשימה של <type> , כאשר <type> הוא אחד מהסוגים לעיל. שים לב ש- list(list(<type>)) אינו חוקי.

ראה גם: op_def_builder.cc:FinalizeAttr לרשימה סופית.

ערכי ברירת מחדל ואילוצים

ל-attrs עשויים להיות ערכי ברירת מחדל, ולסוגים מסוימים של attrs יכולים להיות אילוצים. כדי להגדיר attr עם אילוצים, אתה יכול להשתמש ב- <attr-type-expr> הבאות:

{'<string1>', '<string2>'} : הערך חייב להיות מחרוזת בעלת הערך <string1> או <string2> . שם הסוג, string , משתמע כאשר אתה משתמש בתחביר זה. זה מחקה מנה:

REGISTER_OP("EnumExample")
    .Attr("e: {'apple', 'orange'}");

{<type1>, <type2>} : הערך הוא מסוג type , וחייב להיות אחד מ- <type1> או <type2> , כאשר <type1> ו- <type2> נתמכים tf.DType . אתה לא מציין שסוג ה-attr הוא type . זה משתמע כאשר יש לך רשימה של סוגים ב {...} . לדוגמה, במקרה זה ה-attr t הוא סוג שחייב להיות int32 , float או bool :

REGISTER_OP("RestrictedTypeExample")
    .Attr("t: {int32, float, bool}");

ישנם קיצורי דרך לאילוצי סוג נפוצים:

  • numbertype : type הסוג מוגבל לסוגים המספריים (ללא מחרוזת ולא בול).
  • realnumbertype : כמו numbertype ללא טיפוסים מורכבים.
  • quantizedtype : כמו numbertype אבל רק סוגי המספרים הקוונטיים.

הרשימות הספציפיות של טיפוסים המותרות על ידי אלה מוגדרות על ידי הפונקציות (כמו NumberTypes() ) ב- tensorflow/core/framework/types.h . בדוגמה זו ה-attr t חייב להיות אחד מהסוגים המספריים:

REGISTER_OP("NumberType")
    .Attr("t: numbertype");

עבור המבצע הזה:

tf.number_type(t=tf.int32)  # Valid
tf.number_type(t=tf.bool)   # Invalid

ניתן לשלב רשימות עם רשימות אחרות וסוגים בודדים. האופ הבא מאפשר ל-attr t להיות כל אחד מהסוגים המספריים, או סוג bool:

REGISTER_OP("NumberOrBooleanType")
    .Attr("t: {numbertype, bool}");

עבור המבצע הזה:

tf.number_or_boolean_type(t=tf.int32)  # Valid
tf.number_or_boolean_type(t=tf.bool)   # Valid
tf.number_or_boolean_type(t=tf.string) # Invalid

int >= <n> : הערך חייב להיות int שהערך שלו גדול או שווה ל- <n> , כאשר <n> הוא מספר טבעי. לדוגמה, רישום הפעולה הבא מציין שה-attr a חייב להיות בעל ערך שהוא לפחות 2 :

REGISTER_OP("MinIntExample")
    .Attr("a: int >= 2");

list(<type>) >= <n> : רשימה מסוג <type> שאורכה גדול או שווה ל- <n> . לדוגמה, רישום ה-op הבא מציין שה-attr a הוא רשימה של סוגים (או int32 או float ), ושחייבים להיות לפחות 3 מהם:

REGISTER_OP("TypeListExample")
    .Attr("a: list({int32, float}) >= 3");

כדי להגדיר ערך ברירת מחדל עבור attr (מה שהופך אותו לאופציונלי בקוד שנוצר), הוסף את = <default> לסוף, כמו ב:

REGISTER_OP("AttrDefaultExample")
    .Attr("i: int = 0");

בנוסף, ניתן לציין גם אילוץ וגם ערך ברירת מחדל:

REGISTER_OP("AttrConstraintAndDefaultExample")
    .Attr("i: int >= 1 = 1");

התחביר הנתמך של ערך ברירת המחדל הוא מה שישמש בייצוג הפרוטו של הגדרת GraphDef המתקבלת.

להלן דוגמאות כיצד לציין ברירת מחדל עבור כל הסוגים:

REGISTER_OP("AttrDefaultExampleForAllTypes")
   .Attr("s: string = 'foo'")
   .Attr("i: int = 0")
   .Attr("f: float = 1.0")
   .Attr("b: bool = true")
   .Attr("ty: type = DT_INT32")
   .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
   .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
   .Attr("l_empty: list(int) = []")
   .Attr("l_int: list(int) = [2, 3, 5, 7]");

שים לב במיוחד שהערכים מסוג type משתמשים tf.DType .

רב צורתיות

סוג פולימורפיזם

עבור ops שיכולים לקחת סוגים שונים כקלט או לייצר סוגי פלט שונים, אתה יכול לציין attr בסוג קלט או פלט ברישום op. בדרך כלל, אז היית רושם OpKernel עבור כל סוג נתמך.

לדוגמה, אם תרצה שהאופ ZeroOut יעבוד על s float בנוסף ל- int32 s, רישום המבצע שלך עשוי להיראות כך:

REGISTER_OP("ZeroOut")
    .Attr("T: {float, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

רישום ה-op שלך מציין כעת שסוג הקלט חייב להיות float , או int32 , ושהפלט שלו יהיה אותו סוג, מכיוון שלשניהם יש סוג T .

שִׁיוּם

יש לתת לכניסות, פלטים ו-attrs בדרך כלל שמות נחש_מקרה. החריג היחיד הוא attrs המשמשים כסוג של קלט או בסוג של פלט. ניתן להסיק את המאפיינים הללו כאשר ה-op נוסף לגרף ולכן אינם מופיעים בפונקציה של ה-op. לדוגמה, ההגדרה האחרונה של ZeroOut תיצור פונקציית Python שנראית כך:

def zero_out(to_zero, name=None):
  """...
  Args:
    to_zero: A `Tensor`. Must be one of the following types:
        `float32`, `int32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `to_zero`.
  """

אם מועבר to_zero טנסור int32 , אז T מוגדר אוטומטית ל- int32 (טוב, למעשה DT_INT32 ). השמות המשוערים הללו מקבלים שמות באותיות רישיות או CamelCase.

השווה זאת לאופ שיש לו סוג attr שקובע את סוג הפלט:

REGISTER_OP("StringToNumber")
    .Input("string_tensor: string")
    .Output("output: out_type")
    .Attr("out_type: {float, int32} = DT_FLOAT");
    .Doc(R"doc(
Converts each string in the input Tensor to the specified numeric type.
)doc");

במקרה זה, המשתמש צריך לציין את סוג הפלט, כמו ב-Python שנוצר:

def string_to_number(string_tensor, out_type=None, name=None):
  """Converts each string in the input Tensor to the specified numeric type.

  Args:
    string_tensor: A `Tensor` of type `string`.
    out_type: An optional `tf.DType` from: `tf.float32, tf.int32`.
      Defaults to `tf.float32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor` of type `out_type`.
  """
הקלד דוגמה לפולימורפיזם
#include "tensorflow/core/framework/op_kernel.h"

class ZeroOutInt32Op : public OpKernel {
  // as before
};

class ZeroOutFloatOp : public OpKernel {
 public:
  explicit ZeroOutFloatOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<float>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<float>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutInt32Op);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutFloatOp);

כדי לשמור על תאימות לאחור , עליך לציין ערך ברירת מחדל בעת הוספת attr לאופציה קיימת:

REGISTER_OP("ZeroOut")
  .Attr("T: {float, int32} = DT_INT32")
  .Input("to_zero: T")
  .Output("zeroed: T")

נניח שרצית להוסיף עוד סוגים, אמור double :

REGISTER_OP("ZeroOut")
    .Attr("T: {float, double, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

במקום לכתוב עוד OpKernel עם קוד מיותר כמו לעיל, לעתים קרובות תוכל להשתמש בתבנית C++ במקום זאת. עדיין תהיה לך רישום ליבה אחד (קריאת REGISTER_KERNEL_BUILDER ) לכל עומס יתר.

template <typename T>
class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<T>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<T>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutOp<int32>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    ZeroOutOp<double>);

אם יש לך יותר מכמה עומסים, אתה יכול לשים את הרישום במאקרו.

#include "tensorflow/core/framework/op_kernel.h"

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

REGISTER_KERNEL(int32);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);

#undef REGISTER_KERNEL

בהתאם לרשימת הסוגים עבורם אתה רושם את הליבה, ייתכן שתוכל להשתמש במאקרו המסופק על ידי tensorflow/core/framework/register_types.h :

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"

REGISTER_OP("ZeroOut")
    .Attr("T: realnumbertype")
    .Input("to_zero: T")
    .Output("zeroed: T");

template <typename T>
class ZeroOutOp : public OpKernel { ... };

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);

#undef REGISTER_KERNEL
רשימת כניסות ויציאות

בנוסף ליכולת לקבל או לייצר סוגים שונים, מבצעים יכולים לצרוך או לייצר מספר משתנה של טנזורים.

בדוגמה הבאה, ה-attr T מחזיק רשימה של טיפוסים, והוא משמש כסוג של ה-input in out out . הקלט והפלט הם רשימות של טנזורים מסוג זה (ומספר וסוגי הטנסורים בפלט זהים לקלט, שכן לשניהם יש סוג T ).

REGISTER_OP("PolymorphicListExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

אתה יכול גם להציב הגבלות על הסוגים שניתן לציין ברשימה. במקרה הבא הזה, הקלט הוא רשימה של טנסורים float double . ה-op מקבל, למשל, סוגי קלט (float, double, float) ובמקרה כזה סוג הפלט יהיה גם (float, double, float) .

REGISTER_OP("ListTypeRestrictionExample")
    .Attr("T: list({float, double})")
    .Input("in: T")
    .Output("out: T");

אם אתה רוצה שכל הטנזורים ברשימה יהיו מאותו סוג, אתה יכול לעשות משהו כמו:

REGISTER_OP("IntListInputExample")
    .Attr("N: int")
    .Input("in: N * int32")
    .Output("out: int32");

זה מקבל רשימה של טנסורים int32 , ומשתמש ב- int attr N כדי לציין את אורך הרשימה.

זה יכול להיעשות גם מסוג פולימורפי . בדוגמה הבאה, הקלט הוא רשימה של טנסורים (עם אורך "N" ) מאותו סוג (אך לא מוגדר) ( "T" ), והפלט הוא טנסור בודד מסוג תואם:

REGISTER_OP("SameListInputExample")
    .Attr("N: int")
    .Attr("T: type")
    .Input("in: N * T")
    .Output("out: T");

כברירת מחדל, לרשימות טנסור יש אורך מינימלי של 1. אתה יכול לשנות את ברירת המחדל הזו באמצעות אילוץ ">=" ב-attr המתאים . בדוגמה הבאה זו, הקלט הוא רשימה של לפחות 2 טנסורים int32 :

REGISTER_OP("MinLengthIntListExample")
    .Attr("N: int >= 2")
    .Input("in: N * int32")
    .Output("out: int32");

אותו תחביר עובד עם "list(type)" attrs:

REGISTER_OP("MinimumLengthPolymorphicListExample")
    .Attr("T: list(type) >= 3")
    .Input("in: T")
    .Output("out: T");

כניסות ויציאות

לסיכום האמור לעיל, רישום הפעלה יכול לכלול מספר כניסות ויציאות:

REGISTER_OP("MultipleInsAndOuts")
    .Input("y: int32")
    .Input("z: float")
    .Output("a: string")
    .Output("b: int32");

כל מפרט קלט או פלט הוא בצורה:

<name>: <io-type-expr>

כאשר <name> מתחיל באות ויכול להיות מורכב מתווים אלפאנומריים וקווים תחתונים. <io-type-expr> הוא אחד מביטויי הסוג הבאים:

  • <type> , כאשר <type> הוא סוג קלט נתמך (למשל float , int32 , string ). זה מציין טנזור בודד מהסוג הנתון.

    ראה tf.DType .

    REGISTER_OP("BuiltInTypesExample")
        .Input("integers: int32")
        .Input("complex_numbers: complex64");
    
  • <attr-type> , כאשר <attr-type> הוא השם של Attr עם סוג type או list(type) (עם הגבלת סוג אפשרית). תחביר זה מאפשר פעולות פולימורפיות .

    REGISTER_OP("PolymorphicSingleInput")
        .Attr("T: type")
        .Input("in: T");
    
    REGISTER_OP("RestrictedPolymorphicSingleInput")
        .Attr("T: {int32, int64}")
        .Input("in: T");
    

    הפניה ל-attr מסוג list(type) מאפשרת לקבל רצף של טנסורים.

    REGISTER_OP("ArbitraryTensorSequenceExample")
        .Attr("T: list(type)")
        .Input("in: T")
        .Output("out: T");
    
    REGISTER_OP("RestrictedTensorSequenceExample")
        .Attr("T: list({int32, int64})")
        .Input("in: T")
        .Output("out: T");
    

    שים לב שמספר וסוגי הטנסורים בפלט out זהים לזו בקלט in , מכיוון ששניהם מסוג T

  • עבור רצף של טנסורים עם אותו סוג: <number> * <type> , כאשר <number> הוא השם של Attr עם הסוג int . ה- <type> יכול להיות tf.DType , או השם של attr עם type type . כדוגמה לראשון, אופציה זו מקבלת רשימה של טנסורים int32 :

    REGISTER_OP("Int32SequenceExample")
        .Attr("NumTensors: int")
        .Input("in: NumTensors * int32")
    

    ואילו אופציה זו מקבלת רשימה של טנסורים מכל סוג, כל עוד כולם זהים:

    REGISTER_OP("SameTypeSequenceExample")
        .Attr("NumTensors: int")
        .Attr("T: type")
        .Input("in: NumTensors * T")
    
  • להתייחסות לטנזור: Ref(<type>) , כאשר <type> הוא אחד מהסוגים הקודמים.

כל attr המשמש בסוג הקלט יוסק. לפי מוסכמה, אותם מושקים משתמשים בשמות רישיות (כמו T או N ). אחרת לכניסות, פלטים ו-attrs יש שמות כמו פרמטרי פונקציה (למשל num_outputs ). לפרטים נוספים, עיין בסעיף המוקדם בנושא מתן שמות .

לפרטים נוספים, ראה tensorflow/core/framework/op_def_builder.h .

תאימות לאחור

נניח שכתבת אופציה נחמדה ומותאמת אישית ושיתפת אותה עם אחרים, כך שיש לך לקוחות מרוצים שמשתמשים במבצע שלך. עם זאת, תרצה לבצע שינויים באופציה בדרך כלשהי.

באופן כללי, שינויים במפרטים קיימים וצ'ק-אין חייבים להיות תואמים לאחור: שינוי המפרט של הפעלה אסור לשבור מאגרים קודמים של פרוטוקול GraphDef שנבנו ממפרטים ישנים יותר. הפרטים של תאימות GraphDef מתוארים כאן .

ישנן מספר דרכים לשמור על תאימות לאחור.

  1. כל attrs חדש שמתווסף לפעולה חייב להיות מוגדר עם ערכי ברירת מחדל, ועם ערך ברירת המחדל הזה ל-op חייב להיות ההתנהגות המקורית. כדי לשנות פעולה מלא פולימורפית לפולימורפית, עליך לתת ערך ברירת מחדל לסוג החדש attr כדי לשמר את החתימה המקורית כברירת מחדל. לדוגמה, אם הפעולה שלך הייתה:

    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: float")
        .Output("out: float");
    

    אתה יכול להפוך אותו לפולימורפי בצורה תואמת לאחור באמצעות:

    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: T")
        .Output("out: T")
        .Attr("T: numerictype = DT_FLOAT");
    
  2. אתה יכול בבטחה להפוך אילוץ על אטרק פחות מגביל. לדוגמה, אתה יכול לשנות מ- {int32, int64} ל- {int32, int64, float} או type . או שאתה יכול לשנות מ- {"apple", "orange"} ל- {"apple", "banana", "orange"} או string .

  3. אתה יכול לשנות כניסות / יציאות בודדות לכניסות / יציאות רשימה, כל עוד ברירת המחדל של סוג הרשימה תואמת את החתימה הישנה.

  4. אתה יכול להוסיף קלט/פלט חדש לרשימה, אם ברירת המחדל היא ריקה.

  5. מרחב שמות לכל אופציות חדשות שאתה יוצר, על ידי הקדמת שמות ההפעלה עם משהו ייחודי לפרויקט שלך. זה ימנע את התנגשות ההפעלה שלך עם כל אופציה שיכולה להיכלל בגרסאות עתידיות של TensorFlow.

  6. לתכנן מראש! נסו לצפות שימושים עתידיים באופ. חלק מהשינויים בחתימה לא יכולים להיעשות בצורה תואמת (לדוגמה, הפיכת רשימה מאותו סוג לרשימה של סוגים שונים).

את הרשימה המלאה של שינויים בטוחים ולא בטוחים ניתן למצוא ב- tensorflow/core/framework/op_compatibility_test.cc . אם אינך יכול להפוך את השינוי שלך לפעולה תואמת לאחור, צור פעולה חדשה עם שם חדש עם הסמנטיקה החדשה.

שים לב גם שבעוד ששינויים אלה יכולים לשמור על תאימות GraphDef , קוד Python שנוצר עשוי להשתנות בצורה שאינה תואמת למתקשרים ישנים. ה-API של Python עשוי להישמר תואם על ידי שינויים זהירים בעטיפת Python בכתב יד, על ידי שמירה על החתימה הישנה למעט הוספת ארגומנטים אופציונליים חדשים עד הסוף. שינויים שאינם תואמים בדרך כלל עשויים להתבצע רק כאשר TensorFlow משנה גרסאות עיקריות, וחייבים להתאים לסמנטיקה של גרסת GraphDef .

תמיכה ב-GPU

אתה יכול ליישם OpKernels שונים ולרשום אחד עבור CPU ואחר עבור GPU, בדיוק כמו שאתה יכול לרשום ליבות עבור סוגים שונים . ישנן מספר דוגמאות של גרעינים עם תמיכה ב-GPU ב- tensorflow/core/kernels/ . שימו לב לחלק מהגרעינים יש גרסת CPU בקובץ .cc , גרסת GPU בקובץ המסתיים ב _gpu.cu.cc וקוד מסוים משותף בקובץ .h .

לדוגמה, ל- tf.pad יש הכל מלבד ליבת ה-GPU ב- tensorflow/core/kernels/pad_op.cc . ליבת ה-GPU נמצאת ב- tensorflow/core/kernels/pad_op_gpu.cu.cc , והקוד המשותף הוא מחלקה תבניתית המוגדרת ב- tensorflow/core/kernels/pad_op.h . אנו מארגנים את הקוד כך משתי סיבות: הוא מאפשר לך לחלוק קוד משותף בין מימושי המעבד וה-GPU, והוא מכניס את מימוש ה-GPU לקובץ נפרד כך שניתן יהיה להדר אותו רק על ידי מהדר ה-GPU.

יש לציין דבר אחד, גם כאשר משתמשים בגרסת ליבת ה-GPU של pad , הוא עדיין זקוק לקלט "paddings" שלו בזיכרון המעבד. כדי לסמן שהכניסות או הפלטים נשמרים במעבד, הוסף קריאה HostMemory() לרישום הליבה, למשל:

#define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                              .Device(DEVICE_GPU)      \
                              .TypeConstraint<T>("T")  \
                              .HostMemory("paddings"), \
                          PadOp<GPUDevice, T>)

קומפילציה של הליבה עבור מכשיר ה-GPU

עיין ב- cuda_op_kernel.cu.cc עבור דוגמה שמשתמשת בליבת CUDA כדי ליישם אופ. tf_custom_op_library מקבל ארגומנט gpu_srcs שבו ניתן לציין את רשימת קובצי המקור המכילים את ליבות ה-CUDA (קבצי *.cu.cc ). לשימוש עם התקנה בינארית של TensorFlow, גרעיני ה-CUDA צריכים להיות קומפילציה עם מהדר nvcc של NVIDIA. להלן רצף הפקודות שבהן אתה יכול להשתמש כדי לקמפל את ה- cuda_op_kernel.cu.cc ו- cuda_op_kernel.cc לספרייה אחת הניתנת לטעינה דינמית:

nvcc -std=c++14 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
  ${TF_CFLAGS[@]} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

g++ -std=c++14 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
  cuda_op_kernel.cu.o ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

ניתן לטעון cuda_op_kernel.so שהופק לעיל כרגיל ב- Python, באמצעות הפונקציה tf.load_op_library .

שים לב שאם ספריות ה-CUDA שלך אינן מותקנות ב- /usr/local/lib64 , תצטרך לציין את הנתיב במפורש בפקודה השנייה (g++) למעלה. לדוגמה, הוסף -L /usr/local/cuda-8.0/lib64/ אם ה-CUDA שלך מותקן ב- /usr/local/cuda-8.0 .

יישם את הגרדיאנט ב-Python

בהינתן גרף של פעולות, TensorFlow משתמש בדיפרנציאציה אוטומטית (הפצה לאחור) כדי להוסיף פעולות חדשות המייצגות שיפועים ביחס לאופציות הקיימות. כדי שהבידול האוטומטי יעבוד עבור אופציות חדשות, עליך לרשום פונקציית שיפוע אשר מחשבת שיפועים בהתייחס לכניסות האופציות הנתונות לשיפועים ביחס לפלטים של המבצעים.

מבחינה מתמטית, אם אופ מחשב את \(y = f(x)\) השיפוע הרשום op ממיר שיפועים \(\partial L/ \partial y\) של אובדן \(L\) ביחס ל-\(y\) להדרגות \(\partial L/ \partial x\) ביחס ל- \(x\) באמצעות כלל השרשרת:

\[\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial f}{\partial x}.\]

במקרה של ZeroOut , רק כניסה אחת בקלט משפיעה על הפלט, כך שהשיפוע ביחס לקלט הוא טנסור "חם אחד" דליל. הדבר מתבטא כך:

from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops

@ops.RegisterGradient("ZeroOut")
def _zero_out_grad(op, grad):
  """The gradients for `zero_out`.

  Args:
    op: The `zero_out` `Operation` that we are differentiating, which we can use
      to find the inputs and outputs of the original op.
    grad: Gradient with respect to the output of the `zero_out` op.

  Returns:
    Gradients with respect to the input of `zero_out`.
  """
  to_zero = op.inputs[0]
  shape = array_ops.shape(to_zero)
  index = array_ops.zeros_like(shape)
  first_grad = array_ops.reshape(grad, [-1])[0]
  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)
  return [to_zero_grad]  # List of one Tensor, since we have one input

פרטים על רישום פונקציות מעבר עם tf.RegisterGradient :

  • עבור אופ עם פלט אחד, פונקציית השיפוע תקח tf.Operation , op ו- tf.Tensor grad ותבנה אופסים חדשים מהטנסורים op.inputs[i] , op.outputs[i] ו- grad . ניתן למצוא מידע על כל attrs דרך tf.Operation.get_attr .

  • אם ל-op יש פלטים מרובים, פונקציית ההדרגה תיקח op ו- grads , כאשר grads היא רשימה של מעברים ביחס לכל פלט. התוצאה של פונקציית ההדרגה חייבת להיות רשימה של אובייקטי Tensor המייצגים את ההדרגות ביחס לכל קלט.

  • אם אין שיפוע מוגדר היטב עבור קלט כלשהו, ​​כגון עבור תשומות שלמים המשמשים כמדדים, השיפוע המוחזר המתאים צריך להיות None . לדוגמה, עבור אופ שלוקח טנסור נקודה צפה x ואינדקס שלם i , פונקציית הגרדיאנט return [x_grad, None] .

  • אם אין שיפוע משמעותי עבור הניתוח, לעתים קרובות לא תצטרכו לרשום שום שיפוע, וכל עוד אין צורך בשיפוע של הניתוח, אתה תהיה בסדר. במקרים מסוימים, לאופ אין שיפוע מוגדר היטב אך יכול להיות מעורב בחישוב השיפוע. כאן אתה יכול להשתמש ops.NotDifferentiable כדי להפיץ אפסים אוטומטית לאחור.

שים לב שבזמן שבו נקראת פונקציית הגרדיאנט, רק גרף זרימת הנתונים של פעולות זמין, לא נתוני הטנזור עצמם. לפיכך, כל החישוב חייב להתבצע באמצעות אופציות זרימת טנסור אחרות, שיופעלו בזמן ביצוע הגרף.

הוסף רמזים לסוג בעת רישום הגרדיאנט המותאם אישית עבור סוג הפעלה כדי להפוך את הקוד לקריא יותר, ניתן לניפוי באגים, קל יותר לתחזוקה וחזק יותר באמצעות אימות נתונים. לדוגמה, כאשר לוקחים op כפרמטר בפונקציה, ציין שפונקציית הגרדיאנט תקבל tf.Operation כסוג הפרמטר שלה.

פונקציות צורה ב-C++

לממשק API של TensorFlow יש תכונה הנקראת "הסקת צורה" המספקת מידע על צורות הטנזורים ללא צורך בביצוע הגרף. הסקת צורה נתמכת על ידי "פונקציות צורה" הרשומות עבור כל סוג אופ בהצהרת C++ REGISTER_OP , ומבצעות שני תפקידים: קביעה שהצורות של התשומות תואמות במהלך בניית הגרף, וציון הצורות עבור הפלטים.

פונקציות Shape מוגדרות כפעולות במחלקה shape_inference::InferenceContext . לדוגמה, בפונקציית הצורה עבור ZeroOut:

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

c->set_output(0, c->input(0)); מכריז שיש להגדיר את צורת הפלט הראשון לצורת הקלט הראשון. אם הפלט נבחר על ידי האינדקס שלו כמו בדוגמה לעיל, הפרמטר השני של set_output צריך להיות אובייקט ShapeHandle . אתה יכול ליצור אובייקט ShapeHandle ריק על ידי בנאי ברירת המחדל שלו. ניתן להשיג את האובייקט ShapeHandle עבור קלט עם index idx על ידי c->input(idx) .

ישנן מספר פונקציות צורה נפוצות החלות על פעולות רבות, כגון shape_inference::UnchangedShape שניתן למצוא ב- common_shape_fns.h ולהשתמש בה באופן הבא:

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);

פונקציית צורה יכולה גם להגביל את הצורה של קלט. עבור הגרסה של ZeroOut עם אילוץ צורה וקטורית , פונקציית הצורה תהיה כדלקמן:

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input));
      c->set_output(0, input);
      return Status::OK();
    });

הקריאה WithRank מאמתת שלצורת הקלט c->input(0) יש צורה בעלת ממד אחד בדיוק (או אם צורת הקלט אינה ידועה, צורת הפלט תהיה וקטור עם ממד אחד לא ידוע).

אם האופציה שלך היא פולימורפית עם מספר כניסות , אתה יכול להשתמש באיברים של InferenceContext כדי לקבוע את מספר הצורות לבדיקה, ו- Merge כדי לאמת שהצורות כולן תואמות (לחלופין, גישה לתכונות שמציינות את האורכים, עם InferenceContext::GetAttr , המספק גישה לתכונות של ה-op).

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      ::tensorflow::shape_inference::ShapeHandle output;
      for (size_t i = 0; i < c->num_inputs(); ++i) {
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 2, &input));
        TF_RETURN_IF_ERROR(c->Merge(output, input, &output));
      }
      c->set_output(0, output);
      return Status::OK();
    });

מכיוון שהסקת צורה היא תכונה אופציונלית, והצורות של הטנזורים עשויות להשתנות באופן דינמי, פונקציות הצורה חייבות להיות חזקות למידע צורה לא שלם עבור כל אחת מהכניסות. שיטת Merge ב- InferenceContext מאפשרת למתקשר לטעון ששתי צורות זהות, גם אם לאחת מהן או לשתיהן אין מידע מלא. פונקציות צורה מוגדרות עבור כל האופציות המרכזיות של TensorFlow ומספקות דוגמאות שימוש רבות ושונות.

למחלקה InferenceContext יש מספר פונקציות שניתן להשתמש בהן כדי להגדיר מניפולציות של פונקציות צורה. לדוגמה, אתה יכול לאמת שלמימד מסוים יש ערך מאוד ספציפי באמצעות InferenceContext::Dim ו- InferenceContext::WithValue ; אתה יכול לציין שממד פלט הוא הסכום/מכפלה של שני ממדי קלט באמצעות InferenceContext::Add ו- InferenceContext::Multiply . עיין במחלקה InferenceContext עבור כל מניפולציות הצורות השונות שאתה יכול לציין. הדוגמה הבאה מגדירה את צורת הפלט הראשון ל-(n, 3), כאשר לקלט הראשון יש צורה (n, ...)

.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));
    return Status::OK();
});

אם יש לך פונקציית צורה מסובכת, עליך לשקול להוסיף מבחן לאימות ששילובי צורות קלט שונים מייצרים את שילובי צורות הפלט הצפויים. תוכל לראות דוגמאות כיצד לכתוב מבחנים אלה בחלק מבחני הליבה שלנו. (התחביר של INFER_OK ו- INFER_ERROR הם קצת קריפטיים, אבל נסו להיות קומפקטיים בייצוג מפרטי צורות קלט ופלט בבדיקות. לעת עתה, ראה את ההערות שמסביב במבחנים אלה כדי לקבל תחושה של מפרט מחרוזת הצורה).

בנה חבילת pip עבור המבצע המותאם אישית שלך

כדי לבנות חבילת pip עבור ההפעלה שלך, עיין בדוגמה של tensorflow/custom-op . מדריך זה מראה כיצד לבנות אופציות מותאמות אישית מחבילת TensorFlow pip במקום לבנות את TensorFlow מהמקור.