相關(guān)閱讀
TensorFlow Lite源碼解析之一

1. 前言

愛迪生說過，人工智能就是是百分之九十九的數(shù)據(jù)加上百分之一的算法。畢竟目前人工智能還沒有達(dá)到T800這種以毀滅人類為己任的終結(jié)者級別，歸根到底還是一個程序。這么一想，是不是覺得市面上說的AI要統(tǒng)治人類了根本就是危言聳聽，對于弱人工智能，我治不住你，難道我的360強(qiáng)力卸載還治不住你？

言歸正傳。很顯然，想要了解一個程序，理解它是怎么管理用于存儲數(shù)據(jù)的內(nèi)存是一個繞不開的話題。想要了解TensorFlow Lite是如何工作的，我們首先要弄清楚它的Tensors都會Flow去哪。如果Tensor是水的話，那么內(nèi)存分配的過程就是挖溝的過程。對于TFLite這種用于端測推理的推理引擎，在內(nèi)存使用上也不能像服務(wù)器那么豪橫，總之一句話，就是要努力做到又要馬兒跑，又不讓馬吃草。

2. 太長；不看

你的時間非常值錢，打開手機(jī)肯定不是為了聽我在這嗶嗶嗶，內(nèi)存分配其實是個繁瑣的過程，為了節(jié)省時間先說說結(jié)論。當(dāng)然，不著急的話看完也是極好的。

結(jié)論就是，TFLite首先會根據(jù)每個張量的大小(size)，為它們分配一個偏移地址(offset)，并且保證不會有任何一個張量的數(shù)據(jù)在錯誤的時間覆蓋任何其他有用的張量。并且，TFlite能夠做到用于中間結(jié)果的內(nèi)存總大小不超過最大張量所用空間 * （1+最大分值數(shù)），這是最壞的情況，實際情況所使用的空間可能更小。對于分支少的模型，節(jié)省的內(nèi)存非常可觀。這極大的緩解的了移動端設(shè)備內(nèi)存的壓力。簡單地說就是內(nèi)存復(fù)用。

之后，通過最后一個張量的偏移 + 大小 + 必要的首尾填充，得到一個總的內(nèi)存大小。一次性向系統(tǒng)請求計算出所需的內(nèi)存，得到一個實際的內(nèi)存的起始地址。

最后，依次將每個張量的偏移地址加上實際申請得到的內(nèi)存起始地址，就得到了每個張量數(shù)據(jù)實際的起始地址，結(jié)合張量的大小，最終就可以確定內(nèi)存中哪塊區(qū)域?qū)儆谀膫€特定的張量。

3. 詳細(xì)過程

首先，讓我們站在高處，先做個總體的認(rèn)識。如圖1所示，為TFLite進(jìn)行內(nèi)存分配的時候的主要的函數(shù)調(diào)用流程。主要的分配邏輯都由一個名為ArenaPlanner的類負(fù)責(zé)。

圖1 內(nèi)存分配過程概覽

通常，在代碼中，如果我們看到有Arena這個詞一出現(xiàn)，那么很可能的情況就是程序會通過庫文件或者系統(tǒng)調(diào)用的方式，直接向內(nèi)存索取一大片內(nèi)存，然后再自己分配。在這里，TFLite遵循了這一不成文的規(guī)定，在后面的介紹中我們會知道在TFLite中確實是這么做的。

直接向操作系統(tǒng)索取一大片內(nèi)存再自己做分配的好處顯而易見：避免了多次系統(tǒng)調(diào)用的開銷，因為系統(tǒng)調(diào)用確實不便宜，特別在推理引擎這種對時間要求很高的程序中；另外一個就是程序可以根據(jù)自己的需要定制分配策略以使得效率更高。這就相當(dāng)于申請預(yù)算，算個大概之后就一次性把下一年度的預(yù)算都申請下來，不然在頻繁打報告上花費的時間暫且不論，能不能申請下來還是未知數(shù)。

在圖1中可以看到，TFLite在內(nèi)存分配上大致可以分成以下幾個步驟：

