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

1. 前言

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

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

2. 太長；不看

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

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

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

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

3. 詳細(xì)過程

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

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

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

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

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

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

實(shí)際上递雀，TFLite申請(qǐng)的內(nèi)存分成兩塊柄延，一塊用于儲(chǔ)存臨時(shí)張量以及其他一些臨時(shí)性數(shù)據(jù)，另一塊則用于存儲(chǔ)永久性數(shù)據(jù)。兩塊內(nèi)存不同的是數(shù)據(jù)的生命周期不同搜吧，其他的完全一樣市俊。

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

3.1. 計(jì)劃階段

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

這三者中前兩個(gè)是ArenaPlanner的成員變量似舵，后續(xù)的操作中需要用到它們?cè)诖瞬襟E設(shè)置好的值。最后一個(gè)refcount是局部變量葱峡，主要用于確定dealloc_node_中的元素該如何賦值啄枕。

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

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

對(duì)于同一個(gè)張量（除了模型的輸蝇庭、輸出張量以及節(jié)點(diǎn)自身的臨時(shí)張量），它既是前一個(gè)張量的輸出捡硅，同時(shí)也是后一個(gè)（也可能是多個(gè)）節(jié)點(diǎn)的輸入哮内。因此，對(duì)于alloc_node_以及dealloc_node_中相同位置的元素而言壮韭，雖然它們對(duì)應(yīng)的是同一個(gè)張量北发，但是值往往不同纹因，其關(guān)系至少滿足alloc_node_[tensor_index] + 1 = dealloc_node_[tensor_index]，如圖2所示琳拨。請(qǐng)記住這一點(diǎn)瞭恰，因?yàn)楹罄m(xù)會(huì)用到這一關(guān)系。

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

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

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

3.2. 計(jì)算階段

每一個(gè)張量所擁有的buffer的信息，通過類ArenaAllocWithUsageInterval來保存合呐，ArenaAllocWithUsageInterval包含偏移（offset）暮的、大小、所屬的張量淌实、張量的輸入節(jié)點(diǎn)冻辩、輸出節(jié)點(diǎn)等信息，如下面代碼所示拆祈。

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;
  }
};

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

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

最終钧敞，通過計(jì)算，包含所有的張量的buffer信息ArenaAllocWithUsageInterval會(huì)按照一定的規(guī)則有序的存儲(chǔ)在一個(gè)名叫ordered_allocs_的列表中麸粮，另外溉苛，也會(huì)按照索引號(hào)從小到大的順序保存在allocs_列表中。ordered_allocs_以及allocs_存儲(chǔ)的都是ArenaAllocWithUsageInterval的指針弄诲。

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

1. 開始的時(shí)候，ordered_allocs_會(huì)被清空凤巨；

2. 為張量的索引值排序视乐，排序的原則為，對(duì)于任意兩個(gè)張量T1以及T2：
   i. 如果T1敢茁、T2兩個(gè)都屬于輸入張量或者輸出張量佑淀，那么這T1、T2張量維持他們的順序不變彰檬；
   ii. 如果T1屬于輸入張量或者輸出張量而T2不是伸刃，則T1排在T2之前；
   iii. 如果T1逢倍、T2都不輸入輸入張量或者輸出張量捧颅，則誰需要的內(nèi)存更多誰排前面；對(duì)于需要內(nèi)存一樣的兩個(gè)張量较雕，則誰先被用到誰排前面碉哑。
   排序的結(jié)果是一個(gè)由張量索引值組成的列表，后續(xù)將按照這個(gè)列表為張量分配內(nèi)存亮蒋。前面說到的ordered_allocs_里面元素的順序有大致這么確定扣典，不過可能會(huì)略微有區(qū)別（當(dāng)存在一個(gè)張量太小，可以插入其他一些張量的空隙慎玖，這點(diǎn)在后面有介紹）贮尖。

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

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

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é)點(diǎn)所用張量中最大的張量所用空間 * (1 + 最大分支數(shù))。假設(shè)所有張量大小都為64字節(jié)（為了計(jì)算方便）扔罪，模型有兩個(gè)輸入張量秉沼、一個(gè)輸出張量、十個(gè)中間結(jié)果張量、不包含分支唬复，如果不需要保留中間結(jié)果矗积，那么所需空寂大小為（2 + 1 ）* 64 + 64 * 2 = 320 bytes，與之對(duì)比的是保留中間結(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í)傳遞的參數(shù)如下所示茵烈。

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

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

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

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

從前面的描述我們知道篮迎，此時(shí)這些張量的索引值已經(jīng)有序的被保存在一個(gè)列表里男图，其順序?yàn)椋?code>輸入輸出張量的索引, 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)存，同理沈条，它會(huì)排在所有輸入輸出后面需忿，目前它和i的偏移地址是一樣的。接下來蜡歹，需要拿v的偏移地址與所有在它之前已分配好了偏移地址并且與它有交集的張量比較以便進(jìn)行偏移地址的調(diào)整屋厘。很顯然，目前只有i分配完成月而，但是從alloc_node_以及dealloc_node_的值來看汗洒，i的范圍是1~2，而v的范圍是4~5父款，和它沒有交集溢谤，因此偏移地址不用調(diào)整，因此v的偏移地址會(huì)和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都有交集，因此一共會(huì)調(diào)整三次摧玫，最終會(huì)接在ii或者iv的結(jié)尾耳奕，就看誰結(jié)束的地方更靠后绑青。我們假設(shè)最終會(huì)接在ii結(jié)束的地方。如圖屋群，最終他們占據(jù)的內(nèi)存最多為i + ii + iii闸婴。v與iv復(fù)用同一片地址空間。

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

我們可以推演下整個(gè)推過程：

1. 首先i會(huì)被賦值芍躏；
2. 接著邪乍，有i計(jì)算得到ii，此時(shí)i已經(jīng)沒用了;
3. 然后对竣，不管我們先計(jì)算iii還是v庇楞，都不會(huì)覆蓋到ii的值，只會(huì)覆蓋i否纬；
4. iii吕晌、v計(jì)算完成后，ii也幾經(jīng)沒用因此它的空間又可以為iv所用临燃。

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

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

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

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

嗷對(duì)蹬昌！為什么對(duì)齊了還會(huì)有空隙呢掩完？還是前面的例子，極端點(diǎn)蝌矛，在i特別大的情況下（Q:請(qǐng)問多大叫特別大菲驴？A: 鯤之大荐吵，一鍋燉不下那種……），當(dāng)我們分配iv的時(shí)候赊瞬，在iv的結(jié)尾到iii的起始這段空間就會(huì)出現(xiàn)一個(gè)大間隙先煎。如圖12所示，這不就可以趁虛而入了巧涧。這個(gè)故事告訴我們薯蝎，在特別堵的路上跟車一定要跟的緊，否則肯定會(huì)有不講武德的人插隊(duì)……

圖12 插隊(duì)

3.3. 實(shí)際分配

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

其實(shí)到了這一步非常簡單了，由于每個(gè)張量都有了偏移地址胎撤，因此只需要簡單的加上申請(qǐng)下來的基地址晓殊，就能得到所有真實(shí)地址。這里就不多介紹了哩照。

接頭暗號(hào)：愛碼士1024

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

4. Resources

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