OneFlow源碼解析之Eager模式下Tensor存儲管理
1不同Tensor類型的存儲管理方式
Lazy Tensor 的存儲是由 Runtime 和 Actor 等對象管理的。靜態(tài)圖完成編譯后,需要多少個對象、多少存儲空間都是確定的,Runtime 等在初始化時會分配存儲,在退出時回收資源。
Eager 模式下,Global Tensor 可以視為對 Local Tensor 的分布式封裝,EagerGlobalTensorImpl 在本地的數(shù)據(jù)是一個EagerLocalTensorImpl 對象??梢酝ㄟ^考察 EagerLocalTensorImpl 來理解 eager 模式下 tensor 的存儲管理。
參考的示例代碼如下:
import numpy as np import oneflow as flow a = np.random.randn(1, 4) flow.tensor(a, device=flow.device("cpu"), dtype=flow.float)
2Tensor 存儲相關類的關系
EagerLocalTensorImpl 的存儲相關的類關系如下。
后續(xù)會順著示例代碼的執(zhí)行過程,看看圖中的對象都是在何時、如何構造的,存儲被誰持有、如何分配并釋放。
3通過虛擬機指令為 Tensor 分配存儲
tensor 的構造函數(shù)通過 Python C API 注冊為 PyTensorObject_init,由
functional::_legacy_tensor_ctor根據(jù)簽名進行轉發(fā)。
示例代碼對應的是 TensorWithDataFunctor,調(diào)用 MakeLocalTensorFromData 構造 tensor,在這個函數(shù)中通過調(diào)用 functional::Empty以及 EmptyFunctor分配存儲。在 EmptyFunctor 中把相關屬性都存到 attrs,然后調(diào)用 OpInterpUtil::Dispatch在 vm 指令的執(zhí)行準備過程中分配存儲。
EmptyFunctor 返回的 tensor 是一個只有存儲空間、不含數(shù)據(jù)的對象。數(shù)據(jù)拷貝在后面由
CopyLocalTensorFromUntypedArray完成。
3.1 存儲相關對象的構造
因為是 eager 模式下的 local tensor,OpInterpUtil::Dispatch 會被轉發(fā)到 NaiveInterpret執(zhí)行。對于示例代碼,這個函數(shù)的輸入?yún)?shù)如下:
- inputs 是一個空數(shù)組
- outputs 只有一個元素、且是空指針
因為 outputs 中的 tensor 指針都是空的,所以需要創(chuàng)建一個 EagerLocalTensorImpl 對象,其 one::TensorStorage 成員變量是空指針。
因為 output_eager_blob_objects 中的元素尚未初始化,會調(diào)用 tensor_impl->InitEagerBlobObject進行初始化。因為 tensor_storage_ 還是空的,這個過程會執(zhí)行如下操作:
- 創(chuàng)建 vm::TensorStorage 對象
- 創(chuàng)建 EagerBlobObject 對象
- set_eager_blob_object
- UpdateTensorStorage
- 創(chuàng)建 one::TensorStorage 對象
- 設置 tensor 存儲釋放的回調(diào)函數(shù)
上述對象的創(chuàng)建,都只是記錄相關信息,還不涉及 tensor 的存儲分配。
需要注意的是,注冊到 one::TensorStorage 的回調(diào)函數(shù)被賦值給了成員變量 releaser_hook_,這個函數(shù)會通過虛擬機指令釋放 tensor。
3.2 在指令執(zhí)行過程中分配 tensor 存儲
分配 tensor 存儲的過程如下:
- vm::Instruction::Compute
- vm::InstructionPolicy::ComputeIf
- vm::OpCallInstructionPolicy::Compute
- OpCallInstructionUtil::Compute
- 獲取內(nèi)存分配器
- OpCallInstructionUtil::AllocateOutputBlobsMemory
- blob_object->TryAllocateBlobBodyMemory
- allocator->Allocate
在EagerBlobObject::TryAllocateBlobBodyMemory 中,allocator 分配的存儲地址會賦值給 dptr,存儲地址 dptr 和 Free 函數(shù)一起構造一個智能指針,并賦值給 vm::TensorStorage 的 blob_dptr_ 變量。
4通過虛擬機指令釋放 Tensor 存儲
在前面的 3.1 節(jié)提到,EagerLocalTensorImpl 在初始化 EagerBlobObject、創(chuàng)建 one::TensorStorage 的同時,會設置一個釋放 tensor 的回調(diào)函數(shù),回調(diào)函數(shù)保存在變量 releaser_hook_ 中,one::TensorStorage 析構時調(diào)用這個回調(diào)函數(shù)。把這些信息綜合整理一下,one::TensorStorage 析構時會執(zhí)行如下操作:
vm::InstructionList instruction_list; InstructionsBuilder instructions_builder(&instruction_list); // JUST(Build(&instructions_builder)); if (eager_blob_object->producer_stream().has_value()) { JUST(instructions_builder->ReleaseTensor(eager_blob_object)); } JUST(vm::Run(instructions_builder.mut_instruction_list()));
在InstructionsBuilder::ReleaseTensor 中,如果有其它 stream 最近使用了 eager_blob_object,會通過 SoftSyncStreamBetween 進行同步。通過這種方式解決存儲的依賴問題。
一般情況下,通過 tensor 的 producer_stream 釋放存儲,根據(jù)這個對象獲取對應的 vm::Stream 對象,并據(jù)此構造指令 instruction(包含 eager_blob_object 和 vm_stream),示例代碼對應的指令類型是 FastReleaseTensorInstructionPolicy,其 Compute 方法執(zhí)行具體的存儲釋放邏輯,過程如下:
- ReleaseTensorInstructionPolicy::Release()
- eager_blob_object->DeallocateBlobDataPtr()
- tensor_storage_->Release()
- tensor_storage_->_Release()
- blob_dptr_.reset()
智能指針重置,調(diào)用分配存儲時指定的 Free 方法
5reshape 等場景的存儲管理
在 reshape、slice、transpose 等場景中,調(diào)用的 EagerLocalTensorImpl 構造函數(shù)的參數(shù)包括 input 的 tensor_storage,所以這個 tensor 的 tensor_storage_ 變量不是空的,在執(zhí)行 InitEagerBlobObject 時,只創(chuàng)建 EagerBlobObject以提供 shape、stride等信息;但不會再創(chuàng)建 one::TensorStorage,而是復用 input 的存儲。
6兩個 TensorStorage 類型可以合并嗎?
為什么在 one::TensorStorage 析構時、由它保存的回調(diào)函數(shù)來觸發(fā)釋放 vm::TensorStorage 中的存儲呢?
one::TensorStorage 只多了一個 releaser,這兩個 Storage 類型是否可以合并呢?
在當前的設計下,這兩個類型不能合并。因為
one::TensorStorage::releaser_hook_ 中持有 EagerBlobObject 的智能指針,EagerBlobObject 中也持有 vm::TensorStorage 的智能指針。如果兩個 Storage 類型合并為一個,就會出現(xiàn)循環(huán)引用、對象無法析構而導致內(nèi)存泄漏。
所以,vm::TensorStorage 只是單純的存儲,可以在多個 tensor 之間共享。EagerBlobObject 既包括存儲、也包括 shape、stride、data_type 等獨特的對象信息。而 one::TensorStorage 是為了避免循環(huán)引用而引入的、專門負責釋放存儲的角色。
7附錄
GDB 斷點示例
break oneflow::one::MakeLocalTensorFromData break oneflow::one::NaiveInterpret break oneflow::vm::VirtualMachineEngine::DispatchInstruction break oneflow::vm::OpCallInstructionUtil::Compute break oneflow::vm::OpCallInstructionUtil::AllocateOutputBlobsMemory break oneflow::vm::EagerBlobObject::TryAllocateBlobBodyMemory break oneflow::vm::ReleaseTensorInstructionPolicy::Release break oneflow/core/eager/eager_blob_object.cpp:107
參考資料
OneFlow(github.com/Oneflow-Inc… )
OneFlow源碼解析:Tensor類型體系與Local Tensor
歡迎 Star、試用 OneFlow 最新版本:
以上就是OneFlow源碼解析:Eager模式下Tensor的存儲管理的詳細內(nèi)容,更多關于OneFlow源碼解析:Eager模式下Tensor的存儲管理的資料請關注腳本之家其它相關文章!
相關文章
pytorch cnn 識別手寫的字實現(xiàn)自建圖片數(shù)據(jù)
這篇文章主要介紹了pytorch cnn 識別手寫的字實現(xiàn)自建圖片數(shù)據(jù),小編覺得挺不錯的,現(xiàn)在分享給大家,也給大家做個參考。一起跟隨小編過來看看吧2018-05-05