1. 計劃階段。確定整個模型所有的張量中，哪些需要進(jìn)行內(nèi)存分配；
2. 計算階段。計算這些張量所需要的內(nèi)存大小之和，并確定張量之間的相對地址。就類似于雖然錢還沒到手但是我們已經(jīng)算好了怎么分；
3. 實際分配階段。一次性向內(nèi)存申請指定大小的內(nèi)存，并且將需要的數(shù)據(jù)進(jìn)行拷貝。

實際上，TFLite申請的內(nèi)存分成兩塊，一塊用于儲存臨時張量以及其他一些臨時性數(shù)據(jù)，另一塊則用于存儲永久性數(shù)據(jù)。兩塊內(nèi)存不同的是數(shù)據(jù)的生命周期不同，其他的完全一樣。

接下來我們會一一探究每一個過程的的細(xì)節(jié)。

3.1. 計劃階段

ArenaPlanner通過多個列表來做記錄，以輔助內(nèi)存分配。在計劃階段，也就是在PlanAllocations()調(diào)用中，會用到其中三個，他們分別是alloc_node_、dealloc_node_以及refcounts，它們都是std::vector<int>類型的列表，且長度等于當(dāng)前所在子圖所包含的所有張量的數(shù)量（length），由于通常整個模型中只有一個子圖，因此也可以說這些列表的長度等于整個模型中包含的張量個數(shù)。由于每個張量在解析的時候都被賦予了在從0~length-1中唯一的數(shù)最為索引，因此這些列表中的每一個元素和模型中的張量是一一對應(yīng)的。

這三者中前兩個是ArenaPlanner的成員變量，后續(xù)的操作中需要用到它們在此步驟設(shè)置好的值。最后一個refcount是局部變量，主要用于確定dealloc_node_中的元素該如何賦值。

其執(zhí)行過程如下：

1. 在開始的時候，首先對這三個列表進(jìn)行初始化：alloc_node_、dealloc_node_的元素都被初始化為0xFFFFFFFF，refcounts的元素都被初始化為0；
2. 為所有類型不是kTfLiteOptionalTensor的張量都都設(shè)置引用計數(shù)，也就是設(shè)置refcounts中與每一個張量對應(yīng)的元素的值都至少是1；
3. 將alloc_node_中代表輸入張量、輸出張量以及屬于變量張量的元素的值都設(shè)置為0；將屬于模型中節(jié)點的輸出的張量的值設(shè)置為對應(yīng)的節(jié)點的編號；
4. 如果不需要保留中間結(jié)果，則將模型中節(jié)點的輸入張量在dealloc_node_中對應(yīng)的元素的值設(shè)置成該節(jié)點的編號；
5. 對于模型中各個節(jié)點所擁有的臨時張量，同時將他們在alloc_node_以及dealloc_node_中對應(yīng)的元素的值設(shè)置成該節(jié)點的編號。

對于同一個張量（除了模型的輸、輸出張量以及節(jié)點自身的臨時張量），它既是前一個張量的輸出，同時也是后一個（也可能是多個）節(jié)點的輸入。因此，對于alloc_node_以及dealloc_node_中相同位置的元素而言，雖然它們對應(yīng)的是同一個張量，但是值往往不同，其關(guān)系至少滿足alloc_node_[tensor_index] + 1 = dealloc_node_[tensor_index]，如圖2所示。請記住這一點，因為后續(xù)會用到這一關(guān)系。

圖2 對應(yīng)同一個張量的不同值

這一步完成之后，alloc_node_以及dealloc_node_中同一個張量所對應(yīng)的元素的值有三種關(guān)系：

1. 對于輸入、輸出以及變量張量：alloc_node_對應(yīng)的值為0，dealloc_node_對應(yīng)的值為為0xFFFFFFFF；
2. 對于結(jié)算過程中的中間結(jié)果張量：如果需要保留中間結(jié)果，則alloc_node_對應(yīng)的值為輸出該張量的節(jié)點編號，dealloc_node_對應(yīng)的值為為0xFFFFFFFF；
3. 對于結(jié)算過程中的中間結(jié)果張量：如果不需要保留中間結(jié)果，則alloc_node_對應(yīng)的值為輸出該張量的節(jié)點編號，dealloc_node_對應(yīng)的值為滿足alloc_node_[tensor_index] + 1 = dealloc_node_[tensor_index]；劃重點，這是TFLite實現(xiàn)內(nèi)存內(nèi)存可以使用所用內(nèi)存空間遠(yuǎn)小于所有張量內(nèi)存大小之和的關(guān)鍵所在

