C/C++中 Python 擴充套件物件的簡介
你可能知道可以藉助 C/C++擴充套件 Python,並開發所謂的「擴充套件」。PyTorch 的所有繁重工作由 C/C++實現,而不是純 Python。為了定義 C/C++中一個新的 Python 物件型別,你需要定義如下例項的一個類似結構:
// Python object that backs torch.autograd.Variable
structTHPVariable{
PyObject_HEAD
torch::autograd::Variable cdata;
PyObject* backward_hooks;
};
如上,在定義的開始有一個稱之為 PyObject_HEAD 的巨集,其目標是標準化 Python 物件,並擴充套件至另一個結構,該結構包含一個指向型別物件的指標,以及一個帶有引用計數的欄位。
Python API 中有兩個額外的巨集,分別稱為 Py_INCREF() 和 Py_DECREF(),可用於增加和減少 Python 物件的引用計數。多實體可以借用或擁有其他物件的引用(因此引用計數被增加),而只有當引用計數達到零,Python 才會自動刪除那個物件的記憶體。想了解更多有關 Python C/++擴充套件的知識,請參見:https://docs.python.org/3/extending/newtypes.html。
有趣的事實:使用小的整數作為索引、計數等在很多應用中非常見。為了提高效率,官方 CPython 直譯器快取從-5 到 256 的整數。正由於此,宣告 a = 200; b = 200; a is b 為真,而宣告 a = 300; b = 300; a is b 為假。
Zero-copy PyTorch 張量到 Numpy,反之亦然
PyTorch 有專屬的張量表徵,分離 PyTorch 的內部表徵和外部表徵。但是,由於 Numpy 陣列的使用非常普遍,尤其是當資料載入源不同時,我們確實需要在 Numpy 和 PyTorch 張量之間做轉換。正由於此,PyTorch 給出了兩個方法(from_numpy() 和 numpy()),從而把 Numpy 陣列轉化為 PyTorch 陣列,反之亦然。如果我們檢視把 Numpy 陣列轉化為 PyTorch 張量的呼叫程式碼,就可以獲得有關 PyTorch 內部表徵的更多洞見:
at:: Tensor tensor_from_numpy(PyObject* obj){
if(!PyArray_Check(obj)) {
throwTypeError( "expected np.ndarray (got %s)", Py_TYPE(obj)->tp_name);
}
autoarray= (PyArrayObject*)obj;
intndim = PyArray_NDIM( array);
autosizes = to_aten_shape(ndim, PyArray_DIMS( array));
autostrides = to_aten_shape(ndim, PyArray_STRIDES( array));
// NumPy strides use bytes. Torch strides use element counts.
autoelement_size_in_bytes = PyArray_ITEMSIZE( array);
for( auto& stride : strides) {
stride /= element_size_in_bytes;
}
// (...) - omitted for brevity
void* data_ptr = PyArray_DATA( array);
auto& type = CPU(dtype_to_aten(PyArray_TYPE( array)));
Py_INCREF(obj);
returntype.tensorFromBlob(data_ptr, sizes, strides, [obj]( void* data) {
AutoGIL gil;
Py_DECREF(obj);
});
}
程式碼摘自(tensor_numpy.cpp:https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_numpy.cpp#L88)
正如你在這段程式碼中看到的,PyTorch 從 Numpy 表徵中獲取所有資訊(陣列後設資料),並建立自己的張量。但是,正如你從被標註的第 18 行所看到的,PyTorch 保留一個指向內部 Numpy 陣列原始資料的指標,而不是複製它。這意味著 PyTorch 將擁有這一資料,並與 Numpy 陣列物件共享同一記憶體區域。
還有一點很重要:當 Numpy 陣列物件越出範圍並獲得零引用(zero reference)計數,它將被當作垃圾回收並銷燬,這就是為什麼 Numpy 陣列物件的引用計數在第 20 行有增加。該行之後,PyTorch 將從這一 Numpy 資料 blob 中建立一個新的張量物件,並且在建立這一新張量的過程中,PyTorch 將會傳遞記憶體資料指標,連同記憶體大小、步幅以及稍後張量儲存將會使用的函式(我們將會在下節討論),從而通過減少 Numpy 陣列物件的引用計數並使 Python 關心這一物件記憶體管理而釋放資料。
tensorFromBlob() 方法將建立一個新張量,但只有在為這一張量建立一個新「儲存」之後。儲存是指儲存資料指標的地方,它並不在張量結構內部。張量儲存正是我們下一節要討論的內容。
張量儲存
張量的實際原始資料並不是立即儲存在張量結構中,而是儲存在我們稱之為「儲存(Storage)」的地方,它是張量結構的一部分。
正如我們前面在 tensor_from_numpy() 中看到的程式碼,它呼叫了 tensorFromBlob() 函式以從原始資料 Blob 中建立一個張量。tensorFromBlob() 函式在內部會呼叫另一個名為 storageFromBlob() 函式,該函式主要根據型別為資料建立一個儲存。例如在 CPU 浮點型的情況下,它會返回一個新的 CPUFloatStorage 例項。
CPUFloatStorage 基本上是包含 utility 函式的包裝類(wrapper),且實際儲存結構如下所示稱為 THFloatStorage:
typedefstructTHStorage
{
real *data;
ptrdiff_tsize;
intrefcount;
charflag;
THAllocator *allocator;
void*allocatorContext;
structTHStorage*view;
} THStorage;
如上所示,THStorage 有一個指向原始資料、原始資料大小、flags 和 allocator 的指標,我們會在後面詳細地討論它們。值得注意的是,THStorage 不包含如何解釋內部資料的後設資料,這是因為儲存對儲存的內容「無處理資訊的能力」,只有張量才知道如何「檢視」資料。
因此,你可能已經意識到多個張量可以指向相同的儲存,而僅僅對資料採用不同的解析。這也就是為什麼我們以不同的形狀或維度,檢視相同元素數量的張量會有很高的效率。下面的 Python 程式碼表明,在改變張量的形狀後,儲存中的資料指標將得到共享。
>>> tensor_a = torch.ones(( 3, 3))
>>> tensor_b = tensor_a.view( 9)
>>> tensor_a.storage().data_ptr() == tensor_b.storage().data_ptr()
True
如 THFloatStorage 結構中的第七行程式碼所示,它有一個指向 THAllocator 結構的指標。它因為給分配器(allocator)帶來靈活性而顯得十分重要,其中 allocator 可以用來分配儲存資料。
typedefstructTHAllocator
{
void* (* malloc)( void*, ptrdiff_t);
void* (* realloc)( void*, void*, ptrdiff_t);
void(* free)( void*, void*);
} THAllocator;
程式碼摘自(THAllocator.h:https://github.com/pytorch/pytorch/blob/master/aten/src/TH/THAllocator.h#L16)
如上所述,該結構有三個函式指標欄位來定義分配器的意義:malloc、realloc 和 free。對於分配給 CPU 的記憶體,這些函式當然與傳統的 malloc/realloc/free POSIX 函式相關。然而當我們希望分配儲存給 GPU,我們最終會使用如 cudaMallocHost() 那樣的 CUDA 分配器,我們可以在下面的 THCudaHostAllocator malloc 函式中看到這一點。
staticvoid*THCudaHostAllocator_malloc(void* ctx, ptrdiff_tsize){
void* ptr;
if(size < 0) THError( "Invalid memory size: %ld", size);
if(size == 0) returnNULL;
THCudaCheck(cudaMallocHost(&ptr, size));
returnptr;
}
程式碼摘自(THCAllocator.c:https://github.com/pytorch/pytorch/blob/master/aten/src/THC/THCAllocator.c#L3)
如上所示,分配器呼叫了一個 cudaMallocHost() 函式。你可能已經注意到版本庫組織中有縮寫的表示模式,在瀏覽版本庫時記住這些約定非常重要,它們在 PyTorch README 檔案中有所總結:
- TH = TorcH
- THC = TorcH Cuda
- THCS = TorcH Cuda Sparse
- THCUNN = TorcH CUda Neural Network
- THD = TorcH Distributed
- THNN = TorcH Neural Network
- THS = TorcH Sparse
該約定同樣存在於函式/類別名和其它物件中,因此瞭解它們十分重要。你可以在 TH 程式碼中找到 CPU 分配器,在 THC 程式碼中找到 CUDA 分配器。最後,我們可以看到主張量 THTensor 結構的組成:
typedefstructTHTensor
{
int64_t*size;
int64_t*stride;
intnDimension;
THStorage *storage;
ptrdiff_tstorageOffset;
intrefcount;
charflag;
} THTensor;
如上,THTensor 的主要結構為張量資料保留了 size/strides/dimensions/offsets/等,同時還有儲存 THStorage。我們可以將所有這些結構總結為以下圖表:
現在,如果我們有多重處理的需求,且希望在多個不同的程式中共享張量資料,那麼我們需要一個共享記憶體的方法。否則每次另一個程式需要張量或我們希望實現 Hogwild 訓練過程以將所有不同的程式寫入相同的記憶體區域時,我們就需要在程式間建立副本,這是非常低效的。因此,我們將在下一節討論共享記憶體的特定儲存方法。
共享記憶體
共享記憶體可以用很多種不同的方法實現(依賴於支援的平臺)。PyTorch 支援部分方法,但為了簡單起見,我將討論在 MacOS 上使用 CPU(而不是 GPU)的情況。由於 PyTorch 支援多種共享記憶體的方法,由於程式碼中包含很多級的間接性,這部分會有點困難。
PyTorch 為 Python multiprocessing 模組提供了一個封裝器,可以從 torch.multiprocessing 匯入。他們對該封裝器中的實現做出了一些變動,以確保每當一個 Tensor 被放在佇列上或和其它程式共享時,PyTorch 可以確保僅有一個控制程式碼的共享記憶體會被共享,而不會共享 Tensor 的完整新副本。現在,很多人都不知道 PyTorch 中的 Tensor 方法是 share_memory_(),然而,該函式正好可以觸發那個特定 Tensor 的儲存記憶體的完整重建。該方法的執行過程是建立共享記憶體的一個區域,其可以在不同的程式中使用。最終,該函式可以呼叫以下的函式:
staticTHStorage* THPStorage_(newFilenameStorage)(ptrdiff_tsize)
{
intflags = TH_ALLOCATOR_MAPPED_SHAREDMEM | TH_ALLOCATOR_MAPPED_EXCLUSIVE;
std:: stringhandle = THPStorage_(__newHandle)();
autoctx = libshm_context_new( NULL, handle.c_str(), flags);
returnTHStorage_(newWithAllocator)(size, &THManagedSharedAllocator, ( void*)ctx);
}
如上所示,該函式使用了一個特殊的分類器 THManagedSharedAllocator 來建立另一個儲存。它首先定義了一些 flags,然後建立了一個格式為 /torch_ [process id] _ [random number] 的字串控制程式碼,最後在使用特殊的 THManagedSharedAllocator 建立新的儲存。該分配器有一個指向 PyTorch 內部庫 libshm 的函式指標,它將實現名為 Unix Domain Socket 的通訊以共享特定 quyu 的記憶體控制程式碼。這種分配器實際上是「smart allocator」的特例,因為它包含通訊控制邏輯單元,並使用了另一個稱之為 THRefcountedMapAllocator 的分配器,它將建立市級共享記憶體區域並呼叫 mmp() 以將該區域對映到程式虛擬地址空間。
現在我們可以通過手動交換共享記憶體控制程式碼而將分配給另一個程式的張量分配給一個程式,如下為 Python 示例:
>>> importtorch
>>> tensor_a = torch.ones(( 5, 5))
>>> tensor_a
11111
11111
11111
11111
11111
[torch.FloatTensor of size 5x5]
>>> tensor_a.is_shared()
False
>>> tensor_a = tensor_a.share_memory_()
>>> tensor_a.is_shared()
True
>>> tensor_a_storage = tensor_a.storage()
>>> tensor_a_storage._share_filename_()
(b'/var/tmp/tmp .0.yowqlr', b'/torch_31258_1218748506', 25)
在這段程式碼中,執行程式 A,我們就建立了一個 5×5,被 1 所填充的張量。在此之後,我們將其共享,並列印 Unix Domain Socket 地址和控制程式碼的元組。現在我們可以從另一個程式 B 中接入這一記憶體區域了:
程式 B 執行程式碼:
>>> importtorch
>>> tensor_a = torch.Tensor()
>>> tuple_info = (b'/var/tmp/tmp .0.yowqlr', b'/torch_31258_1218748506', 25)
>>> storage = torch.Storage._new_shared_filename(*tuple_info)
>>> tensor_a = torch.Tensor(storage).view(( 5, 5))
11111
11111
11111
11111
11111
[torch.FloatTensor of size 5x5]
如你所見,使用 Unix Domain Socket 地址和控制程式碼的元組資訊,我們可以接入另一個程式的張量儲存內容。如果你在程式 B 改變張量,你會看到改動也會反映在程式 A 中,因為張量之間共享著同樣的儲存區域。
DLPack:深度學習框架 Babel 的希望
現在讓我們來看看 PyTorch 程式碼庫最新的一些內容——DLPack(https://github.com/dmlc/dlpack)。DLPack 是一個記憶體張量結構的開放標準,允許張量資料在框架之間交換。非常有趣的是,這種記憶體表示是標準化的——與大多數框架已經在使用的記憶體表示方法非常類似,這就允許我們可以在框架之間共享,且完全無需複製資料。鑑於目前我們還沒有內部通訊的工具,DLPack 是一個非常了不起的創造。
它無疑會幫助我們解決今天存在於 MXNet、PyTorch 等框架上「孤島」一樣的張量表示,並允許開發者在多個深度學習框架之間自由操作,享受標準化為框架帶來的優勢。
DLPack 的核心結構 DLTensor 非常簡單,如下所示:
/*!
* brief Plain C Tensor object, does not manage memory.
*/
typedefstruct{
/*!
* brief The opaque data pointer points to the allocated data.
* This will be CUDA device pointer or cl_mem handle in OpenCL.
* This pointer is always aligns to 256 bytes as in CUDA.
*/
void* data;
/*! brief The device context of the tensor */
DLContext ctx;
/*! brief Number of dimensions */
intndim;
/*! brief The data type of the pointer*/
DLDataType dtype;
/*! brief The shape of the tensor */
int64_t* shape;
/*!
* brief strides of the tensor,
* can be NULL, indicating tensor is compact.
*/
int64_t* strides;
/*! brief The offset in bytes to the beginning pointer to data */
uint64_tbyte_offset;
} DLTensor;
程式碼來自 https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h
如你所見,這裡有一個未加工資料的資料指標,以及形態/步幅/偏移/GPU 或 CPU,以及其他 DLTensor 指向的元資訊。
這裡還有一個被稱為 DLManagedTensor 的受管理版本,其中框架可以提供一個環境,以及「刪除」函式,後者可以從借用張量來通知其他框架不再需要資源。
在 PyTorch 中,如果你想要轉換到 DLTensor 格式,或從 DLTensor 格式轉換,你可以找到 C/C++的方法,甚至 Python 方法來做這件事:
importtorch
from torch.utils importdlpack
t = torch.ones(( 5, 5))
dl = dlpack.to_dlpack(t)
這個 Python 函式會從 ATen 呼叫 toDLPack 函式,如下所示:
DLManagedTensor* toDLPack(constTensor& src){
ATenDLMTensor * atDLMTensor(newATenDLMTensor);
atDLMTensor->handle = src;
atDLMTensor->tensor.manager_ctx = atDLMTensor;
atDLMTensor->tensor.deleter = &deleter;
atDLMTensor->tensor.dl_tensor.data = src.data_ptr();
int64_tdevice_id = 0;
if(src.type().is_cuda()) {
device_id = src.get_device();
}
atDLMTensor->tensor.dl_tensor.ctx = getDLContext(src.type(), device_id);
atDLMTensor->tensor.dl_tensor.ndim = src.dim(www.feihuanyule.com
);
atDLMTensor->tensor.dl_tensor.dtype = getDLDataType(src.type());
atDLMTensor->tensor.dl_tensor.shape = const_cast< int64_t*>(src.sizes().data());
atDLMTensor->tensor.dl_tensor.strides = const_cast< int64_t*>(src.strides().data());
atDLMTensor->tensor.dl_tensor.byte_offset = 0;
return&(atDLMTensor->tensor);
}
如上所示,這是一個非常簡單的轉換,它可以將後設資料的 PyTorch 格式轉換為 DLPack 格式,並將指標指向內部張量的資料表示。