MXNet NDArray 完全指南多维数组创建、运算、上下文管理与稀疏存储实战【免费下载链接】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中承载数据与模型参数的核心多维数组类型几乎所有训练与推理流程都围绕它展开。本文以 docs/python_docs/python/tutorials/packages/legacy/ndarray/ 下的系列教程为主体系统讲解 NDArray 的创建方式、数学运算、原地操作、切片与广播、CPU/GPU 上下文管理并深入剖析 CSRNDArray 与 RowSparseNDArray 两种稀疏存储格式及其在稀疏梯度更新中的应用。读完本文你将能够用 NDArray 完成从数据准备到多设备训练的全部基础操作并理解其异步执行引擎带来的性能特性。1. NDArray 是什么MXNet 的数据心脏在 MXNet 中NDArray是一个表示多维、同构、固定大小元素数组的对象其 Python 侧类型定义位于 python/mxnet/ndarray/ndarray.py#L266-L288。它承担两个核心角色数据容器存放训练样本、中间特征、梯度等一切张量数据模型参数载体神经网络权重与偏置都以 NDArray 形式存在并可通过.context属性被放置在 CPU 或指定 GPU 上。NDArray API 的设计刻意与 NumPy 保持高度相似方便数据科学家快速上手但它与 NumPy 有一个本质区别——计算是异步、非阻塞的。当执行c a * b其中a、b均为 NDArray时操作被提交给 MXNet 的 Execution Engine执行引擎函数立即返回用户线程可继续执行后续代码即使计算尚未完成。引擎会构建计算图在保证依赖顺序的前提下对计算进行重排或合并若后续代码对c有操作引擎会在c的结果就绪后自动启动它们无需手写回调。Engine 的异步调度实现可参考 src/engine/threaded_engine.cc#L317 的PushAsync等接口。需要强调自 MXNet 1.x 起MXNet 提供了 NumPy 兼容的新数组类mx.np.ndarray本文所述的经典NDArray即mx.nd属于 legacy遗留数组类相关教程目录见 docs/python_docs/python/tutorials/packages/legacy/ndarray/index.rst。2. 创建 NDArray 的七种姿势2.1 从 Python 列表创建nd.array最简单的方式是直接由 Python 列表构造一维或二维数组import mxnet as mx from mxnet import nd x nd.array([1, 2, 3]) # 一维数组 print(x) y nd.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) # 3x4 二维矩阵 print(y)2.2 未初始化与填充创建empty/fullnd.empty(shape)只申请内存、不初始化元素值返回的矩阵内容不可预测可能包含任意大数nd.full(shape, value)用指定值填充整个数组。x nd.empty((3, 3)) print(x) x nd.full((3, 3), 7) print(x)2.3 全零与全一zeros/ones绝大多数场景我们希望数组被确定性初始化最常用的是全零矩阵x nd.zeros((3, 10)) print(x)全一矩阵对应nd.ones例如nd.ones((3, 4))会生成一个 3 行 4 列、元素全为 1 的矩阵。2.4 随机初始化random_normal神经网络参数初始化时常需要随机采样。nd.random_normal(loc, scale, shape(...))从均值为loc、方差为scale的正态分布中采样y nd.random_normal(0, 1, shape(3, 4)) print(y)2.5 按形状复制zeros_likezeros_like复制一个数组的形状但内容置零常用于准备与某个张量同形的输出缓冲区z nd.zeros_like(y) print(z)2.6 等间隔序列arangend.arange(n)生成[0, 1, ..., n-1]的等间隔一维数组配合.reshape可一步生成自定义形状数据y nd.arange(6) y y.reshape((3, 1)) # 或直接链式调用 y nd.arange(6).reshape((3, 1))2.7 查询数组属性shape/size/dtype/context每个 NDArray 都可通过属性查询其元信息y.shape # 各维度大小组成的元组 y.size # 元素总数等于 shape 各分量之积 y.dtype # 数据类型 y.context # 数据所在设备CPU 或某块 GPUfloat32是默认数据类型。可以显式指定其他精度——低精度可提升性能高精度保证数值稳定性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)完整创建 API 的文档化说明可对照 python/mxnet/ndarray/ndarray.py 中对应工厂函数。3. 数组运算从逐元素到矩阵乘法NDArray 支持大量标准数学运算且运算符被重载写起来与 NumPy 几乎一致。__add__、__mul__等运算符重载的实现可见 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)) # 逐元素指数运算3.2 转置与矩阵乘法对于二维矩阵先转置再点积即可完成真正的矩阵乘法nd.dot(x, y.T)nd.dot是 MXNet 中最常用的矩阵运算算子之一其底层在 CPU/GPU 上均有高度优化实现。4. 原地操作与内存管理4.1 每次运算都会分配新内存每次执行y x y都会分配一块新内存存放结果然后让y指向新对象旧内存被释放。用 Python 内置的id()函数可以验证变量引用的对象是否改变print(y, y) print(id(y):, id(y)) y y x print(after yyx, y, y) print(id(y):, id(y)) # id 已经改变4.2 切片写入复用缓冲区result[:] ...若想复用已分配的内存可用切片赋值语法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这一步仍会分配一个临时缓冲区先把结果算出来再拷入z。4.3out参数彻底消除临时缓冲区每个算子都支持out关键字参数直接把结果写入指定数组nd.elemwise_add(x, y, outz) print(after nd.elemwise_add(x, y, outz), z, z, is in id(z):, id(z))z的id始终不变全程无临时分配。__iadd__等原地运算符x y内部正是通过op.broadcast_add(self, other, outself)实现的见 python/mxnet/ndarray/ndarray.py#L326-L335效果与显式out相同。4.4 两种原地写法不打算复用x时可把结果写回x本身MXNet 提供两种方式# 方式一切片赋值 x[:] x y # 方式二op-equals 运算符 x y在这类原地操作中MXNet 还会校验数组是否可写writable只读数组会抛出ValueError。5. 切片与广播5.1 切片语法速查NDArray 完整支持 NumPy 风格的切片a[start:end]取start到end-1的元素a[start:]取start到末尾a[:end]取开头到end-1a[:]整个数组的副本。一维与二维读取示例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 # 替换一块子区域多维切片同样支持按行列抽取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当低维数组与高维数组做逐元素运算时MXNet 会触发广播机制低维数组沿维度为 1 的轴复制扩展直至形状匹配。x nd.ones(shape(3, 6)) y nd.arange(6) print(x y , x y)初始时y的形状是(6)MXNet 会将其推断为(1, 6)然后沿行方向广播成(3, 6)再相加。广播优先沿最左侧的轴复制因此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 每个数组都有上下文MXNet 中每个数组都有一个 context上下文可以是 CPU也可以是某块 GPU甚至分布式场景下的多台服务器。合理分配数组所在设备能最小化设备间数据传输时间——例如在带 GPU 的服务器上训练时模型参数最好常驻 GPU。将数组创建在指定设备上通过ctx参数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把一个数组复制到目标上下文x_gpu x.copyto(gpu(0)) print(x_gpu)运算符的结果与输入处于同一上下文x_gpu z # 结果仍在 GPU(0) 上6.3 条件复制as_in_context与copyto的区别注意即使数组已经在目标设备上z.copyto(gpu(0))仍会复制一份并分配新内存。如果只是想“确保数组在正确设备上”应使用as_in_context()——当变量已经在目标设备时它是一次 no-op空操作不产生复制print(id(z):, id(z)) z z.copyto(gpu(0)) # 总是分配新内存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的实现可参考 python/mxnet/ndarray/ndarray.py#L2716copyto与 python/mxnet/ndarray/ndarray.py#L2862as_in_context。训练循环中建议始终用as_in_context处理数据批次与参数避免无谓的跨设备拷贝。7. 与 NumPy 互转及性能陷阱7.1 互转接口NDArray 与 NumPy 互转非常方便且转换后的数组不共享内存a x.asnumpy() # NDArray - numpy.ndarray type(a) y nd.array(a) # numpy.ndarray - NDArray7.2 阻塞调用会打断异步流水线.asnumpy()、.asscalar()、.wait_to_read()、.waitall()都是阻塞操作调用时 MXNet 必须等待 Execution Engine 完成此前提交的所有异步计算才能取回结果.asnumpy()的 C 侧同步实现即MXNDArraySyncCopyToCPU见 python/mxnet/ndarray/ndarray.py#L2652-L2659.asscalar()位于 python/mxnet/ndarray/ndarray.py#L2661。这带来的实际体验是如果计算图很长某处调用.asnumpy()时会“感觉”耗时很久——真正耗时的并非转换本身而是引擎要在此刻集中完成积压的全部异步计算。在 GPU 上这个问题更突出数据必须从 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需要传入切分数。再如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 指标时尽量在“该值大概率已算完”的时刻再取值。经典做法是用一个缓冲类缓存上一轮的 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 为什么需要稀疏存储现实世界大量数据集是高维稀疏的。以推荐系统为例类别与用户数量可达百万量级但每个用户实际购买的类目极少绝大多数元素为 0。用默认的稠密结构存储这类矩阵内存与计算都被浪费在 0 上。CSRNDArray以压缩稀疏行CSR格式存储二维矩阵并让算子使用专门算法。该格式面向列数众多、且每一行都很稀疏非零元素少的 2D 矩阵。对高稀疏度矩阵如约 1% 非零、密度约 1%相比稠密NDArray有两个主要优势内存占用显著降低部分运算如矩阵-向量乘法明显更快。CSR 格式与 SciPy 的实现相似但CSRNDArray额外继承了 NDArray 的异步非阻塞求值与自动并行化能力。同时NDArray家族新增了stype属性用于记录存储类型稠密 NDArray 的stype为defaultCSRNDArray为csr。8.2 CSR 三数组表示一个 CSRNDArray 把二维矩阵表示为三个一维数组data、indptr、indicesdata矩阵非零元素按行主序排列indicesdata中每个非零元素对应的列索引indptr每行第一个非零元素在data中的偏移指针indptr[0]恒为 0indptr[i1]是到第i行为止非零元素总数的累计值。行i的列索引位于indices[indptr[i]:indptr[i1]]对应值位于data[indptr[i]:indptr[i1]]。同一行内列索引必须升序排列且不允许出现重复列索引。以矩阵为例[[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 创建 CSRNDArray方式一由data/indices/indptr三元组创建可传 Python 列表或 NumPy 数组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参数及其默认行为定义在 python/mxnet/ndarray/sparse.py#L838默认上下文为当前默认上下文默认 dtype 为float32当输入为 NDArray/NumPy 数组时沿用其 dtype。方式二从 SciPy CSR 对象创建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())方式三从稠密数组压缩转换。已有数据但未计算indices/indptr时可用tostype(csr)一步完成压缩并直接访问压缩后的内部数组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)生成半精度数组。8.4 检查与存储类型转换检查 CSR 数组的常用方法.asnumpy()转成稠密numpy.ndarray查看内容.data/.indices/.indptr查看内部三个存储数组.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()深拷贝返回新数组copyto(dest)或切片赋值dest[:] src深拷贝到既有数组注意若源与目标存储类型不一致copyto/[]不会改变目标数组的存储类型源数组会被临时转换。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 内容均为全 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 稀疏算子与存储类型推断对稀疏数组有专门实现的算子集中在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) # 调用专门针对 (csr, dense) 的稀疏 dot # [[ 15.], [ 0.], [ 9.]] NDArray 3x1 cpu(0)存储类型推断规则稀疏算子的输出存储类型由输入推断。例如a * 2的结果仍是 CSR0 乘 2 还是 0而a ones((3,4))的结果变为稠密。可通过输出对象的.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对未实现稀疏版本的稠密算子仍可传入稀疏输入但会有性能代价——MXNet 会把稀疏输入临时转成稠密再计算若提供稀疏输出则把稠密结果转回指定稀疏格式。回退发生时会打印警告信息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 稀疏数据加载从 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、数据为[[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 动机稀疏梯度训练大规模稀疏模型时权重梯度往往也是稀疏的。设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。真实世界中与稀疏输入交互的参数其梯度通常存在大量全零行切片。稠密存储与计算这些 0 行是浪费更重要的是SGD、AdaGrad、Adam 等基于梯度的优化方法可以充分利用稀疏梯度实现高效更新。RowSparseNDArray以行稀疏row sparse格式存储矩阵专为“绝大多数行切片全为零”的数组设计典型场景就是权重的稀疏梯度。9.2 行稀疏格式一个形状为[LARGE0, D1, ..., Dn]的多维 NDArray 用两个一维数组表示data任意 dtype、形状[D0, D1, ..., Dn]indices一维 int64 数组、形状[D0]元素升序排列存储非零行切片的行索引。对应稠密数组满足dense[rsp.indices[i], :, :, ...] rsp.data[i, :, :, ...]。典型使用场景是LARGE0 D0且大多数行切片为零。二维示例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 创建与检查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。可用方法与 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)均可实现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因为转置点积产生行稀疏结构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与稠密数组相加则退化为稠密default。对非稀疏算子输入/输出会临时转成稠密再回退并打印警告行为与 CSRNDArray 一致。9.6 稀疏优化器与 lazy update当梯度为row_sparse存储且优化器以lazy_updateTrue创建时MXNet 执行稀疏梯度更新。稀疏优化器只更新gradient.indices中出现的行切片对应的权重与状态。以 SGD 为例稠密更新规则为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非零时惰性更新与稠密更新的优化结果不同。如需禁用惰性更新创建优化器时将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 中支持稀疏更新的优化器有mxnet.optimizer.SGD、mxnet.optimizer.Adam与mxnet.optimizer.AdaGrad。9.7 GPU 支持默认情况下稀疏数组算子包括 CSR 与 row_sparse在 CPU 上执行。在 GPU 上创建需显式指定上下文无 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. 学习路径与延伸阅读本系列教程的完整目录含各章节导航卡片见 docs/python_docs/python/tutorials/packages/legacy/ndarray/index.rst对应文档源文件均位于 docs/python_docs/python/tutorials/packages/legacy/ndarray/ 下01-ndarray-intro.mdNDArray 基础创建与属性02-ndarray-operations.md运算、原地操作、切片、广播与 NumPy 互转03-ndarray-contexts.mdCPU/GPU 上下文管理gotchas_numpy_in_mxnet.mdNumPy 使用中的常见误区与性能优化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稀疏工厂函数csr_matrix#L838、row_sparse_array#L1036、zeros#L1523与array#L1595src/engine/threaded_engine.ccExecution Engine 的异步任务调度PushAsync#L317解释了 NDArray 运算为何异步非阻塞。11. 结语NDArray 是 MXNet 一切计算的基础从最简单的数组创建、逐元素运算、原地写回到 CPU/GPU 上下文迁移再到为高维稀疏数据量身定制的 CSRNDArray 与 RowSparseNDArray以及与之配合的惰性稀疏优化器。理解其异步执行模型与存储类型推断规则是在 MXNet 上写出高性能训练代码的前提。如果你正在用新项目对接 MXNet也可以留意官方后续推出的 NumPy 兼容数组类mx.np.ndarray但经典 NDArray 所承载的存储抽象、上下文管理与异步执行思想依然一脉相承。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考