3.2. 計算階段

每一個張量所擁有的buffer的信息，通過類ArenaAllocWithUsageInterval來保存，ArenaAllocWithUsageInterval包含偏移（offset）、大小、所屬的張量、張量的輸入節(jié)點、輸出節(jié)點等信息，如下面代碼所示。

struct ArenaAllocWithUsageInterval {
  ArenaAllocWithUsageInterval() { reset(); }

  size_t offset;
  size_t size;
  int32_t tensor;
  int32_t first_node;
  int32_t last_node;

  inline void reset() {
    offset = 0;
    size = 0;
    tensor = -1;
    first_node = -1;
    last_node = -1;
  }

  inline bool operator<(const ArenaAllocWithUsageInterval& other) const {
    return offset < other.offset;
  }
};

在計算階段要做的就是確定每一個ArenaAllocWithUsageInterval實例中這些屬性的值以及總的內(nèi)存大小。最終TFLite申請內(nèi)存的時候是按照總的內(nèi)存大小來申請一塊連續(xù)的內(nèi)存，便可以得到這塊內(nèi)存的起始地址，起始地址加上偏移地址就能得到每一個張量所擁有的buffer的起始地址，結(jié)合size屬性就能確定整個buffer的邊界，如下圖3所示。

圖3 張量內(nèi)存分布示例

最終，通過計算，包含所有的張量的buffer信息ArenaAllocWithUsageInterval會按照一定的規(guī)則有序的存儲在一個名叫ordered_allocs_的列表中，另外，也會按照索引號從小到大的順序保存在allocs_列表中。ordered_allocs_以及allocs_存儲的都是ArenaAllocWithUsageInterval的指針。

那么，TFLite具體是怎么確定每個張量的偏移地址的呢？其過程具體如下：

1. 開始的時候，ordered_allocs_會被清空；

2. 為張量的索引值排序，排序的原則為，對于任意兩個張量T1以及T2：
   i. 如果T1、T2兩個都屬于輸入張量或者輸出張量，那么這T1、T2張量維持他們的順序不變；
   ii. 如果T1屬于輸入張量或者輸出張量而T2不是，則T1排在T2之前；
   iii. 如果T1、T2都不輸入輸入張量或者輸出張量，則誰需要的內(nèi)存更多誰排前面；對于需要內(nèi)存一樣的兩個張量，則誰先被用到誰排前面。
   排序的結(jié)果是一個由張量索引值組成的列表，后續(xù)將按照這個列表為張量分配內(nèi)存。前面說到的ordered_allocs_里面元素的順序有大致這么確定，不過可能會略微有區(qū)別（當(dāng)存在一個張量太小，可以插入其他一些張量的空隙，這點在后面有介紹）。

3. 接著，根據(jù)第二步得到的列表依次為張量分配內(nèi)存偏移地址，分配過程如下所示，整個過程很簡單，從頭到尾依次查看是否有空隙容納當(dāng)前張量，如果沒有則再已分配的張量后面進(jìn)行分配。

你肯定好奇，既然每個張量的起始地址都進(jìn)行了內(nèi)存對齊，那么哪來的空間進(jìn)行插入操作？豈不是脫褲子放那啥——多此一舉。別急，下面就是見證奇跡的時刻：

