摘要：訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)需要大量的內(nèi)存，用戶使用這個(gè)工具包，可以在計(jì)算時(shí)間成本僅增加的基礎(chǔ)上，在上運(yùn)行規(guī)模大倍的前饋模型。使用導(dǎo)入此功能，與使用方法相同，使用梯度函數(shù)來計(jì)算參數(shù)的損失梯度。隨后，在反向傳播中重新計(jì)算檢查點(diǎn)之間的節(jié)點(diǎn)。

OpenAI是電動(dòng)汽車制造商特斯拉創(chuàng)始人 Elon Musk和著名的科技孵化器公司 Y Combinator總裁 Sam Altman于 2016年聯(lián)合創(chuàng)立的 AI公司，使命是建立安全的人工智能，確保 AGI（Artificial general intelligence，通用人工智能）的全球效益“盡可能廣泛和均勻分配”。

目前，這個(gè)非營利組織已經(jīng)有 60名全職研究人員和工程人員，“摒棄隨之產(chǎn)生的利己機(jī)會(huì)，為實(shí)現(xiàn)自己的使命而奮斗”。

Y Combinator和一家電動(dòng)車制造商能夠走到一起創(chuàng)辦人工智能公司 OpenAI，是因?yàn)閮杉夜竟餐哪繕?biāo)：創(chuàng)造一種新的人工智能實(shí)驗(yàn)室，這個(gè)實(shí)驗(yàn)室不受包括 Google在內(nèi)的任何人的控制，讓人類以安全的方式開發(fā)真正的人工智能。人工智能威脅論的 Musk建立 OpenAI，除了他所說讓人類以安全的方式開發(fā)真正的人工智能這一目標(biāo)外，也是出于其公司未來業(yè)務(wù)的需求，特斯拉、SpaceX都需要人工智能。

他同時(shí)也是十分積極的人工智能威脅論支持者，多次直言不諱對(duì)于人工智能失去人類控制導(dǎo)致災(zāi)難性后果的擔(dān)憂。

研究成果：使用梯度檢查點(diǎn)節(jié)約內(nèi)存

OpenAI自創(chuàng)立以來開發(fā)了分層強(qiáng)化學(xué)習(xí)算法、機(jī)器人模擬訓(xùn)練閉環(huán)系統(tǒng)等項(xiàng)目，近日研究員 Tim Salimans和 Yaroslav Bulatov又推出了一個(gè)工具包，可以讓內(nèi)存計(jì)算和模型更加適配。

訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)需要大量的內(nèi)存，用戶使用這個(gè)工具包，可以在計(jì)算時(shí)間成本僅增加 20％的基礎(chǔ)上，在 GPU上運(yùn)行規(guī)模大 10倍的前饋模型。

同樣地，計(jì)算反向傳播的梯度損失會(huì)在訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)的過程中占用大量內(nèi)存。通過我們這個(gè)模型定義的計(jì)算圖檢查點(diǎn)節(jié)點(diǎn)，并在反向傳播期間重新計(jì)算這些節(jié)點(diǎn)之間的圖形部分，可以降低計(jì)算梯度的內(nèi)存成本。當(dāng)訓(xùn)練由 n層組成的深度前饋神經(jīng)網(wǎng)絡(luò)時(shí)，我們可以以這種方式將存儲(chǔ)器消耗降低到 O（sqrt（n）），代價(jià)是執(zhí)行一次額外的正向傳播（參見例子：Training Deep Nets with Sublinear Memory Cost, by Chen et al（2016）。這個(gè)存儲(chǔ)庫在 Tensorflow中提供了這個(gè)功能，使用 Tensorflow圖形編輯器自動(dòng)重寫反向傳遞的計(jì)算圖。

?

使用常規(guī) tf.gradients函數(shù)和我們的內(nèi)存優(yōu)化梯度方法訓(xùn)練同樣批量大小的 ResNet模型時(shí)所使用的內(nèi)存

如何運(yùn)行？

對(duì)于具有 n層的簡單前饋神經(jīng)網(wǎng)絡(luò)，用于獲得梯度的計(jì)算圖如下所示：

