組創(chuàng)建、運算、上下文管理與稀疏存儲實戰(zhàn))
MXNet NDArray 完全指南多維數(shù)組創(chuàng)建、運算、上下文管理與稀疏存儲實戰(zhàn)【免費下載鏈接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more項目地址: https://gitcode.com/gh_mirrors/mxne/mxnetNDArray 是 Apache MXNet本倉庫 mxnet中承載數(shù)據(jù)與模型參數(shù)的核心多維數(shù)組類型幾乎所有訓練與推理流程都圍繞它展開。本文以 docs/python_docs/python/tutorials/packages/legacy/ndarray/ 下的系列教程為主體系統(tǒng)講解 NDArray 的創(chuàng)建方式、數(shù)學運算、原地操作、切片與廣播、CPU/GPU 上下文管理并深入剖析 CSRNDArray 與 RowSparseNDArray 兩種稀疏存儲格式及其在稀疏梯度更新中的應用。讀完本文你將能夠用 NDArray 完成從數(shù)據(jù)準備到多設備訓練的全部基礎操作并理解其異步執(zhí)行引擎帶來的性能特性。1. NDArray 是什么MXNet 的數(shù)據(jù)心臟在 MXNet 中NDArray是一個表示多維、同構、固定大小元素數(shù)組的對象其 Python 側類型定義位于 python/mxnet/ndarray/ndarray.py#L266-L288。它承擔兩個核心角色數(shù)據(jù)容器存放訓練樣本、中間特征、梯度等一切張量數(shù)據(jù)模型參數(shù)載體神經(jīng)網(wǎng)絡權重與偏置都以 NDArray 形式存在并可通過.context屬性被放置在 CPU 或指定 GPU 上。NDArray API 的設計刻意與 NumPy 保持高度相似方便數(shù)據(jù)科學家快速上手但它與 NumPy 有一個本質(zhì)區(qū)別——計算是異步、非阻塞的。當執(zhí)行c a * b其中a、b均為 NDArray時操作被提交給 MXNet 的 Execution Engine執(zhí)行引擎函數(shù)立即返回用戶線程可繼續(xù)執(zhí)行后續(xù)代碼即使計算尚未完成。引擎會構建計算圖在保證依賴順序的前提下對計算進行重排或合并若后續(xù)代碼對c有操作引擎會在c的結果就緒后自動啟動它們無需手寫回調(diào)。Engine 的異步調(diào)度實現(xiàn)可參考 src/engine/threaded_engine.cc#L317 的PushAsync等接口。需要強調(diào)自 MXNet 1.x 起MXNet 提供了 NumPy 兼容的新數(shù)組類mx.np.ndarray本文所述的經(jīng)典NDArray即mx.nd屬于 legacy遺留數(shù)組類相關教程目錄見 docs/python_docs/python/tutorials/packages/legacy/ndarray/index.rst。2. 創(chuàng)建 NDArray 的七種姿勢2.1 從 Python 列表創(chuàng)建nd.array最簡單的方式是直接由 Python 列表構造一維或二維數(shù)組import mxnet as mx from mxnet import nd x nd.array([1, 2, 3]) # 一維數(shù)組 print(x) y nd.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) # 3x4 二維矩陣 print(y)2.2 未初始化與填充創(chuàng)建empty/fullnd.empty(shape)只申請內(nèi)存、不初始化元素值返回的矩陣內(nèi)容不可預測可能包含任意大數(shù)nd.full(shape, value)用指定值填充整個數(shù)組。x nd.empty((3, 3)) print(x) x nd.full((3, 3), 7) print(x)2.3 全零與全一zeros/ones絕大多數(shù)場景我們希望數(shù)組被確定性初始化最常用的是全零矩陣x nd.zeros((3, 10)) print(x)全一矩陣對應nd.ones例如nd.ones((3, 4))會生成一個 3 行 4 列、元素全為 1 的矩陣。2.4 隨機初始化random_normal神經(jīng)網(wǎng)絡參數(shù)初始化時常需要隨機采樣。nd.random_normal(loc, scale, shape(...))從均值為loc、方差為scale的正態(tài)分布中采樣y nd.random_normal(0, 1, shape(3, 4)) print(y)2.5 按形狀復制zeros_likezeros_like復制一個數(shù)組的形狀但內(nèi)容置零常用于準備與某個張量同形的輸出緩沖區(qū)z nd.zeros_like(y) print(z)2.6 等間隔序列arangend.arange(n)生成[0, 1, ..., n-1]的等間隔一維數(shù)組配合.reshape可一步生成自定義形狀數(shù)據(jù)y nd.arange(6) y y.reshape((3, 1)) # 或直接鏈式調(diào)用 y nd.arange(6).reshape((3, 1))2.7 查詢數(shù)組屬性shape/size/dtype/context每個 NDArray 都可通過屬性查詢其元信息y.shape # 各維度大小組成的元組 y.size # 元素總數(shù)等于 shape 各分量之積 y.dtype # 數(shù)據(jù)類型 y.context # 數(shù)據(jù)所在設備CPU 或某塊 GPUfloat32是默認數(shù)據(jù)類型。可以顯式指定其他精度——低精度可提升性能高精度保證數(shù)值穩(wěn)定性import numpy as np a nd.array([1, 2, 3]) # 默認 float32 b nd.array([1, 2, 3], dtypenp.int32) # 32 位整型 c nd.array([1.2, 2.3], dtypenp.float16) # 16 位半精度浮點 (a.dtype, b.dtype, c.dtype)完整創(chuàng)建 API 的文檔化說明可對照 python/mxnet/ndarray/ndarray.py 中對應工廠函數(shù)。3. 數(shù)組運算從逐元素到矩陣乘法NDArray 支持大量標準數(shù)學運算且運算符被重載寫起來與 NumPy 幾乎一致。__add__、__mul__等運算符重載的實現(xiàn)可見 python/mxnet/ndarray/ndarray.py#L322-L379。3.1 逐元素運算x nd.ones((3, 4)) y nd.random_normal(0, 1, shape(3, 4)) print(x, x) print(y, y) x x y # 逐元素加法 print(x x y, x, x) x nd.array([1, 2, 3]) y nd.array([2, 2, 2]) print(x * y) # 逐元素乘法 print(nd.exp(x)) # 逐元素指數(shù)運算3.2 轉置與矩陣乘法對于二維矩陣先轉置再點積即可完成真正的矩陣乘法nd.dot(x, y.T)nd.dot是 MXNet 中最常用的矩陣運算算子之一其底層在 CPU/GPU 上均有高度優(yōu)化實現(xiàn)。4. 原地操作與內(nèi)存管理4.1 每次運算都會分配新內(nèi)存每次執(zhí)行y x y都會分配一塊新內(nèi)存存放結果然后讓y指向新對象舊內(nèi)存被釋放。用 Python 內(nèi)置的id()函數(shù)可以驗證變量引用的對象是否改變print(y, y) print(id(y):, id(y)) y y x print(after yyx, y, y) print(id(y):, id(y)) # id 已經(jīng)改變4.2 切片寫入復用緩沖區(qū)result[:] ...若想復用已分配的內(nèi)存可用切片賦值語法z nd.zeros_like(x) print(id(z):, id(z)) z[:] x y print(z[:] x y, z, z) print(id(z) is the same as before:, id(z)) # id 保持不變不過x y這一步仍會分配一個臨時緩沖區(qū)先把結果算出來再拷入z。4.3out參數(shù)徹底消除臨時緩沖區(qū)每個算子都支持out關鍵字參數(shù)直接把結果寫入指定數(shù)組nd.elemwise_add(x, y, outz) print(after nd.elemwise_add(x, y, outz), z, z, is in id(z):, id(z))z的id始終不變?nèi)虩o臨時分配。__iadd__等原地運算符x y內(nèi)部正是通過op.broadcast_add(self, other, outself)實現(xiàn)的見 python/mxnet/ndarray/ndarray.py#L326-L335效果與顯式out相同。4.4 兩種原地寫法不打算復用x時可把結果寫回x本身MXNet 提供兩種方式# 方式一切片賦值 x[:] x y # 方式二op-equals 運算符 x y在這類原地操作中MXNet 還會校驗數(shù)組是否可寫writable只讀數(shù)組會拋出ValueError。5. 切片與廣播5.1 切片語法速查NDArray 完整支持 NumPy 風格的切片a[start:end]取start到end-1的元素a[start:]取start到末尾a[:end]取開頭到end-1a[:]整個數(shù)組的副本。一維與二維讀取示例x nd.array([1, 2, 3]) s x[1:3] # 取第 2、3 個元素 print(slicing the 2nd and 3rd elements, s, s) x nd.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) s x[1:3] # 取第 2、3 行 print(slicing the 2nd and 3rd rows, s, s)5.2 寫入與多維切片切片不僅能讀還能寫x[2] 9.0 # 整行替換 x[0, 2] 9.0 # 替換單個元素 x[1:2, 1:3] 5.0 # 替換一塊子區(qū)域多維切片同樣支持按行列抽取x nd.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) s x[1:2, 1:3] # 取第 2 行第 2~3 列 s x[:, :1] # 第一列 s x[:1, :] # 第一行 s x[:, 3:] # 最后一列 s x[2:, :] # 最后一行5.3 廣播Broadcasting當?shù)途S數(shù)組與高維數(shù)組做逐元素運算時MXNet 會觸發(fā)廣播機制低維數(shù)組沿維度為 1 的軸復制擴展直至形狀匹配。x nd.ones(shape(3, 6)) y nd.arange(6) print(x y , x y)初始時y的形狀是(6)MXNet 會將其推斷為(1, 6)然后沿行方向廣播成(3, 6)再相加。廣播優(yōu)先沿最左側的軸復制因此y被解釋為(1, 6)而非(6, 1)。若想按列廣播需用.reshape顯式給出二維形狀y y.reshape((3, 1)) print(x y , x y) # 按列廣播 y nd.arange(6).reshape((3, 1)) # arange 與 reshape 一步鏈式完成6. 上下文管理CPU 與 GPU6.1 每個數(shù)組都有上下文MXNet 中每個數(shù)組都有一個 context上下文可以是 CPU也可以是某塊 GPU甚至分布式場景下的多臺服務器。合理分配數(shù)組所在設備能最小化設備間數(shù)據(jù)傳輸時間——例如在帶 GPU 的服務器上訓練時模型參數(shù)最好常駐 GPU。將數(shù)組創(chuàng)建在指定設備上通過ctx參數(shù)from mxnet import gpu from mxnet import nd z nd.ones(shape(3, 3), ctxgpu(0)) # 放在第一塊 GPU 上 print(z)沒有 GPU 時把ctxgpu(0)替換為ctxmx.cpu()即可。6.2 跨設備復制copytocopyto把一個數(shù)組復制到目標上下文x_gpu x.copyto(gpu(0)) print(x_gpu)運算符的結果與輸入處于同一上下文x_gpu z # 結果仍在 GPU(0) 上6.3 條件復制as_in_context與copyto的區(qū)別注意即使數(shù)組已經(jīng)在目標設備上z.copyto(gpu(0))仍會復制一份并分配新內(nèi)存。如果只是想“確保數(shù)組在正確設備上”應使用as_in_context()——當變量已經(jīng)在目標設備時它是一次 no-op空操作不產(chǎn)生復制print(id(z):, id(z)) z z.copyto(gpu(0)) # 總是分配新內(nèi)存id 改變 print(id(z):, id(z)) z z.as_in_context(gpu(0)) # 已在 gpu(0) 上no-opid 不變 print(id(z):, id(z)) print(z)as_in_context與copyto的實現(xiàn)可參考 python/mxnet/ndarray/ndarray.py#L2716copyto與 python/mxnet/ndarray/ndarray.py#L2862as_in_context。訓練循環(huán)中建議始終用as_in_context處理數(shù)據(jù)批次與參數(shù)避免無謂的跨設備拷貝。7. 與 NumPy 互轉及性能陷阱7.1 互轉接口NDArray 與 NumPy 互轉非常方便且轉換后的數(shù)組不共享內(nèi)存a x.asnumpy() # NDArray - numpy.ndarray type(a) y nd.array(a) # numpy.ndarray - NDArray7.2 阻塞調(diào)用會打斷異步流水線.asnumpy()、.asscalar()、.wait_to_read()、.waitall()都是阻塞操作調(diào)用時 MXNet 必須等待 Execution Engine 完成此前提交的所有異步計算才能取回結果.asnumpy()的 C 側同步實現(xiàn)即MXNDArraySyncCopyToCPU見 python/mxnet/ndarray/ndarray.py#L2652-L2659.asscalar()位于 python/mxnet/ndarray/ndarray.py#L2661。這帶來的實際體驗是如果計算圖很長某處調(diào)用.asnumpy()時會“感覺”耗時很久——真正耗時的并非轉換本身而是引擎要在此刻集中完成積壓的全部異步計算。在 GPU 上這個問題更突出數(shù)據(jù)必須從 GPU 拷貝回 CPU 才能生成np.array。7.3 用 NDArray 原生算子替代 NumPy性能最佳實踐是直接在 NDArray 上進行所有運算完全繞開 NumPy。當某個 NumPy 算子缺失時可采取三種策略策略一用若干低層算子組合出高層算子。例如 NumPy 有np.full_likeNDArray API 沒有但可以用ones乘以填充值替代from mxnet import nd import numpy as np np_y np.full_like(anp.arange(6, dtypeint), fill_value10) nd_y nd.ones(shape(6,)) * 10 np.array_equal(np_y, nd_y.asnumpy()) # True策略二尋找名稱或簽名相近的算子。例如nd.ravel_multi_index類似np.ravelnp.split按索引切分而nd.split需要傳入切分數(shù)。再如nd.pad只能處理 4 維張量低維輸入需先擴維再還原def pad_array(data, max_length): # 擴展到 4 維因為 nd.pad 只支持 4 維張量 data_expanded data.reshape(1, 1, 1, data.shape[0]) # 用常量 0 填充全部 4 個維度 data_padded nd.pad(data_expanded, modeconstant, pad_width[0, 0, 0, 0, 0, 0, 0, max_length - data.shape[0]], constant_value0) # 移除臨時維度 data_reshaped_back data_padded.reshape(max_length) return data_reshaped_back pad_array(nd.array([1, 2, 3]), max_length10) # 輸出: [ 1. 2. 3. 0. 0. 0. 0. 0. 0. 0.] # NDArray 10 cpu(0)7.4 最小化阻塞影響延遲取值的 LossBuffer 模式當不得不使用.asnumpy()或.asscalar()例如打印 loss 指標時盡量在“該值大概率已算完”的時刻再取值。經(jīng)典做法是用一個緩沖類緩存上一輪的 loss把打印推遲一個迭代讓引擎有充足時間并行完成上一輪計算from __future__ import print_function import mxnet as mx from mxnet import gluon, nd, autograd from mxnet.ndarray import NDArray from mxnet.gluon import HybridBlock import numpy as np class LossBuffer(object): 存儲 loss 值的簡單緩沖new_loss 返回上一輪 loss def __init__(self): self._loss None def new_loss(self, loss): ret self._loss self._loss loss return ret property def loss(self): return self._loss net gluon.nn.Dense(10) ce gluon.loss.SoftmaxCELoss() net.initialize() data nd.random.uniform(shape(1024, 100)) label nd.array(np.random.randint(0, 10, (1024,)), dtypeint32) train_dataset gluon.data.ArrayDataset(data, label) train_data gluon.data.DataLoader(train_dataset, batch_size128, shuffleTrue, num_workers2) trainer gluon.Trainer(net.collect_params(), optimizersgd) loss_buffer LossBuffer() for data, label in train_data: with autograd.record(): out net(data) # 保存新 loss返回上一輪 loss prev_loss loss_buffer.new_loss(ce(out, label)) loss_buffer.loss.backward() trainer.step(data.shape[0]) if prev_loss is not None: print(Loss: {}.format(np.mean(prev_loss.asnumpy())))運行時會觀察到 loss 輸出大約如下每次迭代延遲一個周期打印阻塞時間被攤薄Loss: 2.310760974884033 Loss: 2.334498643875122 Loss: 2.3244147300720215 ...8. 稀疏 NDArray一CSRNDArray 壓縮稀疏行格式8.1 為什么需要稀疏存儲現(xiàn)實世界大量數(shù)據(jù)集是高維稀疏的。以推薦系統(tǒng)為例類別與用戶數(shù)量可達百萬量級但每個用戶實際購買的類目極少絕大多數(shù)元素為 0。用默認的稠密結構存儲這類矩陣內(nèi)存與計算都被浪費在 0 上。CSRNDArray以壓縮稀疏行CSR格式存儲二維矩陣并讓算子使用專門算法。該格式面向列數(shù)眾多、且每一行都很稀疏非零元素少的 2D 矩陣。對高稀疏度矩陣如約 1% 非零、密度約 1%相比稠密NDArray有兩個主要優(yōu)勢內(nèi)存占用顯著降低部分運算如矩陣-向量乘法明顯更快。CSR 格式與 SciPy 的實現(xiàn)相似但CSRNDArray額外繼承了 NDArray 的異步非阻塞求值與自動并行化能力。同時NDArray家族新增了stype屬性用于記錄存儲類型稠密 NDArray 的stype為defaultCSRNDArray為csr。8.2 CSR 三數(shù)組表示一個 CSRNDArray 把二維矩陣表示為三個一維數(shù)組data、indptr、indicesdata矩陣非零元素按行主序排列indicesdata中每個非零元素對應的列索引indptr每行第一個非零元素在data中的偏移指針indptr[0]恒為 0indptr[i1]是到第i行為止非零元素總數(shù)的累計值。行i的列索引位于indices[indptr[i]:indptr[i1]]對應值位于data[indptr[i]:indptr[i1]]。同一行內(nèi)列索引必須升序排列且不允許出現(xiàn)重復列索引。以矩陣為例[[7, 0, 8, 0] [0, 0, 0, 0] [0, 9, 0, 0]]按行主序去除全部 0 得到data [7, 8, 9]每個非零元素所在列依次為第 0、2、1 列即indices [0, 2, 1]第一行有 2 個非零元素、第二行 0 個、第三行 1 個累計偏移得到indptr [0, 2, 2, 3]。重建時用data[0:2]indices[0:2]還原第一行data[2:2]indices[2:2]還原全零行data[2:3]indices[2:3]還原第三行。8.3 創(chuàng)建 CSRNDArray方式一由data/indices/indptr三元組創(chuàng)建可傳 Python 列表或 NumPy 數(shù)組import mxnet as mx shape (3, 4) data_list [7, 8, 9] indices_list [0, 2, 1] indptr_list [0, 2, 2, 3] a mx.nd.sparse.csr_matrix((data_list, indices_list, indptr_list), shapeshape) a.asnumpy() # array([[ 7., 0., 8., 0.], # [ 0., 0., 0., 0.], # [ 0., 9., 0., 0.]], dtypefloat32)mx.nd.sparse.csr_matrix的完整簽名含shape/ctx/dtype參數(shù)及其默認行為定義在 python/mxnet/ndarray/sparse.py#L838默認上下文為當前默認上下文默認 dtype 為float32當輸入為 NDArray/NumPy 數(shù)組時沿用其 dtype。方式二從 SciPy CSR 對象創(chuàng)建import numpy as np import scipy.sparse as spsp data_np np.array([7, 8, 9]) indptr_np np.array([0, 2, 2, 3]) indices_np np.array([0, 2, 1]) c spsp.csr.csr_matrix((data_np, indices_np, indptr_np), shapeshape) d mx.nd.sparse.array(c) # scipy csr - CSRNDArray print(d.asnumpy())方式三從稠密數(shù)組壓縮轉換。已有數(shù)據(jù)但未計算indices/indptr時可用tostype(csr)一步完成壓縮并直接訪問壓縮后的內(nèi)部數(shù)組big_array mx.nd.round(mx.nd.random.uniform(low0, high1, shape(1000, 100))) big_array_csr big_array.tostype(csr) indices big_array_csr.indices indptr big_array_csr.indptr data big_array_csr.data # data indices indptr 的總大小遠小于稠密 big_array指定元素類型mx.nd.sparse.array(a)默認float32也可通過dtype指定例如mx.nd.array(a, dtypenp.float16)生成半精度數(shù)組。8.4 檢查與存儲類型轉換檢查 CSR 數(shù)組的常用方法.asnumpy()轉成稠密numpy.ndarray查看內(nèi)容.data/.indices/.indptr查看內(nèi)部三個存儲數(shù)組.stype查看存儲類型csr。{a.stype: a.stype, data: a.data, indices: a.indices, indptr: a.indptr} # {a.stype: csr, data: [ 7. 8. 9.] NDArray 3 cpu(0), # indices: [0 2 1] NDArray 3 cpu(0), indptr: [0 2 2 3] NDArray 4 cpu(0)}存儲類型轉換有兩種等價途徑# 途徑一tostype 方法 ones mx.nd.ones((2, 2)) csr ones.tostype(csr) # default - csr dense csr.tostype(default) # csr - default # 途徑二cast_storage 算子 csr mx.nd.sparse.cast_storage(ones, csr) dense mx.nd.sparse.cast_storage(csr, default)8.5 復制與索引copy()深拷貝返回新數(shù)組copyto(dest)或切片賦值dest[:] src深拷貝到既有數(shù)組注意若源與目標存儲類型不一致copyto/[]不會改變目標數(shù)組的存儲類型源數(shù)組會被臨時轉換。a mx.nd.ones((2, 2)).tostype(csr) b a.copy() c mx.nd.sparse.zeros(csr, (2, 2)) c[:] a d mx.nd.sparse.zeros(csr, (2, 2)) a.copyto(d) # b/c/d 內(nèi)容均為全 1且 b is a 為 FalseCSRNDArray 支持沿 axis 0 切片切片返回新 CSRNDArray多維索引或沿特定軸的切片目前不支持a mx.nd.array(np.arange(6).reshape(3, 2)).tostype(csr) b a[1:2].asnumpy() # array([[ 2., 3.]]) c a[:].asnumpy() # array([[ 0., 1.], [ 2., 3.], [ 4., 5.]])8.6 稀疏算子與存儲類型推斷對稀疏數(shù)組有專門實現(xiàn)的算子集中在mx.nd.sparse下。例如dot(csr, dense)shape (3, 4) data [7, 8, 9] indptr [0, 2, 2, 3] indices [0, 2, 1] a mx.nd.sparse.csr_matrix((data, indices, indptr), shapeshape) rhs mx.nd.ones((4, 1)) out mx.nd.sparse.dot(a, rhs) # 調(diào)用專門針對 (csr, dense) 的稀疏 dot # [[ 15.], [ 0.], [ 9.]] NDArray 3x1 cpu(0)存儲類型推斷規(guī)則稀疏算子的輸出存儲類型由輸入推斷。例如a * 2的結果仍是 CSR0 乘 2 還是 0而a ones((3,4))的結果變?yōu)槌砻??？赏ㄟ^輸出對象的.stype屬性確認b a * 2 c a mx.nd.ones(shape(3, 4)) {b.stype: b.stype, c.stype: c.stype} # {b.stype: csr, c.stype: default}存儲類型回退fallback對未實現(xiàn)稀疏版本的稠密算子仍可傳入稀疏輸入但會有性能代價——MXNet 會把稀疏輸入臨時轉成稠密再計算若提供稀疏輸出則把稠密結果轉回指定稀疏格式?；赝税l(fā)生時會打印警告信息Jupyter 中顯示在終端控制臺e mx.nd.sparse.zeros(csr, a.shape) d mx.nd.log(a) # 稠密算子 稀疏輸入 - d.stype 為 default e mx.nd.log(a, oute) # 稠密算子 稀疏輸出 - e.stype 保持 csr8.7 稀疏數(shù)據(jù)加載從 CSRNDArray 批量加載mx.io.NDArrayIterdata mx.nd.array(np.arange(36).reshape((9, 4))).tostype(csr) labels np.ones([9, 1]) batch_size 3 dataiter mx.io.NDArrayIter(data, labels, batch_size, last_batch_handlediscard) [batch.data[0] for batch in dataiter] # 每個 batch 都是 CSRNDArray 3x4 cpu(0)從 libsvm 格式文件加載mx.io.LibSVMIter。libsvm 每行格式為label col_idx1:value1 col_idx2:value2 ...例如矩陣有 6 列時1 2:1.5 4:-3.5表示 label 為 1、數(shù)據(jù)為[[0, 0, 1.5, 0, -3.5, 0]]。注意列索引按行升序且從 0 開始非 1 開始data_path data.t with open(data_path, w) as fout: fout.write(1.0 0:1 2:2\n) fout.write(1.0 0:3 5:4\n) fout.write(1.0 2:5 8:6 9:7\n) fout.write(1.0 3:8\n) fout.write(-1 0:0.5 9:1.5\n) fout.write(-2.0\n) fout.write(-3.0 0:-0.6 1:2.25 2:1.25\n) fout.write(-3.0 1:2 2:-1.25\n) fout.write(4 2:-1.2\n) data_train mx.io.LibSVMIter(data_libsvmdata_path, data_shape(10,), label_shape(1,), batch_size3) for batch in data_train: print(data_train.getdata()) # CSRNDArray 3x10 cpu(0) print(data_train.getlabel()) # NDArray 3 cpu(0)9. 稀疏 NDArray二RowSparseNDArray 與稀疏梯度更新9.1 動機稀疏梯度訓練大規(guī)模稀疏模型時權重梯度往往也是稀疏的。設X為 1x2 矩陣、W為 2x3 矩陣Y dot(X, W)import mxnet as mx X mx.nd.array([[1, 0]]) W mx.nd.array([[3, 4, 5], [6, 7, 8]]) Y mx.nd.dot(X, W)則grad_W dot(X.T, ones_like(Y))為grad_W mx.nd.dot(X, mx.nd.ones_like(Y), transpose_aTrue) # [[ 1. 1. 1.] # [ 0. 0. 0.]]由于X的第 1 列全為 0grad_W的第 1 行也全為 0。真實世界中與稀疏輸入交互的參數(shù)其梯度通常存在大量全零行切片。稠密存儲與計算這些 0 行是浪費更重要的是SGD、AdaGrad、Adam 等基于梯度的優(yōu)化方法可以充分利用稀疏梯度實現(xiàn)高效更新。RowSparseNDArray以行稀疏row sparse格式存儲矩陣專為“絕大多數(shù)行切片全為零”的數(shù)組設計典型場景就是權重的稀疏梯度。9.2 行稀疏格式一個形狀為[LARGE0, D1, ..., Dn]的多維 NDArray 用兩個一維數(shù)組表示data任意 dtype、形狀[D0, D1, ..., Dn]indices一維 int64 數(shù)組、形狀[D0]元素升序排列存儲非零行切片的行索引。對應稠密數(shù)組滿足dense[rsp.indices[i], :, :, ...] rsp.data[i, :, :, ...]。典型使用場景是LARGE0 D0且大多數(shù)行切片為零。二維示例5x3 矩陣第 0、2 行非零data [[1, 2, 3], [4, 0, 5]] indices [0, 2]三維同樣支持一個 3x3x2 張量中第 0、1 個“行切片”非零則data [[[1,0],[0,2],[3,4]], [[5,0],[6,0],[0,0]]]、indices [0, 1]。RowSparseNDArray是NDArray的子類其.stype屬性值為row_sparse。9.3 創(chuàng)建與檢查import mxnet as mx import numpy as np shape (6, 2) data_list [[1, 2], [3, 4]] indices_list [1, 4] a mx.nd.sparse.row_sparse_array((data_list, indices_list), shapeshape) # RowSparseNDArray 6x2 cpu(0) b mx.nd.sparse.row_sparse_array((np.array([[1, 2], [3, 4]]), np.array([1, 4])), shapeshape)row_sparse_array定義于 python/mxnet/ndarray/sparse.py#L1036?？捎梅椒ㄅc CSRNDArray 基本一致.dtype、.asnumpy()、.data、.indices、.tostype、.cast_storage、.copy、.copyto。a.asnumpy() # array([[ 0., 0.], # [ 1., 2.], # [ 0., 0.], # [ 0., 0.], # [ 3., 4.], # [ 0., 0.]], dtypefloat32) {a.stype: a.stype, data: a.data, indices: a.indices} # {a.stype: row_sparse, data: [[1., 2.],[3., 4.]] NDArray 2x2 cpu(0), # indices: [1 4] NDArray 2 cpu(0)}存儲類型轉換與 CSR 相同tostype(row_sparse)/mx.nd.sparse.cast_storage(ones, row_sparse)均可實現(xiàn)default與row_sparse互轉。9.4 保留部分行切片retainmx.nd.sparse.retain(rsp, indices)按行索引保留指定行切片其余置零data [[1, 2], [3, 4], [5, 6]] indices [0, 2, 3] rsp mx.nd.sparse.row_sparse_array((data, indices), shape(5, 2)) rsp_retained mx.nd.sparse.retain(rsp, mx.nd.array([0, 1])) # 保留行 0、1rsp_retained.asnumpy() 中僅第 0 行非零9.5 存儲類型推斷與回退稀疏算子輸出類型同樣由輸入推斷。例如sparse.dot(csr, dense, transpose_aTrue)的輸出會被推斷為row_sparse因為轉置點積產(chǎn)生行稀疏結構lhs mx.nd.sparse.csr_matrix((data, indices, indptr), shape(3, 5)) rhs mx.nd.ones((3, 2)) transpose_dot mx.nd.sparse.dot(lhs, rhs, transpose_aTrue) # RowSparseNDArray 5x2 cpu(0)標量運算保持行稀疏a * 2結果為row_sparse與稠密數(shù)組相加則退化為稠密default。對非稀疏算子輸入/輸出會臨時轉成稠密再回退并打印警告行為與 CSRNDArray 一致。9.6 稀疏優(yōu)化器與 lazy update當梯度為row_sparse存儲且優(yōu)化器以lazy_updateTrue創(chuàng)建時MXNet 執(zhí)行稀疏梯度更新。稀疏優(yōu)化器只更新gradient.indices中出現(xiàn)的行切片對應的權重與狀態(tài)。以 SGD 為例稠密更新規(guī)則為rescaled_grad learning_rate * rescale_grad * clip(grad, clip_gradient) weight_decay * weight state momentum * state rescaled_grad weight weight - state而稀疏梯度的默認惰性更新lazy update為for row in grad.indices: rescaled_grad[row] learning_rate * rescale_grad * clip(grad[row], clip_gradient) weight_decay * weight[row] state[row] momentum[row] * state[row] rescaled_grad[row] weight[row] weight[row] - state[row]注意當weight_decay或momentum非零時惰性更新與稠密更新的優(yōu)化結果不同。如需禁用惰性更新創(chuàng)建優(yōu)化器時將lazy_update設為False。實測示例shape (4, 2) weight mx.nd.ones(shape).tostype(row_sparse) data [[1, 2], [4, 5]] indices [1, 2] grad mx.nd.sparse.row_sparse_array((data, indices), shapeshape) sgd mx.optimizer.SGD(learning_rate0.01, momentum0.01) momentum sgd.create_state(0, weight) sgd.update(0, weight, grad, momentum) # 只有行 1、2 的 weight 與 momentum 被更新 # weight [[1, 1], [0.99, 0.98], [0.96, 0.95], [1, 1]] # momentum [[0, 0], [-0.01, -0.02], [-0.04, -0.05], [0, 0]]目前 MXNet 中支持稀疏更新的優(yōu)化器有mxnet.optimizer.SGD、mxnet.optimizer.Adam與mxnet.optimizer.AdaGrad。9.7 GPU 支持默認情況下稀疏數(shù)組算子包括 CSR 與 row_sparse在 CPU 上執(zhí)行。在 GPU 上創(chuàng)建需顯式指定上下文無 GPU 時把gpu_device改為mx.cpu()import sys gpu_device mx.gpu() # 無 GPU 時改為 mx.cpu() try: a mx.nd.sparse.zeros(row_sparse, (100, 100), ctxgpu_device) a except mx.MXNetError as err: sys.stderr.write(str(err))10. 學習路徑與延伸閱讀本系列教程的完整目錄含各章節(jié)導航卡片見 docs/python_docs/python/tutorials/packages/legacy/ndarray/index.rst對應文檔源文件均位于 docs/python_docs/python/tutorials/packages/legacy/ndarray/ 下01-ndarray-intro.mdNDArray 基礎創(chuàng)建與屬性02-ndarray-operations.md運算、原地操作、切片、廣播與 NumPy 互轉03-ndarray-contexts.mdCPU/GPU 上下文管理gotchas_numpy_in_mxnet.mdNumPy 使用中的常見誤區(qū)與性能優(yōu)化sparse/csr.md 與 sparse/row_sparse.md兩種稀疏存儲格式的完整教程。如需深入源碼可重點閱讀python/mxnet/ndarray/ndarray.pyNDArray類定義#L266及asnumpy#L2635、asscalar#L2661、copyto#L2716、as_in_context#L2862等核心方法python/mxnet/ndarray/sparse.py稀疏工廠函數(shù)csr_matrix#L838、row_sparse_array#L1036、zeros#L1523與array#L1595src/engine/threaded_engine.ccExecution Engine 的異步任務調(diào)度PushAsync#L317解釋了 NDArray 運算為何異步非阻塞。11. 結語NDArray 是 MXNet 一切計算的基礎從最簡單的數(shù)組創(chuàng)建、逐元素運算、原地寫回到 CPU/GPU 上下文遷移再到為高維稀疏數(shù)據(jù)量身定制的 CSRNDArray 與 RowSparseNDArray以及與之配合的惰性稀疏優(yōu)化器。理解其異步執(zhí)行模型與存儲類型推斷規(guī)則是在 MXNet 上寫出高性能訓練代碼的前提。如果你正在用新項目對接 MXNet也可以留意官方后續(xù)推出的 NumPy 兼容數(shù)組類mx.np.ndarray但經(jīng)典 NDArray 所承載的存儲抽象、上下文管理與異步執(zhí)行思想依然一脈相承?！久赓M下載鏈接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more項目地址: https://gitcode.com/gh_mirrors/mxne/mxnet創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 MXNet NDArray 完全指南：多維數(shù)組創(chuàng)建、運算、上下文管理與稀疏存儲實戰(zhàn)
MXNet NDArray 完全指南多維數(shù)組創(chuàng)建、運算、上下文管理與稀疏存儲實戰(zhàn)【免費下載鏈接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more項目地址: https://gitcode.com/gh_mirrors/mxne/mxnetNDArray 是 Apache MXNet本倉庫 mxnet中承載數(shù)據(jù)與模型參數(shù)的核心多維數(shù)組類型幾乎所有訓練與推理流程都圍繞它展開。本文以 docs/python_docs/python/tutorials/packages/legacy/ndarray/ 下的系列教程為主體系統(tǒng)講解 NDArray 的創(chuàng)建方式、數(shù)學運算、原地操作、切片與廣播、CPU/GPU 上下文管理并深入剖析 CSRNDArray 與 RowSparseNDArray 兩種稀疏存儲格式及其在稀疏梯度更新中的應用。讀完本文你將能夠用 NDArray 完成從數(shù)據(jù)準備到多設備訓練的全部基礎操作并理解其異步執(zhí)行引擎帶來的性能特性。1. NDArray 是什么MXNet 的數(shù)據(jù)心臟在 MXNet 中NDArray是一個表示多維、同構、固定大小元素數(shù)組的對象其 Python 側類型定義位于 python/mxnet/ndarray/ndarray.py#L266-L288。它承擔兩個核心角色數(shù)據(jù)容器存放訓練樣本、中間特征、梯度等一切張量數(shù)據(jù)模型參數(shù)載體神經(jīng)網(wǎng)絡權重與偏置都以 NDArray 形式存在并可通過.context屬性被放置在 CPU 或指定 GPU 上。NDArray API 的設計刻意與 NumPy 保持高度相似方便數(shù)據(jù)科學家快速上手但它與 NumPy 有一個本質(zhì)區(qū)別——計算是異步、非阻塞的。當執(zhí)行c a * b其中a、b均為 NDArray時操作被提交給 MXNet 的 Execution Engine執(zhí)行引擎函數(shù)立即返回用戶線程可繼續(xù)執(zhí)行后續(xù)代碼即使計算尚未完成。引擎會構建計算圖在保證依賴順序的前提下對計算進行重排或合并若后續(xù)代碼對c有操作引擎會在c的結果就緒后自動啟動它們無需手寫回調(diào)。Engine 的異步調(diào)度實現(xiàn)可參考 src/engine/threaded_engine.cc#L317 的PushAsync等接口。需要強調(diào)自 MXNet 1.x 起MXNet 提供了 NumPy 兼容的新數(shù)組類mx.np.ndarray本文所述的經(jīng)典NDArray即mx.nd屬于 legacy遺留數(shù)組類相關教程目錄見 docs/python_docs/python/tutorials/packages/legacy/ndarray/index.rst。2. 創(chuàng)建 NDArray 的七種姿勢2.1 從 Python 列表創(chuàng)建nd.array最簡單的方式是直接由 Python 列表構造一維或二維數(shù)組import mxnet as mx from mxnet import nd x nd.array([1, 2, 3]) # 一維數(shù)組 print(x) y nd.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) # 3x4 二維矩陣 print(y)2.2 未初始化與填充創(chuàng)建empty/fullnd.empty(shape)只申請內(nèi)存、不初始化元素值返回的矩陣內(nèi)容不可預測可能包含任意大數(shù)nd.full(shape, value)用指定值填充整個數(shù)組。x nd.empty((3, 3)) print(x) x nd.full((3, 3), 7) print(x)2.3 全零與全一zeros/ones絕大多數(shù)場景我們希望數(shù)組被確定性初始化最常用的是全零矩陣x nd.zeros((3, 10)) print(x)全一矩陣對應nd.ones例如nd.ones((3, 4))會生成一個 3 行 4 列、元素全為 1 的矩陣。2.4 隨機初始化random_normal神經(jīng)網(wǎng)絡參數(shù)初始化時常需要隨機采樣。nd.random_normal(loc, scale, shape(...))從均值為loc、方差為scale的正態(tài)分布中采樣y nd.random_normal(0, 1, shape(3, 4)) print(y)2.5 按形狀復制zeros_likezeros_like復制一個數(shù)組的形狀但內(nèi)容置零常用于準備與某個張量同形的輸出緩沖區(qū)z nd.zeros_like(y) print(z)2.6 等間隔序列arangend.arange(n)生成[0, 1, ..., n-1]的等間隔一維數(shù)組配合.reshape可一步生成自定義形狀數(shù)據(jù)y nd.arange(6) y y.reshape((3, 1)) # 或直接鏈式調(diào)用 y nd.arange(6).reshape((3, 1))2.7 查詢數(shù)組屬性shape/size/dtype/context每個 NDArray 都可通過屬性查詢其元信息y.shape # 各維度大小組成的元組 y.size # 元素總數(shù)等于 shape 各分量之積 y.dtype # 數(shù)據(jù)類型 y.context # 數(shù)據(jù)所在設備CPU 或某塊 GPUfloat32是默認數(shù)據(jù)類型。可以顯式指定其他精度——低精度可提升性能高精度保證數(shù)值穩(wěn)定性import numpy as np a nd.array([1, 2, 3]) # 默認 float32 b nd.array([1, 2, 3], dtypenp.int32) # 32 位整型 c nd.array([1.2, 2.3], dtypenp.float16) # 16 位半精度浮點 (a.dtype, b.dtype, c.dtype)完整創(chuàng)建 API 的文檔化說明可對照 python/mxnet/ndarray/ndarray.py 中對應工廠函數(shù)。3. 數(shù)組運算從逐元素到矩陣乘法NDArray 支持大量標準數(shù)學運算且運算符被重載寫起來與 NumPy 幾乎一致。__add__、__mul__等運算符重載的實現(xiàn)可見 python/mxnet/ndarray/ndarray.py#L322-L379。3.1 逐元素運算x nd.ones((3, 4)) y nd.random_normal(0, 1, shape(3, 4)) print(x, x) print(y, y) x x y # 逐元素加法 print(x x y, x, x) x nd.array([1, 2, 3]) y nd.array([2, 2, 2]) print(x * y) # 逐元素乘法 print(nd.exp(x)) # 逐元素指數(shù)運算3.2 轉置與矩陣乘法對于二維矩陣先轉置再點積即可完成真正的矩陣乘法nd.dot(x, y.T)nd.dot是 MXNet 中最常用的矩陣運算算子之一其底層在 CPU/GPU 上均有高度優(yōu)化實現(xiàn)。4. 原地操作與內(nèi)存管理4.1 每次運算都會分配新內(nèi)存每次執(zhí)行y x y都會分配一塊新內(nèi)存存放結果然后讓y指向新對象舊內(nèi)存被釋放。用 Python 內(nèi)置的id()函數(shù)可以驗證變量引用的對象是否改變print(y, y) print(id(y):, id(y)) y y x print(after yyx, y, y) print(id(y):, id(y)) # id 已經(jīng)改變4.2 切片寫入復用緩沖區(qū)result[:] ...若想復用已分配的內(nèi)存可用切片賦值語法z nd.zeros_like(x) print(id(z):, id(z)) z[:] x y print(z[:] x y, z, z) print(id(z) is the same as before:, id(z)) # id 保持不變不過x y這一步仍會分配一個臨時緩沖區(qū)先把結果算出來再拷入z。4.3out參數(shù)徹底消除臨時緩沖區(qū)每個算子都支持out關鍵字參數(shù)直接把結果寫入指定數(shù)組nd.elemwise_add(x, y, outz) print(after nd.elemwise_add(x, y, outz), z, z, is in id(z):, id(z))z的id始終不變?nèi)虩o臨時分配。__iadd__等原地運算符x y內(nèi)部正是通過op.broadcast_add(self, other, outself)實現(xiàn)的見 python/mxnet/ndarray/ndarray.py#L326-L335效果與顯式out相同。4.4 兩種原地寫法不打算復用x時可把結果寫回x本身MXNet 提供兩種方式# 方式一切片賦值 x[:] x y # 方式二op-equals 運算符 x y在這類原地操作中MXNet 還會校驗數(shù)組是否可寫writable只讀數(shù)組會拋出ValueError。5. 切片與廣播5.1 切片語法速查NDArray 完整支持 NumPy 風格的切片a[start:end]取start到end-1的元素a[start:]取start到末尾a[:end]取開頭到end-1a[:]整個數(shù)組的副本。一維與二維讀取示例x nd.array([1, 2, 3]) s x[1:3] # 取第 2、3 個元素 print(slicing the 2nd and 3rd elements, s, s) x nd.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) s x[1:3] # 取第 2、3 行 print(slicing the 2nd and 3rd rows, s, s)5.2 寫入與多維切片切片不僅能讀還能寫x[2] 9.0 # 整行替換 x[0, 2] 9.0 # 替換單個元素 x[1:2, 1:3] 5.0 # 替換一塊子區(qū)域多維切片同樣支持按行列抽取x nd.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) s x[1:2, 1:3] # 取第 2 行第 2~3 列 s x[:, :1] # 第一列 s x[:1, :] # 第一行 s x[:, 3:] # 最后一列 s x[2:, :] # 最后一行5.3 廣播Broadcasting當?shù)途S數(shù)組與高維數(shù)組做逐元素運算時MXNet 會觸發(fā)廣播機制低維數(shù)組沿維度為 1 的軸復制擴展直至形狀匹配。x nd.ones(shape(3, 6)) y nd.arange(6) print(x y , x y)初始時y的形狀是(6)MXNet 會將其推斷為(1, 6)然后沿行方向廣播成(3, 6)再相加。廣播優(yōu)先沿最左側的軸復制因此y被解釋為(1, 6)而非(6, 1)。若想按列廣播需用.reshape顯式給出二維形狀y y.reshape((3, 1)) print(x y , x y) # 按列廣播 y nd.arange(6).reshape((3, 1)) # arange 與 reshape 一步鏈式完成6. 上下文管理CPU 與 GPU6.1 每個數(shù)組都有上下文MXNet 中每個數(shù)組都有一個 context上下文可以是 CPU也可以是某塊 GPU甚至分布式場景下的多臺服務器。合理分配數(shù)組所在設備能最小化設備間數(shù)據(jù)傳輸時間——例如在帶 GPU 的服務器上訓練時模型參數(shù)最好常駐 GPU。將數(shù)組創(chuàng)建在指定設備上通過ctx參數(shù)from mxnet import gpu from mxnet import nd z nd.ones(shape(3, 3), ctxgpu(0)) # 放在第一塊 GPU 上 print(z)沒有 GPU 時把ctxgpu(0)替換為ctxmx.cpu()即可。6.2 跨設備復制copytocopyto把一個數(shù)組復制到目標上下文x_gpu x.copyto(gpu(0)) print(x_gpu)運算符的結果與輸入處于同一上下文x_gpu z # 結果仍在 GPU(0) 上6.3 條件復制as_in_context與copyto的區(qū)別注意即使數(shù)組已經(jīng)在目標設備上z.copyto(gpu(0))仍會復制一份并分配新內(nèi)存。如果只是想“確保數(shù)組在正確設備上”應使用as_in_context()——當變量已經(jīng)在目標設備時它是一次 no-op空操作不產(chǎn)生復制print(id(z):, id(z)) z z.copyto(gpu(0)) # 總是分配新內(nèi)存id 改變 print(id(z):, id(z)) z z.as_in_context(gpu(0)) # 已在 gpu(0) 上no-opid 不變 print(id(z):, id(z)) print(z)as_in_context與copyto的實現(xiàn)可參考 python/mxnet/ndarray/ndarray.py#L2716copyto與 python/mxnet/ndarray/ndarray.py#L2862as_in_context。訓練循環(huán)中建議始終用as_in_context處理數(shù)據(jù)批次與參數(shù)避免無謂的跨設備拷貝。7. 與 NumPy 互轉及性能陷阱7.1 互轉接口NDArray 與 NumPy 互轉非常方便且轉換后的數(shù)組不共享內(nèi)存a x.asnumpy() # NDArray - numpy.ndarray type(a) y nd.array(a) # numpy.ndarray - NDArray7.2 阻塞調(diào)用會打斷異步流水線.asnumpy()、.asscalar()、.wait_to_read()、.waitall()都是阻塞操作調(diào)用時 MXNet 必須等待 Execution Engine 完成此前提交的所有異步計算才能取回結果.asnumpy()的 C 側同步實現(xiàn)即MXNDArraySyncCopyToCPU見 python/mxnet/ndarray/ndarray.py#L2652-L2659.asscalar()位于 python/mxnet/ndarray/ndarray.py#L2661。這帶來的實際體驗是如果計算圖很長某處調(diào)用.asnumpy()時會“感覺”耗時很久——真正耗時的并非轉換本身而是引擎要在此刻集中完成積壓的全部異步計算。在 GPU 上這個問題更突出數(shù)據(jù)必須從 GPU 拷貝回 CPU 才能生成np.array。7.3 用 NDArray 原生算子替代 NumPy性能最佳實踐是直接在 NDArray 上進行所有運算完全繞開 NumPy。當某個 NumPy 算子缺失時可采取三種策略策略一用若干低層算子組合出高層算子。例如 NumPy 有np.full_likeNDArray API 沒有但可以用ones乘以填充值替代from mxnet import nd import numpy as np np_y np.full_like(anp.arange(6, dtypeint), fill_value10) nd_y nd.ones(shape(6,)) * 10 np.array_equal(np_y, nd_y.asnumpy()) # True策略二尋找名稱或簽名相近的算子。例如nd.ravel_multi_index類似np.ravelnp.split按索引切分而nd.split需要傳入切分數(shù)。再如nd.pad只能處理 4 維張量低維輸入需先擴維再還原def pad_array(data, max_length): # 擴展到 4 維因為 nd.pad 只支持 4 維張量 data_expanded data.reshape(1, 1, 1, data.shape[0]) # 用常量 0 填充全部 4 個維度 data_padded nd.pad(data_expanded, modeconstant, pad_width[0, 0, 0, 0, 0, 0, 0, max_length - data.shape[0]], constant_value0) # 移除臨時維度 data_reshaped_back data_padded.reshape(max_length) return data_reshaped_back pad_array(nd.array([1, 2, 3]), max_length10) # 輸出: [ 1. 2. 3. 0. 0. 0. 0. 0. 0. 0.] # NDArray 10 cpu(0)7.4 最小化阻塞影響延遲取值的 LossBuffer 模式當不得不使用.asnumpy()或.asscalar()例如打印 loss 指標時盡量在“該值大概率已算完”的時刻再取值。經(jīng)典做法是用一個緩沖類緩存上一輪的 loss把打印推遲一個迭代讓引擎有充足時間并行完成上一輪計算from __future__ import print_function import mxnet as mx from mxnet import gluon, nd, autograd from mxnet.ndarray import NDArray from mxnet.gluon import HybridBlock import numpy as np class LossBuffer(object): 存儲 loss 值的簡單緩沖new_loss 返回上一輪 loss def __init__(self): self._loss None def new_loss(self, loss): ret self._loss self._loss loss return ret property def loss(self): return self._loss net gluon.nn.Dense(10) ce gluon.loss.SoftmaxCELoss() net.initialize() data nd.random.uniform(shape(1024, 100)) label nd.array(np.random.randint(0, 10, (1024,)), dtypeint32) train_dataset gluon.data.ArrayDataset(data, label) train_data gluon.data.DataLoader(train_dataset, batch_size128, shuffleTrue, num_workers2) trainer gluon.Trainer(net.collect_params(), optimizersgd) loss_buffer LossBuffer() for data, label in train_data: with autograd.record(): out net(data) # 保存新 loss返回上一輪 loss prev_loss loss_buffer.new_loss(ce(out, label)) loss_buffer.loss.backward() trainer.step(data.shape[0]) if prev_loss is not None: print(Loss: {}.format(np.mean(prev_loss.asnumpy())))運行時會觀察到 loss 輸出大約如下每次迭代延遲一個周期打印阻塞時間被攤薄Loss: 2.310760974884033 Loss: 2.334498643875122 Loss: 2.3244147300720215 ...8. 稀疏 NDArray一CSRNDArray 壓縮稀疏行格式8.1 為什么需要稀疏存儲現(xiàn)實世界大量數(shù)據(jù)集是高維稀疏的。以推薦系統(tǒng)為例類別與用戶數(shù)量可達百萬量級但每個用戶實際購買的類目極少絕大多數(shù)元素為 0。用默認的稠密結構存儲這類矩陣內(nèi)存與計算都被浪費在 0 上。CSRNDArray以壓縮稀疏行CSR格式存儲二維矩陣并讓算子使用專門算法。該格式面向列數(shù)眾多、且每一行都很稀疏非零元素少的 2D 矩陣。對高稀疏度矩陣如約 1% 非零、密度約 1%相比稠密NDArray有兩個主要優(yōu)勢內(nèi)存占用顯著降低部分運算如矩陣-向量乘法明顯更快。CSR 格式與 SciPy 的實現(xiàn)相似但CSRNDArray額外繼承了 NDArray 的異步非阻塞求值與自動并行化能力。同時NDArray家族新增了stype屬性用于記錄存儲類型稠密 NDArray 的stype為defaultCSRNDArray為csr。8.2 CSR 三數(shù)組表示一個 CSRNDArray 把二維矩陣表示為三個一維數(shù)組data、indptr、indicesdata矩陣非零元素按行主序排列indicesdata中每個非零元素對應的列索引indptr每行第一個非零元素在data中的偏移指針indptr[0]恒為 0indptr[i1]是到第i行為止非零元素總數(shù)的累計值。行i的列索引位于indices[indptr[i]:indptr[i1]]對應值位于data[indptr[i]:indptr[i1]]。同一行內(nèi)列索引必須升序排列且不允許出現(xiàn)重復列索引。以矩陣為例[[7, 0, 8, 0] [0, 0, 0, 0] [0, 9, 0, 0]]按行主序去除全部 0 得到data [7, 8, 9]每個非零元素所在列依次為第 0、2、1 列即indices [0, 2, 1]第一行有 2 個非零元素、第二行 0 個、第三行 1 個累計偏移得到indptr [0, 2, 2, 3]。重建時用data[0:2]indices[0:2]還原第一行data[2:2]indices[2:2]還原全零行data[2:3]indices[2:3]還原第三行。8.3 創(chuàng)建 CSRNDArray方式一由data/indices/indptr三元組創(chuàng)建可傳 Python 列表或 NumPy 數(shù)組import mxnet as mx shape (3, 4) data_list [7, 8, 9] indices_list [0, 2, 1] indptr_list [0, 2, 2, 3] a mx.nd.sparse.csr_matrix((data_list, indices_list, indptr_list), shapeshape) a.asnumpy() # array([[ 7., 0., 8., 0.], # [ 0., 0., 0., 0.], # [ 0., 9., 0., 0.]], dtypefloat32)mx.nd.sparse.csr_matrix的完整簽名含shape/ctx/dtype參數(shù)及其默認行為定義在 python/mxnet/ndarray/sparse.py#L838默認上下文為當前默認上下文默認 dtype 為float32當輸入為 NDArray/NumPy 數(shù)組時沿用其 dtype。方式二從 SciPy CSR 對象創(chuàng)建import numpy as np import scipy.sparse as spsp data_np np.array([7, 8, 9]) indptr_np np.array([0, 2, 2, 3]) indices_np np.array([0, 2, 1]) c spsp.csr.csr_matrix((data_np, indices_np, indptr_np), shapeshape) d mx.nd.sparse.array(c) # scipy csr - CSRNDArray print(d.asnumpy())方式三從稠密數(shù)組壓縮轉換。已有數(shù)據(jù)但未計算indices/indptr時可用tostype(csr)一步完成壓縮并直接訪問壓縮后的內(nèi)部數(shù)組big_array mx.nd.round(mx.nd.random.uniform(low0, high1, shape(1000, 100))) big_array_csr big_array.tostype(csr) indices big_array_csr.indices indptr big_array_csr.indptr data big_array_csr.data # data indices indptr 的總大小遠小于稠密 big_array指定元素類型mx.nd.sparse.array(a)默認float32也可通過dtype指定例如mx.nd.array(a, dtypenp.float16)生成半精度數(shù)組。8.4 檢查與存儲類型轉換檢查 CSR 數(shù)組的常用方法.asnumpy()轉成稠密numpy.ndarray查看內(nèi)容.data/.indices/.indptr查看內(nèi)部三個存儲數(shù)組.stype查看存儲類型csr。{a.stype: a.stype, data: a.data, indices: a.indices, indptr: a.indptr} # {a.stype: csr, data: [ 7. 8. 9.] NDArray 3 cpu(0), # indices: [0 2 1] NDArray 3 cpu(0), indptr: [0 2 2 3] NDArray 4 cpu(0)}存儲類型轉換有兩種等價途徑# 途徑一tostype 方法 ones mx.nd.ones((2, 2)) csr ones.tostype(csr) # default - csr dense csr.tostype(default) # csr - default # 途徑二cast_storage 算子 csr mx.nd.sparse.cast_storage(ones, csr) dense mx.nd.sparse.cast_storage(csr, default)8.5 復制與索引copy()深拷貝返回新數(shù)組copyto(dest)或切片賦值dest[:] src深拷貝到既有數(shù)組注意若源與目標存儲類型不一致copyto/[]不會改變目標數(shù)組的存儲類型源數(shù)組會被臨時轉換。a mx.nd.ones((2, 2)).tostype(csr) b a.copy() c mx.nd.sparse.zeros(csr, (2, 2)) c[:] a d mx.nd.sparse.zeros(csr, (2, 2)) a.copyto(d) # b/c/d 內(nèi)容均為全 1且 b is a 為 FalseCSRNDArray 支持沿 axis 0 切片切片返回新 CSRNDArray多維索引或沿特定軸的切片目前不支持a mx.nd.array(np.arange(6).reshape(3, 2)).tostype(csr) b a[1:2].asnumpy() # array([[ 2., 3.]]) c a[:].asnumpy() # array([[ 0., 1.], [ 2., 3.], [ 4., 5.]])8.6 稀疏算子與存儲類型推斷對稀疏數(shù)組有專門實現(xiàn)的算子集中在mx.nd.sparse下。例如dot(csr, dense)shape (3, 4) data [7, 8, 9] indptr [0, 2, 2, 3] indices [0, 2, 1] a mx.nd.sparse.csr_matrix((data, indices, indptr), shapeshape) rhs mx.nd.ones((4, 1)) out mx.nd.sparse.dot(a, rhs) # 調(diào)用專門針對 (csr, dense) 的稀疏 dot # [[ 15.], [ 0.], [ 9.]] NDArray 3x1 cpu(0)存儲類型推斷規(guī)則稀疏算子的輸出存儲類型由輸入推斷。例如a * 2的結果仍是 CSR0 乘 2 還是 0而a ones((3,4))的結果變?yōu)槌砻??？赏ㄟ^輸出對象的.stype屬性確認b a * 2 c a mx.nd.ones(shape(3, 4)) {b.stype: b.stype, c.stype: c.stype} # {b.stype: csr, c.stype: default}存儲類型回退fallback對未實現(xiàn)稀疏版本的稠密算子仍可傳入稀疏輸入但會有性能代價——MXNet 會把稀疏輸入臨時轉成稠密再計算若提供稀疏輸出則把稠密結果轉回指定稀疏格式?；赝税l(fā)生時會打印警告信息Jupyter 中顯示在終端控制臺e mx.nd.sparse.zeros(csr, a.shape) d mx.nd.log(a) # 稠密算子 稀疏輸入 - d.stype 為 default e mx.nd.log(a, oute) # 稠密算子 稀疏輸出 - e.stype 保持 csr8.7 稀疏數(shù)據(jù)加載從 CSRNDArray 批量加載mx.io.NDArrayIterdata mx.nd.array(np.arange(36).reshape((9, 4))).tostype(csr) labels np.ones([9, 1]) batch_size 3 dataiter mx.io.NDArrayIter(data, labels, batch_size, last_batch_handlediscard) [batch.data[0] for batch in dataiter] # 每個 batch 都是 CSRNDArray 3x4 cpu(0)從 libsvm 格式文件加載mx.io.LibSVMIter。libsvm 每行格式為label col_idx1:value1 col_idx2:value2 ...例如矩陣有 6 列時1 2:1.5 4:-3.5表示 label 為 1、數(shù)據(jù)為[[0, 0, 1.5, 0, -3.5, 0]]。注意列索引按行升序且從 0 開始非 1 開始data_path data.t with open(data_path, w) as fout: fout.write(1.0 0:1 2:2\n) fout.write(1.0 0:3 5:4\n) fout.write(1.0 2:5 8:6 9:7\n) fout.write(1.0 3:8\n) fout.write(-1 0:0.5 9:1.5\n) fout.write(-2.0\n) fout.write(-3.0 0:-0.6 1:2.25 2:1.25\n) fout.write(-3.0 1:2 2:-1.25\n) fout.write(4 2:-1.2\n) data_train mx.io.LibSVMIter(data_libsvmdata_path, data_shape(10,), label_shape(1,), batch_size3) for batch in data_train: print(data_train.getdata()) # CSRNDArray 3x10 cpu(0) print(data_train.getlabel()) # NDArray 3 cpu(0)9. 稀疏 NDArray二RowSparseNDArray 與稀疏梯度更新9.1 動機稀疏梯度訓練大規(guī)模稀疏模型時權重梯度往往也是稀疏的。設X為 1x2 矩陣、W為 2x3 矩陣Y dot(X, W)import mxnet as mx X mx.nd.array([[1, 0]]) W mx.nd.array([[3, 4, 5], [6, 7, 8]]) Y mx.nd.dot(X, W)則grad_W dot(X.T, ones_like(Y))為grad_W mx.nd.dot(X, mx.nd.ones_like(Y), transpose_aTrue) # [[ 1. 1. 1.] # [ 0. 0. 0.]]由于X的第 1 列全為 0grad_W的第 1 行也全為 0。真實世界中與稀疏輸入交互的參數(shù)其梯度通常存在大量全零行切片。稠密存儲與計算這些 0 行是浪費更重要的是SGD、AdaGrad、Adam 等基于梯度的優(yōu)化方法可以充分利用稀疏梯度實現(xiàn)高效更新。RowSparseNDArray以行稀疏row sparse格式存儲矩陣專為“絕大多數(shù)行切片全為零”的數(shù)組設計典型場景就是權重的稀疏梯度。9.2 行稀疏格式一個形狀為[LARGE0, D1, ..., Dn]的多維 NDArray 用兩個一維數(shù)組表示data任意 dtype、形狀[D0, D1, ..., Dn]indices一維 int64 數(shù)組、形狀[D0]元素升序排列存儲非零行切片的行索引。對應稠密數(shù)組滿足dense[rsp.indices[i], :, :, ...] rsp.data[i, :, :, ...]。典型使用場景是LARGE0 D0且大多數(shù)行切片為零。二維示例5x3 矩陣第 0、2 行非零data [[1, 2, 3], [4, 0, 5]] indices [0, 2]三維同樣支持一個 3x3x2 張量中第 0、1 個“行切片”非零則data [[[1,0],[0,2],[3,4]], [[5,0],[6,0],[0,0]]]、indices [0, 1]。RowSparseNDArray是NDArray的子類其.stype屬性值為row_sparse。9.3 創(chuàng)建與檢查import mxnet as mx import numpy as np shape (6, 2) data_list [[1, 2], [3, 4]] indices_list [1, 4] a mx.nd.sparse.row_sparse_array((data_list, indices_list), shapeshape) # RowSparseNDArray 6x2 cpu(0) b mx.nd.sparse.row_sparse_array((np.array([[1, 2], [3, 4]]), np.array([1, 4])), shapeshape)row_sparse_array定義于 python/mxnet/ndarray/sparse.py#L1036?？捎梅椒ㄅc CSRNDArray 基本一致.dtype、.asnumpy()、.data、.indices、.tostype、.cast_storage、.copy、.copyto。a.asnumpy() # array([[ 0., 0.], # [ 1., 2.], # [ 0., 0.], # [ 0., 0.], # [ 3., 4.], # [ 0., 0.]], dtypefloat32) {a.stype: a.stype, data: a.data, indices: a.indices} # {a.stype: row_sparse, data: [[1., 2.],[3., 4.]] NDArray 2x2 cpu(0), # indices: [1 4] NDArray 2 cpu(0)}存儲類型轉換與 CSR 相同tostype(row_sparse)/mx.nd.sparse.cast_storage(ones, row_sparse)均可實現(xiàn)default與row_sparse互轉。9.4 保留部分行切片retainmx.nd.sparse.retain(rsp, indices)按行索引保留指定行切片其余置零data [[1, 2], [3, 4], [5, 6]] indices [0, 2, 3] rsp mx.nd.sparse.row_sparse_array((data, indices), shape(5, 2)) rsp_retained mx.nd.sparse.retain(rsp, mx.nd.array([0, 1])) # 保留行 0、1rsp_retained.asnumpy() 中僅第 0 行非零9.5 存儲類型推斷與回退稀疏算子輸出類型同樣由輸入推斷。例如sparse.dot(csr, dense, transpose_aTrue)的輸出會被推斷為row_sparse因為轉置點積產(chǎn)生行稀疏結構lhs mx.nd.sparse.csr_matrix((data, indices, indptr), shape(3, 5)) rhs mx.nd.ones((3, 2)) transpose_dot mx.nd.sparse.dot(lhs, rhs, transpose_aTrue) # RowSparseNDArray 5x2 cpu(0)標量運算保持行稀疏a * 2結果為row_sparse與稠密數(shù)組相加則退化為稠密default。對非稀疏算子輸入/輸出會臨時轉成稠密再回退并打印警告行為與 CSRNDArray 一致。9.6 稀疏優(yōu)化器與 lazy update當梯度為row_sparse存儲且優(yōu)化器以lazy_updateTrue創(chuàng)建時MXNet 執(zhí)行稀疏梯度更新。稀疏優(yōu)化器只更新gradient.indices中出現(xiàn)的行切片對應的權重與狀態(tài)。以 SGD 為例稠密更新規(guī)則為rescaled_grad learning_rate * rescale_grad * clip(grad, clip_gradient) weight_decay * weight state momentum * state rescaled_grad weight weight - state而稀疏梯度的默認惰性更新lazy update為for row in grad.indices: rescaled_grad[row] learning_rate * rescale_grad * clip(grad[row], clip_gradient) weight_decay * weight[row] state[row] momentum[row] * state[row] rescaled_grad[row] weight[row] weight[row] - state[row]注意當weight_decay或momentum非零時惰性更新與稠密更新的優(yōu)化結果不同。如需禁用惰性更新創(chuàng)建優(yōu)化器時將lazy_update設為False。實測示例shape (4, 2) weight mx.nd.ones(shape).tostype(row_sparse) data [[1, 2], [4, 5]] indices [1, 2] grad mx.nd.sparse.row_sparse_array((data, indices), shapeshape) sgd mx.optimizer.SGD(learning_rate0.01, momentum0.01) momentum sgd.create_state(0, weight) sgd.update(0, weight, grad, momentum) # 只有行 1、2 的 weight 與 momentum 被更新 # weight [[1, 1], [0.99, 0.98], [0.96, 0.95], [1, 1]] # momentum [[0, 0], [-0.01, -0.02], [-0.04, -0.05], [0, 0]]目前 MXNet 中支持稀疏更新的優(yōu)化器有mxnet.optimizer.SGD、mxnet.optimizer.Adam與mxnet.optimizer.AdaGrad。9.7 GPU 支持默認情況下稀疏數(shù)組算子包括 CSR 與 row_sparse在 CPU 上執(zhí)行。在 GPU 上創(chuàng)建需顯式指定上下文無 GPU 時把gpu_device改為mx.cpu()import sys gpu_device mx.gpu() # 無 GPU 時改為 mx.cpu() try: a mx.nd.sparse.zeros(row_sparse, (100, 100), ctxgpu_device) a except mx.MXNetError as err: sys.stderr.write(str(err))10. 學習路徑與延伸閱讀本系列教程的完整目錄含各章節(jié)導航卡片見 docs/python_docs/python/tutorials/packages/legacy/ndarray/index.rst對應文檔源文件均位于 docs/python_docs/python/tutorials/packages/legacy/ndarray/ 下01-ndarray-intro.mdNDArray 基礎創(chuàng)建與屬性02-ndarray-operations.md運算、原地操作、切片、廣播與 NumPy 互轉03-ndarray-contexts.mdCPU/GPU 上下文管理gotchas_numpy_in_mxnet.mdNumPy 使用中的常見誤區(qū)與性能優(yōu)化sparse/csr.md 與 sparse/row_sparse.md兩種稀疏存儲格式的完整教程。如需深入源碼可重點閱讀python/mxnet/ndarray/ndarray.pyNDArray類定義#L266及asnumpy#L2635、asscalar#L2661、copyto#L2716、as_in_context#L2862等核心方法python/mxnet/ndarray/sparse.py稀疏工廠函數(shù)csr_matrix#L838、row_sparse_array#L1036、zeros#L1523與array#L1595src/engine/threaded_engine.ccExecution Engine 的異步任務調(diào)度PushAsync#L317解釋了 NDArray 運算為何異步非阻塞。11. 結語NDArray 是 MXNet 一切計算的基礎從最簡單的數(shù)組創(chuàng)建、逐元素運算、原地寫回到 CPU/GPU 上下文遷移再到為高維稀疏數(shù)據(jù)量身定制的 CSRNDArray 與 RowSparseNDArray以及與之配合的惰性稀疏優(yōu)化器。理解其異步執(zhí)行模型與存儲類型推斷規(guī)則是在 MXNet 上寫出高性能訓練代碼的前提。如果你正在用新項目對接 MXNet也可以留意官方后續(xù)推出的 NumPy 兼容數(shù)組類mx.np.ndarray但經(jīng)典 NDArray 所承載的存儲抽象、上下文管理與異步執(zhí)行思想依然一脈相承?！久赓M下載鏈接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more項目地址: https://gitcode.com/gh_mirrors/mxne/mxnet創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考