TfLiteStatus SimpleMemoryArena::Allocate(
    TfLiteContext* context, size_t alignment, size_t size, int32_t tensor,
    int32_t first_node, int32_t last_node,
    ArenaAllocWithUsageInterval* new_alloc) {
  TF_LITE_ENSURE(context, alignment <= arena_alignment_);
  new_alloc->tensor = tensor;
  new_alloc->first_node = first_node;
  new_alloc->last_node = last_node;
  new_alloc->size = size;
  if (size == 0) {
    new_alloc->offset = 0;
    return kTfLiteOk;
  }

  // If we don't find a better gap just allocate at the end of the buffer.
  const size_t kOffsetNotAssigned = std::numeric_limits<size_t>::max();
  size_t best_offset = kOffsetNotAssigned;
  size_t best_offset_fit = kOffsetNotAssigned;

  // Go through the sorted allocs and look at the gaps between them.
  size_t current_offset = 0;
  for (const auto& alloc : ordered_allocs_) {
    if (alloc.last_node < first_node || alloc.first_node > last_node) {
      // Usage interval of alloc doesn't intersect with current tensor's usage
      // interval, so we skip it.
      continue;
    }
    size_t aligned_current_offset = AlignTo(alignment, current_offset);
    // If we found a gap larger than required size, and smaller than previous
    // best fit, take it.
    if (aligned_current_offset + size <= alloc.offset &&
        alloc.offset - aligned_current_offset < best_offset_fit) {
      best_offset = aligned_current_offset;
      best_offset_fit = alloc.offset - current_offset;
    }
    current_offset = std::max(current_offset, alloc.offset + alloc.size);
  }
  if (best_offset == kOffsetNotAssigned) {
    best_offset = AlignTo(alignment, current_offset);
  }

  // Update the required buffer size.
  high_water_mark_ = std::max(high_water_mark_, best_offset + size);
  new_alloc->offset = best_offset;

  auto insertion_it = ordered_allocs_.begin();
  while (insertion_it != ordered_allocs_.end() && *insertion_it < *new_alloc) {
    ++insertion_it;
  }
  ordered_allocs_.insert(insertion_it, *new_alloc);
  return kTfLiteOk;
}

上面展示的是進(jìn)行內(nèi)存分配的核心代碼——Allocate()，通過上面的代碼我們知道，如果不需要保存中間結(jié)果，那么TFLite用于張量分配的空間大小為輸入張量所需空間之和 + 輸出張量所需空間之和 + 節(jié)點所用張量中最大的張量所用空間 * (1 + 最大分支數(shù))。假設(shè)所有張量大小都為64字節(jié)（為了計算方便），模型有兩個輸入張量、一個輸出張量、十個中間結(jié)果張量、不包含分支，如果不需要保留中間結(jié)果，那么所需空寂大小為（2 + 1 ）* 64 + 64 * 2 = 320 bytes，與之對比的是保留中間結(jié)果所需的內(nèi)存（2 + 1 + 10）* 64 = 832 bytes，可以看到節(jié)約的內(nèi)存還是相當(dāng)可觀的，減少了將近2/3的內(nèi)存使用。

那么，TFLite是怎么做到的呢？
還記得我們前面我們前面說過的alloc_node_和dealloc_node_嗎？他們的用處就在這體現(xiàn)出來了。Allocate()調(diào)用時傳遞的參數(shù)如下所示。

          arena_.Allocate(context_, tensor_alignment_, tensor.bytes,
                          tensor_index, alloc_node_[tensor_index],
                          dealloc_node_[tensor_index], &allocs_[tensor_index]));

結(jié)合上面的Allocate()的源碼，我們舉個栗子：

假設(shè)有一個模型，一共有6和節(jié)點編號0~6，用圓圈，以及五個用于保存中間結(jié)果Tensor，編號i~v，用方形表示。如圖4所示，假設(shè)這五個張量的大小關(guān)系為 i > v > ii > iv > iii。現(xiàn)在我們要為所有張量分配偏移地址。

圖4 模型結(jié)構(gòu)示意圖

從前面的描述我們知道，此時這些張量的索引值已經(jīng)有序的被保存在一個列表里，其順序為：輸入輸出張量的索引, i, v, ii, iv, iii，alloc_node_以及dealloc_node_列表的內(nèi)容如圖5所示，輸入輸出的值前面已經(jīng)提到過，分別是0 和 0xFFFFFFFF，就不再圖中畫出。

圖5 alloc_node_以及dealloc_node_