標(biāo)記了 f的節(jié)點(diǎn)激活神經(jīng)網(wǎng)絡(luò)層。在正向傳遞期間，所有這些節(jié)點(diǎn)按順序一一被評(píng)估。相對(duì)于這些層的激活和參數(shù)的損失梯度由用 b標(biāo)記的節(jié)點(diǎn)表示。在反向傳遞期間，所有這些節(jié)點(diǎn)都按照相反的順序進(jìn)行評(píng)估。計(jì)算 b節(jié)點(diǎn)需要 f節(jié)點(diǎn)的計(jì)算結(jié)果，因此所有 f節(jié)點(diǎn)在正向傳遞之后都被保存在存儲(chǔ)器中。只有當(dāng)反向傳播進(jìn)行到足以計(jì)算出一個(gè) f節(jié)點(diǎn)的所有依賴關(guān)系或子節(jié)點(diǎn)時(shí)，它才能從內(nèi)存中消失。這意味著簡單反向傳播所需的內(nèi)存隨神經(jīng)網(wǎng)絡(luò)層數(shù) n線性增長。下面我們顯示這些節(jié)點(diǎn)的計(jì)算順序。紫色陰影圓圈表示在給定時(shí)間需要保存在內(nèi)存中的節(jié)點(diǎn)。

?

Vanilla backprop

如上所述的簡單反向傳播的計(jì)算效果較佳：它僅對(duì)每個(gè)節(jié)點(diǎn)進(jìn)行一次計(jì)算。但是，如果我們?cè)敢庵匦掠?jì)算節(jié)點(diǎn)，則可以節(jié)省大量內(nèi)存。例如，當(dāng)需要時(shí)我們可以簡單地重新計(jì)算每個(gè)節(jié)點(diǎn)的正向傳遞。執(zhí)行順序和使用的內(nèi)存如下所示：

使用這種策略，圖中計(jì)算神經(jīng)網(wǎng)絡(luò)層數(shù)為 n的梯度所需的內(nèi)存不變，這在內(nèi)存方面是最優(yōu)的。但是，請(qǐng)注意，節(jié)點(diǎn)評(píng)估的數(shù)量現(xiàn)在擴(kuò)大到 n ? ^2，而之前為 n：n個(gè)節(jié)點(diǎn)中的每一個(gè)都按 n次的順序重新計(jì)算。因此，對(duì)于深度網(wǎng)絡(luò)來說，圖形計(jì)算變得慢得多，這使得該方法在深度學(xué)習(xí)中不切實(shí)際。

為了在內(nèi)存和計(jì)算之間取得平衡，我們需要提出一個(gè)允許節(jié)點(diǎn)重新計(jì)算但不至于太頻繁的策略。我們?cè)谶@里使用的策略是將神經(jīng)網(wǎng)絡(luò)激活的一個(gè)子集標(biāo)記為檢查點(diǎn)節(jié)點(diǎn)。

?

我們選擇的檢查點(diǎn)節(jié)點(diǎn)

這些檢查點(diǎn)節(jié)點(diǎn)在正向傳遞后保留在內(nèi)存中，而其余節(jié)點(diǎn)至多重新計(jì)算一次。重新計(jì)算后，非檢查點(diǎn)節(jié)點(diǎn)將保存在內(nèi)存中，直到不再需要它們。對(duì)于簡單前饋神經(jīng)網(wǎng)絡(luò)來說，所有的神經(jīng)元激活節(jié)點(diǎn)都是由正向傳播定義的圖形分隔符或關(guān)節(jié)點(diǎn)。這意味著我們只需要在 backprop期間計(jì)算 b節(jié)點(diǎn)時(shí)，重新計(jì)算 b節(jié)點(diǎn)和它之前的最后一個(gè)檢查點(diǎn)之間的節(jié)點(diǎn)。當(dāng) backprop計(jì)算進(jìn)展到檢查點(diǎn)節(jié)點(diǎn)時(shí)，所有從它重新計(jì)算的節(jié)點(diǎn)都可以從內(nèi)存中刪除。計(jì)算和內(nèi)存使用的結(jié)果順序如下所示：

?

圖 3.檢查節(jié)點(diǎn) backprop

對(duì)于本例中的簡單前饋網(wǎng)絡(luò)來說，較佳選擇是將每個(gè) sqrt（n）節(jié)點(diǎn)標(biāo)記為檢查點(diǎn)。這樣，檢查點(diǎn)節(jié)點(diǎn)的數(shù)量和檢查點(diǎn)之間的節(jié)點(diǎn)數(shù)目都是 sqrt（n），這意味著所需的內(nèi)存現(xiàn)在也與我們網(wǎng)絡(luò)中層數(shù)的平方根成比例。由于每個(gè)節(jié)點(diǎn)最多只能重新計(jì)算一次，因此該策略所需的額外計(jì)算量僅相當(dāng)于網(wǎng)絡(luò)的單個(gè)正向傳遞。

我們的軟件包可以實(shí)現(xiàn)檢查點(diǎn) backprop，如上圖 3所示。這個(gè)過程是通過采用將上圖（圖 1）標(biāo)準(zhǔn) backprop，并使用 Tensorflow圖形編輯器對(duì)其自動(dòng)重寫來實(shí)現(xiàn)的。對(duì)于包含關(guān)節(jié)點(diǎn)（單節(jié)點(diǎn)圖分隔符）的圖，我們使用 sqrt（n）策略自動(dòng)選擇檢查點(diǎn)，為前饋網(wǎng)絡(luò)提供 sqrt（n）所需內(nèi)存。對(duì)于只包含多節(jié)點(diǎn)圖分隔符的普通圖像，使用我們的檢查點(diǎn) backprop仍然有效，但是目前還需要用戶手動(dòng)選擇檢查點(diǎn)。

在我們軟件包的 博客文章中可以找到計(jì)算圖、內(nèi)存使用情況和梯度計(jì)算策略的附加說明文件。

?

安裝要求

pip install tf-nightly-gpu

pip install toposort networkx pytest

另外，在運(yùn)行測試時(shí)，要確保可以找到 CUDA分析工具接口（CUPTI），例如，通過運(yùn)行

```pip install toposort networkx pytestexport LD_LIBRARY_PATH =“$ {LD_LIBRARY_PATH}：/ usr / local / cuda / extras / CUPTI / lib64”

?

使用方法

這個(gè)存儲(chǔ)庫提供了基于 Tensorflow的 tf.gradient的替代方案。使用from memory_saving_gradients import gradients導(dǎo)入此功能，與使用 tf.gradients方法相同，使用梯度函數(shù)來計(jì)算參數(shù)的損失梯度。 （假設(shè)你明確地調(diào)用 tf.gradients，而不是在 tf.train.Optimizer中隱式調(diào)用）。

除了 tf.gradients的常規(guī)參數(shù)之外，我們的梯度函數(shù)還有一個(gè)額外的參數(shù)：檢查點(diǎn)。檢查點(diǎn)參數(shù)將通過計(jì)算圖告訴您向前傳遞過程中所需檢查點(diǎn)的梯度函數(shù)。隨后，在反向傳播中重新計(jì)算檢查點(diǎn)之間的節(jié)點(diǎn)。您可以使用很多 tensor、梯度（ys，xs，檢查點(diǎn) = [tensor1，tensor2]），或者可以使用以下幾個(gè)關(guān)鍵字之一：

"collection"（默認(rèn)值）：檢查由 tf.get_collection（"checkpoints"）返回的所有張量。之后在定義模型時(shí)，需要確保使用 tf.add_to_collection（“檢查點(diǎn)”，張量）將張量添加到此集合中。

"內(nèi)存"：使用啟發(fā)式自動(dòng)選擇一組節(jié)點(diǎn)作為檢查點(diǎn)，實(shí)現(xiàn)我們所需的 O（sqrt（n））內(nèi)存使用。啟發(fā)式方法通過自動(dòng)識(shí)別圖中的關(guān)節(jié)點(diǎn)，即在移除時(shí)將圖分成兩個(gè)部分的張量，然后找出這些張量的合理數(shù)目。目前，它適用于很多（但不是全部）模型。

"速度"：這個(gè)選項(xiàng)通過檢查所有操作的輸出來使運(yùn)行速度較大化，且通常成本高昂，即卷積和矩陣乘法。

?

覆蓋tf.gradients

直接使用新梯度函數(shù)的一個(gè)有效果的替代方法，是覆寫 Python已經(jīng)注冊(cè)到 tf.gradients的函數(shù)，如下所示：

import tensorflow as tf

import memory_saving_gradients

# monkey patch tf.gradients to point to our custom version, with automatic checkpoint selection

def gradients_memory(ys, xs, grad_ys=None, **kwargs):

? return memory_saving_gradients.gradients(ys, xs, grad_ys, checkpoints="memory", **kwargs)

tf.__dict__["gradients"] = gradients_memory

之后，使用節(jié)約內(nèi)存的版本對(duì) tf.gradients的調(diào)用都將使用內(nèi)存保存版本。

?

測試

測試文件夾包含用于測試代碼的正確性，并分析各種模型內(nèi)存使用情況的腳本。修改代碼后，您可以從該文件夾運(yùn)行./run_all_tests.sh來執(zhí)行測試。

?

在 CIFAR10上測試不同層數(shù)的 ResNet的內(nèi)存使用情況和運(yùn)行時(shí)間。大小為 1280，GTX1080

限制

在運(yùn)行處理大型圖形過慢的模型之前，這些代碼用 Python執(zhí)行所有的圖形操作。當(dāng)前用于自動(dòng)選擇檢查點(diǎn)的算法純粹是啟發(fā)式的，預(yù)計(jì)在我們測試之外的某些模型上會(huì)失敗。在這種情況下，應(yīng)使用手動(dòng)模式選擇檢查點(diǎn)。

?

參考

Academic papers describing checkpointed backpropagation: Training Deep Nets with Sublinear Memory Cost, by Chen et al. (2016), https://arxiv.org/pdf/1604.06174.pdfMemory-Efficient Backpropagation Through Time, by Gruslys et al. (2016) https://arxiv.org/pdf/1604.06174.pdf

Explanation of using graph_editor to implement checkpointing on TensorFlow graphs: https://github.com/tensorflow/tensorflow/issues/4359#issuecomment-269241038, https://github.com/yaroslavvb/stuff/blob/master/simple_rewiring.ipynb

Experiment code/details: https://medium.com/@yaroslavvb/testing-memory-saving-on-v100-8aa716bbdf00

TensorFlow memory tracking package:

https://github.com/yaroslavvb/chain_constant_memory/blob/master/mem_util_test.py

Implementation of "memory-poor" backprop strategy in TensorFlow for a simple feed-forward net:

https://github.com/yaroslavvb/chain_constant_memory/

?

再度招兵買馬

在過去的幾個(gè)月，OpenAI一直保持沉默。而近期，這家公司似乎渴望把更多的人才招攬到團(tuán)隊(duì)中去，發(fā)布了多個(gè)招聘職位。

從公司發(fā)布的工作崗位上看，團(tuán)隊(duì)不僅需要招聘協(xié)調(diào)員，還要招聘更多的人員。

https://twitter.com/OpenAI/status/951498942541283328?ref_src=twsrc%5Etfw&ref_url=https%3A%2F%2Finterestingengineering.com%2Felon-musk-and-openai-want-to-create-an-artificial-intelligence-that-wont-spell-doom-for-humanity

目前，OpenAI官網(wǎng)招聘頁面放出的招聘崗位（以舊金山為例）主要包括技術(shù)人員（基礎(chǔ)架構(gòu)工程師，機(jī)器學(xué)習(xí)方面工程師、研究院、工作人員、實(shí)習(xí)生）、運(yùn)營人員（業(yè)務(wù)經(jīng)理、招聘協(xié)調(diào)員）。

招聘頁面：https://jobs.lever.co/openai

當(dāng)初，OpenAI花了大價(jià)錢邀請(qǐng)一眾大佬加盟，其中就包括曾在谷歌和 Facebook實(shí)習(xí)過的研究員 Wojciech Zaremba，說 OpenAI開出的價(jià)錢“近乎瘋狂”。那么，瘋狂的 offer價(jià)錢在 2016年時(shí)到底是什么水平呢？據(jù)微軟研究院副總裁 Peter Lee稱，當(dāng)時(shí)一名較高級(jí)研究員的 offer遠(yuǎn)超過美國全國橄欖球聯(lián)賽的較高級(jí)四分衛(wèi)的酬勞，這還是在沒有硅谷跟你搶人才的情況下。Zaremba說道，OpenAI給他開出的價(jià)格是市場平均水平的兩至三倍。

聯(lián)想到最近 OpenAI人才流失比較嚴(yán)重，比如加州伯克利大學(xué)教授、機(jī)器人學(xué)習(xí)大牛 Pieter Abbeel離開 OpenAI創(chuàng)立智能機(jī)器人公司 Embodied Intelligence，想必這次招聘他們出的價(jià)錢也不會(huì)低。

原文鏈接：

https://github.com/openai/gradient-checkpointing

https://openai.com/

https://interestingengineering.com/elon-musk-and-openai-want-to-create-an-artificial-intelligence-that-wont-spell-doom-for-humanity

商業(yè)智能與數(shù)據(jù)分析群

興趣范圍包括各種讓數(shù)據(jù)產(chǎn)生價(jià)值的辦法，實(shí)際應(yīng)用案例分享與討論，分析工具，ETL工具，數(shù)據(jù)倉庫，數(shù)據(jù)挖掘工具，報(bào)表系統(tǒng)等全方位知識(shí)

QQ群：81035754