PyTorch官方發(fā)布推薦系統(tǒng)庫(kù)：TorchRec
點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”
設(shè)為星標(biāo)，干貨直達(dá)！
我們很高興地宣布 TorchRec，一個(gè)用于推薦系統(tǒng)的 PyTorch 域庫(kù)。這個(gè)新庫(kù)提供了通用的稀疏性和并行性原語(yǔ)，使研究人員能夠構(gòu)建最先進(jìn)的個(gè)性化模型并將其部署到生產(chǎn)中。https://github.com/pytorch/torchrec
我們是怎么來(lái)到這里的?
推薦系統(tǒng) (RecSys) 包含當(dāng)今大量部署在生產(chǎn)環(huán)境中的 AI，但您可能不會(huì)通過(guò)查看 Github 了解它。與 Vision 和 NLP 等領(lǐng)域不同，RecSys 的大部分持續(xù)創(chuàng)新和開(kāi)發(fā)都是在封閉的公司門(mén)后進(jìn)行的。對(duì)于研究這些技術(shù)的學(xué)術(shù)研究人員或構(gòu)建個(gè)性化用戶(hù)體驗(yàn)的公司而言，該領(lǐng)域遠(yuǎn)未實(shí)現(xiàn)民主化。此外，RecSys 作為一個(gè)領(lǐng)域，很大程度上是由稀疏和/或順序事件上的學(xué)習(xí)模型定義的，這與 AI 的其他領(lǐng)域有很大的重疊。許多技術(shù)是可轉(zhuǎn)移的，特別是對(duì)于擴(kuò)展和分布式執(zhí)行。全球?qū)θ斯ぶ悄艿暮艽笠徊糠滞顿Y用于開(kāi)發(fā)這些 RecSys 技術(shù)，因此封鎖它們會(huì)阻止這項(xiàng)投資流入更廣泛的人工智能領(lǐng)域。
到 2020 年年中，PyTorch 團(tuán)隊(duì)收到了大量反饋，稱(chēng)開(kāi)源 PyTorch 生態(tài)系統(tǒng)中還沒(méi)有大規(guī)模的生產(chǎn)質(zhì)量推薦系統(tǒng)包。當(dāng)我們?cè)噲D找到一個(gè)好的答案時(shí)，Meta 的一組工程師希望將 Meta 的生產(chǎn) RecSys 堆棧作為 PyTorch 域庫(kù)貢獻(xiàn)出來(lái)，并堅(jiān)定地致力于圍繞它發(fā)展一個(gè)生態(tài)系統(tǒng)。這似乎是一個(gè)讓 RecSys 領(lǐng)域的研究人員和公司受益的好主意。因此，從 Meta 的堆棧開(kāi)始，我們開(kāi)始模塊化和設(shè)計(jì)一個(gè)完全可擴(kuò)展的代碼庫(kù)，該代碼庫(kù)適用于各種推薦用例。我們的目標(biāo)是從 Meta 的軟件堆棧中提取關(guān)鍵構(gòu)建塊，以同時(shí)實(shí)現(xiàn)創(chuàng)造性探索和擴(kuò)展。經(jīng)過(guò)近兩年的跨 Meta 基準(zhǔn)測(cè)試、遷移和測(cè)試，我們很高興終于與 RecSys 社區(qū)一起踏上了這一旅程。我們希望這個(gè)軟件包在 RecSys 行業(yè)開(kāi)啟對(duì)話(huà)和協(xié)作，從 Meta 作為第一個(gè)相當(dāng)大的貢獻(xiàn)者開(kāi)始。
TorchRec介紹
TorchRec 包括一個(gè)可擴(kuò)展的低級(jí)建?；A(chǔ)以及豐富的電池模塊。我們最初的目標(biāo)是“雙塔”（[1]，[2]）架構(gòu)，這些架構(gòu)具有單獨(dú)的子模塊來(lái)學(xué)習(xí)候選項(xiàng)目和查詢(xún)或上下文的表示。輸入信號(hào)可以是浮點(diǎn)“密集”特征或需要訓(xùn)練大型嵌入表的高基數(shù)分類(lèi)“稀疏”特征的混合。此類(lèi)架構(gòu)的有效訓(xùn)練涉及將復(fù)制計(jì)算的“密集”部分的數(shù)據(jù)并行性和將大型嵌入表分區(qū)到多個(gè)節(jié)點(diǎn)的模型并行性相結(jié)合。
特別地，這個(gè)庫(kù)包含：
建模原語(yǔ)，例如嵌入包和鋸齒張量，可以使用混合數(shù)據(jù)并行和模型并行輕松創(chuàng)作大型、高性能的多設(shè)備/多節(jié)點(diǎn)模型。
由 FBGEMM 提供支持的優(yōu)化 RecSys 內(nèi)核，包括對(duì)稀疏和量化操作的支持。
一個(gè)分片器，它可以使用各種不同的策略對(duì)嵌入表進(jìn)行分區(qū)，包括數(shù)據(jù)并行、逐表、逐行、逐表行和逐列分片。
可以自動(dòng)為模型生成優(yōu)化的分片計(jì)劃的規(guī)劃器。
流水線(xiàn)以重疊數(shù)據(jù)加載設(shè)備傳輸（復(fù)制到 GPU）、設(shè)備間通信（input_dist）和計(jì)算（前向、后向）以提高性能。
GPU 推理支持。
RecSys 的通用模塊，例如模型和公共數(shù)據(jù)集（Criteo 和 Movielens）
為了展示此工具的靈活性，讓我們看一下從我們的 DLRM 事件預(yù)測(cè)示例中提取的以下代碼片段：
#?Specify?the?sparse?embedding?layers
eb_configs?=?[
???EmbeddingBagConfig(
???????name=f"t_{feature_name}",
???????embedding_dim=64,
???????num_embeddings=100_000,
???????feature_names=[feature_name],
???)
???for?feature_idx,?feature_name?in?enumerate(DEFAULT_CAT_NAMES)
]
#?Import?and?instantiate?the?model?with?the?embedding?configuration
#?The?"meta"?device?indicates?lazy?instantiation,?with?no?memory?allocated
train_model?=?DLRM(
???embedding_bag_collection=EmbeddingBagCollection(
???????tables=eb_configs,?device=torch.device("meta")
???),
???dense_in_features=len(DEFAULT_INT_NAMES),
???dense_arch_layer_sizes=[512,?256,?64],
???over_arch_layer_sizes=[512,?512,?256,?1],
???dense_device=device,
)
#?Distribute?the?model?over?many?devices,?just?as?one?would?with?DDP.
model?=?DistributedModelParallel(
???module=train_model,
???device=device,
)
optimizer?=?torch.optim.SGD(params,?lr=args.learning_rate)
#?Optimize?the?model?in?a?standard?loop?just?as?you?would?any?other?model!
#?Or,?you?can?use?the?pipeliner?to?synchronize?communication?and?compute
for?epoch?in?range(epochs):
???#?Train擴(kuò)展性能
TorchRec 擁有最先進(jìn)的大規(guī)模推薦 AI 基礎(chǔ)設(shè)施，為 Meta 的一些最大模型提供支持。它用于訓(xùn)練一個(gè) 1.25 萬(wàn)億參數(shù)的模型，于 1 月投入生產(chǎn)，還有一個(gè)即將投入生產(chǎn)的 3 萬(wàn)億參數(shù)模型。這應(yīng)該很好地表明 PyTorch 完全有能力解決業(yè)界最大規(guī)模的 RecSys 問(wèn)題。我們從社區(qū)中的許多人那里聽(tīng)說(shuō)，分片嵌入是一個(gè)痛點(diǎn)。TorchRec 清楚地解決了這個(gè)問(wèn)題。不幸的是，使用公共數(shù)據(jù)集提供大規(guī)?；鶞?zhǔn)測(cè)試具有挑戰(zhàn)性，因?yàn)榇蠖鄶?shù)開(kāi)源基準(zhǔn)測(cè)試太小而無(wú)法大規(guī)模顯示性能。
展望
開(kāi)源和開(kāi)放技術(shù)具有普遍的好處。Meta 正在為 PyTorch 社區(qū)播種最先進(jìn)的 RecSys 軟件包，希望許多人加入到構(gòu)建它的過(guò)程中，支持新的研究并幫助許多公司。TorchRec 背后的團(tuán)隊(duì)計(jì)劃無(wú)限期地繼續(xù)這個(gè)計(jì)劃，建立 TorchRec 以滿(mǎn)足 RecSys 社區(qū)的需求，歡迎新的貢獻(xiàn)者，并繼續(xù)推動(dòng) Meta 的個(gè)性化。我們很高興開(kāi)始這一旅程，并期待貢獻(xiàn)、想法和反饋！
參考
[1] Sampling-Bias-Corrected Neural Modeling for Large Corpus Item Recommendations
[2] DLRM: An advanced, open source deep learning recommendation model
本文翻譯自PyTorch blog：https://pytorch.org/blog/introducing-torchrec/
推薦閱讀
輔助模塊加速收斂，精度大幅提升！移動(dòng)端實(shí)時(shí)的NanoDet-Plus來(lái)了！
SSD的torchvision版本實(shí)現(xiàn)詳解
機(jī)器學(xué)習(xí)算法工程師
? ??? ? ? ? ? ? ? ? ? ? ????????? ??一個(gè)用心的公眾號(hào)