根據(jù)排序結(jié)果，輸入輸出的內(nèi)存分配完成后緊接著的就是分配張量 i。由于i與所有輸入輸出張量都有交集，因此只能排在他們后面。和自然的，i的偏移地址就從輸入輸出張量結(jié)束的地方開始。

圖6 張量i所占空間示意圖

緊接著，需要分配張量v的內(nèi)存，同理，它會排在所有輸入輸出后面，目前它和i的偏移地址是一樣的。接下來，需要拿v的偏移地址與所有在它之前已分配好了偏移地址并且與它有交集的張量比較以便進(jìn)行偏移地址的調(diào)整。很顯然，目前只有i分配完成，但是從alloc_node_以及dealloc_node_的值來看，i的范圍是1~2，而v的范圍是4~5，和它沒有交集，因此偏移地址不用調(diào)整，因此v的偏移地址會和i的一樣。

圖7 張量v所占空間示意圖

同樣的，接下來分配的是ii的偏移地址，由于ii與i、v有都有交集，因此需要根據(jù)他們調(diào)整偏移地址，由于i所需空間大于ii，因此在有著相同的起始地址情況下，肯定是i往后延伸跟多，因此ii直接回跟在i結(jié)束的地方。

圖8 張量ii所占空間示意圖

接下來，分配iv的偏移地址。iv只與v有交集，因此只根據(jù)v進(jìn)行調(diào)整，所iv緊隨v的后面。

圖9 張量iv所占空間示意圖

最后，分配iii。iii與ii、v、iv都有交集，因此一共會調(diào)整三次，最終會接在ii或者iv的結(jié)尾，就看誰結(jié)束的地方更靠后。我們假設(shè)最終會接在ii結(jié)束的地方。如圖，最終他們占據(jù)的內(nèi)存最多為i + ii + iii。v與iv復(fù)用同一片地址空間。

圖10 張量iii所占空間示意圖

我們可以推演下整個推過程：

1. 首先i會被賦值；
2. 接著，有i計算得到ii，此時i已經(jīng)沒用了;
3. 然后，不管我們先計算iii還是v，都不會覆蓋到ii的值，只會覆蓋i；
4. iii、v計算完成后，ii也幾經(jīng)沒用因此它的空間又可以為iv所用。

可以看到，整個過程不僅節(jié)省了空間，而且也能保證我們的數(shù)據(jù)都是安全的。如果需要保存中間結(jié)果，那么整個內(nèi)存排布將變成下面這樣：

圖11 保存中間結(jié)果張量所占空間示意圖

好了今天就到這了。歡1迎2關(guān)3注4個5人6微7信8公9眾10號【愛碼士1024】，此地有崇山峻嶺，茂林修竹；又有清流激……

等等！你還沒解釋為什么內(nèi)存對齊了還會存在空隙插入比較小的張量！

嗷對！為什么對齊了還會有空隙呢？還是前面的例子，極端點，在i特別大的情況下（Q:請問多大叫特別大？A: 鯤之大，一鍋燉不下那種……），當(dāng)我們分配iv的時候，在iv的結(jié)尾到iii的起始這段空間就會出現(xiàn)一個大間隙。如圖12所示，這不就可以趁虛而入了。這個故事告訴我們，在特別堵的路上跟車一定要跟的緊，否則肯定會有不講武德的人插隊……

圖12 插隊

3.3. 實際分配

這個內(nèi)存分配過程挺巧妙的。完成這一步，距離成功只有一步之遙，最后只要成功向系統(tǒng)申請一塊足夠大的內(nèi)存就行了。就相當(dāng)于我們常說的“等我有了錢，我要怎么怎么樣”，計劃已經(jīng)做好了，現(xiàn)在就差兩塊錢去買體彩。唯一的區(qū)別就是操作系統(tǒng)大概率會滿足TFLite的需求。

其實到了這一步非常簡單了，由于每個張量都有了偏移地址，因此只需要簡單的加上申請下來的基地址，就能得到所有真實地址。這里就不多介紹了。

接頭暗號：愛碼士1024

歡1迎2關(guān)3注4個5人6微7信8公9眾10號：愛碼士1024

4. Resources

[1] https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite