尧图网站建设 
从 PyTorch 迁移到 Apache MXNetGluon API 逐模块对照实战指南【免费下载链接】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/mx/mxnetApache MXNet 的 Gluon API 为开发者提供了与 PyTorch 同样简单、直观的命令式imperative编程体验同时额外支持通过 hybridize 将网络编译为符号图以获取性能优化。本文以 MNIST 手写数字识别MLP 模型为贯穿案例从安装、数据操作、模型构建、训练循环到多 GPU/分布式训练、监控与模型部署逐模块对照 PyTorch 与 MXNet 的写法差异并深入 Gluon 的源码级实现细节帮助有 PyTorch 经验的读者在 10 分钟内掌握迁移要点。读完本文你将能够独立把 PyTorch 代码改写成等价或更简的 Gluon 代码理解 NDArray 与 Tensor、autograd 作用域与 zero_grad、Trainer 与 optimizer 等核心概念的对应关系并利用官方 对比指南目录 持续深化迁移技能。为什么值得从 PyTorch 迁移到 MXNetPyTorch 以易于理解的 API 和完全命令式的编程范式广受欢迎Apache MXNet 的 Gluon API 既保留了这种简单性与灵活性又允许你通过 hybridize 将命令式网络转换为符号图symbolic graph从而利用符号执行的性能优化算子融合、静态内存规划等。需要说明的是官方文档中引用的 NVIDIA 训练/推理性能基准数据ResNet-50 上 MXNet 每秒训练 10,925 张图像比 PyTorch 的 6,175 张快约 77%是 2019 年 4 月的快照硬件、软件版本与今日环境差异较大实际性能需以自身环境实测为准。对迁移决策而言更可靠的参考是 API 易用性与功能覆盖度这正是本文的主体。安装pip 一条命令搞定PyTorch 默认使用 conda 安装例如# !conda install pytorch-cpu -c pytorch, torchvision而 MXNet 使用 pip 安装 CPU 版本# !pip install mxnet如果需要 GPU 支持需要指定 CUDA 版本例如安装 CUDA 10.2 版本# !pip install mxnet-cu102从当前仓库的发布脚本tools/pip和 python/setup.py 可以看出MXNet 的 pip 包按 CUDA 版本拆分为mxnet-cuXXX系列安装时务必使 CUDA 版本与本地驱动/运行时匹配。也可以直接基于本仓库源码构建参见 config 目录下的 cmake 配置如 linux_cpu.cmake。数据操作Tensor 与 NDArray 的对照两个框架都以多维矩阵作为数据载体PyTorch 沿用 Torch 的命名习惯称之为 TensorMXNet 则遵循 NumPy 惯例称之为 NDArray。以下代码创建一个 5×3、元素全为 1 的二维矩阵并将每个元素加 1 后输出PyTorchimport torch x torch.ones(5, 3) y x 1 yMXNetfrom mxnet import np x np.ones((5, 3)) y x 1 y两者最主要的差异在于MXNet 的形状shape参数需要像 NumPy 一样以元组形式传入即np.ones((5, 3))而不是torch.ones(5, 3)。当前仓库的 python/mxnet/numpy 模块正是对 NumPy 风格的完整实现其 NDArray 与 PyTorch Tensor 在绝大多数张量操作上保持一一对应详见后文附录的映射表。模型训练MNIST 上的 MLP 四步走下面用多层感知机MLP在 MNIST 数据集上做图像分类逐步对比两个框架的写法。整个任务分为 4 步。第 1 步读取数据先从网络下载 MNIST 数据集并加载到内存以便按批次读取。PyTorchfrom torchvision import datasets, transforms trans transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.13,), (0.31,))]) pt_train_data torch.utils.data.DataLoader(datasets.MNIST( root., trainTrue, downloadTrue, transformtrans), batch_size128, shuffleTrue, num_workers4)MXNetfrom mxnet import gluon from mxnet.gluon.data.vision import datasets, transforms trans transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.13, 0.31)]) mx_train_data gluon.data.DataLoader( datasets.MNIST(trainTrue).transform_first(trans), batch_size128, shuffleTrue, num_workers4)两个框架都支持从各自数据源自动下载 MNIST并只取训练集部分。核心差异在于 MXNet 使用了transform_first方法——它的作用是把数据变换只应用到每个数据批次中的第一个元素即 MNIST 图像而不作用于第二个元素标签。这一点可以从 python/mxnet/gluon/data/dataset.py 的实现得到印证transform_first(fn, lazyTrue)与transform(fn)的区别正是对样本各字段的变换范围控制。API 文档见 gluon/data/index.rst。第 2 步创建模型定义一个带单个隐藏层256 个单元和 10 个输出单元的 MLP。PyTorchimport torch.nn as pt_nn pt_net pt_nn.Sequential( pt_nn.Linear(28 * 28, 256), pt_nn.ReLU(), pt_nn.Linear(256, 10))MXNetimport mxnet.gluon.nn as mx_nn mx_net mx_nn.Sequential() mx_net.add(mx_nn.Dense(256, activationrelu), mx_nn.Dense(10)) mx_net.initialize()两者都用 Sequential 容器逐层堆叠网络但 MXNet 有三个显著差异输入尺寸自动推断PyTorch 必须在Linear的第一个参数中显式给出输入大小MXNet 的Dense只需声明输出单元数输入尺寸会在第一次前向传播后自动推断网络结构更灵活。这一机制在 Gluon 的 Block 实现中表现为惰性参数分配deferred parameter allocation。激活函数内联MXNet 可以在全连接层、卷积层中直接指定激活函数如activationrelu无需单独插入 ReLU 层。显式初始化模型结构定义完成后MXNet 要求显式调用mx_net.initialize()完成参数初始化。Sequential 的层按顺序依次执行若需要自定义执行流程PyTorch 通过继承nn.Module并重写.forward()实现MXNet 则通过继承 gluon.Block 达到同样效果__init__中声明子层、forward中定义前向计算逻辑。第 3 步损失函数与优化算法选用交叉熵损失与随机梯度下降SGD优化器。PyTorchpt_loss_fn pt_nn.CrossEntropyLoss() pt_trainer torch.optim.SGD(pt_net.parameters(), lr0.1)MXNetmx_loss_fn gluon.loss.SoftmaxCrossEntropyLoss() mx_trainer gluon.Trainer(mx_net.collect_params(), sgd, {learning_rate: 0.1})两者代码差异很小。MXNet 使用 Trainer 类优化算法作为构造参数传入字符串sgd并通过 .collect_params() 收集网络全部参数交给 Trainer 管理。SoftmaxCrossEntropyLoss定义于 python/mxnet/gluon/loss.py它将 softmax 与交叉熵合并为单一算子数值上更稳定。第 4 步训练循环实现完整的训练算法。注意由于权重初始化值不同、shuffle 导致的数据读取顺序不同每次运行的结果会略有差异。PyTorchimport time for epoch in range(5): total_loss .0 tic time.time() for X, y in pt_train_data: pt_trainer.zero_grad() loss pt_loss_fn(pt_net(X.view(-1, 28 * 28)), y) loss.backward() pt_trainer.step() total_loss loss.mean() print(epoch %d, avg loss %.4f, time %.2f % ( epoch, total_loss / len(pt_train_data), time.time() - tic))MXNetfrom mxnet import autograd for epoch in range(5): total_loss .0 tic time.time() for X, y in mx_train_data: with autograd.record(): loss mx_loss_fn(mx_net(X), y) loss.backward() mx_trainer.step(batch_size128) total_loss loss.mean().item() print(epoch %d, avg loss %.4f, time %.2f % ( epoch, total_loss / len(mx_train_data), time.time() - tic))MXNet 侧需要留意的差异点无需手动展平把 4 维输入喂给前向传播时不需要手动view(-1, 28*28)展平成 2 维Dense会自动展平输入。autograd 作用域计算必须在 autograd.record() 上下文内进行这样反向传播时才能自动求导。不需要 zero_grad默认grad_reqwrite时新梯度是覆盖写入而非累加所以无需像 PyTorch 那样每步调用trainer.zero_grad()。step 必须指定更新步长调用 step() 时需要传入 batch sizebatch_size128用于梯度归一化。标量转换需要调用 .item() 把多维数组转成 Python 标量才能参与total_loss 这类运算。官方文档提到在该示例中 MXNet 比 PyTorch 快一倍但这类玩具级对比需要谨慎看待不能作为普遍性能结论。总结两者的核心差异从上面的完整流程可以看到Gluon API 与 PyTorch 的相似度非常高迁移成本很低。主要差异集中在两点术语不同Tensor 对应 NDArraynn.Module对应gluon.Blockoptimizer对应Trainer梯度行为不同PyTorch 中梯度是累加的需要手动zero_grad()MXNet 中默认是覆盖写入的grad_reqwrite通常无需清梯度。其余代码几乎可以一一对照平移。进一步加速Hybridize 与自定义层Gluon API 与 PyTorch 相似但还提供了一些让代码更快的额外功能阅读 Hybridize 教程学习如何编写既能以命令式运行、又能通过net.hybridize()转换为符号图执行的代码HybridBlock/HybridSequential阅读 自定义层教程学习如何继承gluon.Block/HybridBlock扩展自己的网络组件。附录PyTorch 与 Gluon API 逐函数对照表以下对照表覆盖了张量操作、GPU 使用、跨设备拷贝、自动求导、模块定义、优化器、监控、数据加载与序列化等全部常用场景。张量操作对照以下列出 PyTorch Tensor 与 MXNet NDArray 中函数名不同的部分功能PyTorchMXNet Gluon逐元素反余弦x.acos()或torch.acos(x)nd.arccos(x)批量矩阵乘并累加torch.addbmm(M, batch1, batch2)nd.linalg_gemm(M, batch1, batch2)前 n-2 维被约减按元素除法后乘 v 再加 ttorch.addcdiv(t, v, t1, t2)t v*(t1/t2)矩阵乘并累加torch.addmm(M, mat1, mat2)nd.linalg_gemm(M, mat1, mat2)两向量外积加矩阵m.addr(vec1, vec2)不可用逐元素套用函数x.apply_(callable)不可用可用nd.custom(x, op)代替逐元素反正弦x.asin()或torch.asin(x)nd.arcsin(x)逐元素反正切x.atan()或torch.atan(x)nd.arctan(x)两张量的反正切x.atan2(y)或torch.atan2(x, y)不可用批量矩阵乘x.bmm(y)或torch.bmm(x, x)nd.linalg_gemm2(x, y)从伯努利分布采样x.bernoulli()不可用用柯西分布随机数填充x.cauchy_()不可用沿给定维度切分x.chunk(num_of_chunk)nd.split(x, num_outputsnum_of_chunk)数值裁剪到 [min, max]x.clamp(min, max)nd.clip(x, min, max)返回张量副本x.clone()x.copy()叉积x.cross(y)不可用沿轴累积求积x.cumprod(1)不可用沿轴累积求和x.cumsum(1)不可用首元素地址x.data_ptr()不可用创建对角张量x.diag()不可用计算张量范数x.dist()nd.norm(x)仅 L2 范数高斯误差函数x.erf()不可用广播/扩展到新形状x.expand(3, 4)x.broadcast_to([3, 4])指数分布随机填充x.exponential_()nd.random_exponential()逐元素取模x.fmod(3)nd.module(x, 3)取小数部分x.frac()x - nd.trunc(x)按索引收集元素torch.gather(x, 1, torch.LongTensor([[0,0],[1,0]]))nd.gather_nd(x, nd.array([[[0,0],[1,1]],[[0,0],[1,0]]]))最小二乘/最小范数求解B.gels(A)不可用几何分布采样x.geometric_(p)不可用张量所在设备print(x)会打印所在设备x.context重复张量x.repeat(4, 2)x.tile(4, 2)张量数据类型x.type()x.dtype散点写入torch.zeros(2, 4).scatter_(1, torch.LongTensor([[2], [3]]), 1.23)nd.scatter_nd(nd.array([1.23, 1.23]), nd.array([[0,1],[2,3]]), (2,4))返回形状x.size()x.shape元素个数x.numel()x.size转为 NumPy 数组x.numpy()x.asnumpy()对称矩阵特征分解e, v a.symeig()v, e nd.linalg.syevd(a)转置x.t()x.T均匀分布采样torch.uniform_()nd.sample_uniform()插入新维度x.unsqueeze()nd.expand_dims(x)重塑x.view(16)x.reshape((16,))按指定张量形状重塑x.view_as(y)x.reshape_like(y)类型转换x.type(type)x.astype(dtype)拷贝值到另一张量dst.copy_(src)src.copyto(dst)创建全零张量x torch.zeros(2, 3)x nd.zeros((2, 3))创建全一张量x torch.ones(2, 3)x nd.ones((2, 3))与输入同形状的全一张量y torch.ones_like(x)y nd.ones_like(x)GPU 使用与 Tensor 类似MXNet 的 NDArray 也可以通过指定 context 拷贝到 GPU 上并执行运算功能PyTorchMXNet Gluon拷贝到 GPUy torch.FloatTensor(1).cuda()y mx.nd.ones((1,), ctxmx.gpu(0))转成 NumPy 数组x y.cpu().numpy()x y.asnumpy()上下文作用域with torch.cuda.device(1):y torch.cuda.FloatTensor(1)with mx.gpu(1):y mx.nd.ones((3, 5))跨设备拷贝NDArray 可以跨多块 GPU 拷贝功能PyTorchMXNet GluonGPU 0 拷贝到 GPU 1x torch.cuda.FloatTensor(1)y x.cuda(1)x mx.nd.ones((1,), ctxmx.gpu(0))y x.as_in_context(mx.gpu(1))不同 GPU 间拷贝数据y.copy_(x)x.copyto(y)自动求导Variable 包装 vs autograd 作用域PyTorch 与 MXNet 的 autograd 包都实现了张量的自动微分。功能PyTorchMXNet Gluon记录计算图x Variable(torch.FloatTensor(1), requires_gradTrue)y x * 2y.backward()x mx.nd.ones((1,))x.attach_grad()with mx.autograd.record():y x * 2y.backward()作用域覆盖pause、train_mode、predict_mode部分算子Dropout、BatchNorm 等在训练与推理阶段行为不同MXNet 通过train_mode与predict_mode作用域控制pause作用域用于不需要计算梯度的代码段。功能PyTorchMXNet Gluon作用域覆盖不可用x mx.nd.ones((1,))with autograd.train_mode():y mx.nd.Dropout(x)with autograd.predict_mode():z mx.nd.Dropout(y)w mx.nd.ones((1,))w.attach_grad()with autograd.record():y x * wy.backward()with autograd.pause():w w.grad批末同步MXNet 采用惰性求值lazy evaluation以获得更优性能Python 线程只把操作推入后端引擎队列便立即返回。因此在训练阶段每个 batch 结束需要显式同步例如调用asnumpy()、wait_to_read()、metric.update(...)等。功能PyTorchMXNet Gluon批末同步不可用for (data, label) in train_data:with autograd.record():output net(data)L loss(output, label)L.backward()trainer.step(data.shape[0])metric.update([label], [output])PyTorch 模块与 Gluon Block自定义 Block写法几乎一致功能PyTorchMXNet Gluon自定义块class Net(torch.nn.Module):def __init__(self, D_in, D_out):super(Net, self).__init__()self.linear torch.nn.Linear(D_in, D_out)def forward(self, x):return self.linear(x)class Net(mx.gluon.Block):def __init__(self, D_in, D_out):super(Net, self).__init__()self.dense mx.gluon.nn.Dense(D_out, in_unitsD_in)def forward(self, x):return self.dense(x)参数与初始化器PyTorch 创建层时无需指定参数初始化器不同层有各自的默认初始化方式Gluon 中创建层时可以指定初始化器也可以留空。调用net.initialize(init method)后所有参数完成初始化——除显式指定了初始化器的层外其余层统一使用initialize传入的初始化方法。相关 API 见 initializer/index.rst如mx.init.Xavier()、mx.init.Normal(0.01)。功能PyTorchMXNet Gluon获取全部参数net.parameters()net.collect_params()初始化网络不可用net.initialize(mx.init.Xavier())指定层初始化器layer torch.nn.Linear(20, 10)torch.nn.init.normal(layer.weight, 0, 0.01)layer mx.gluon.nn.Dense(10, weight_initializermx.init.Normal(0.01))已有 Block 的调用方式一致功能PyTorchMXNet Gluon调用已有块y net(x)y net(x)HybridBlock可混合化、支持部分形状信息HybridBlock 同时支持以 Symbol 和 NDArray 前向执行。混合化hybridize后HybridBlock 会创建并缓存代表前向计算的符号图。仓库中绝大多数内置块Dense、Conv2D、MaxPool2D、BatchNorm 等都是 HybridBlock。无需显式声明层的输入数量只需给出输出数量形状会在网络收到输入时自动推断。功能PyTorchMXNet Gluon部分形状 混合化不可用net mx.gluon.nn.HybridSequential()net.add(mx.gluon.nn.Dense(10))net.hybridize()SymbolBlock从符号图构建模型SymbolBlock 可以从 Symbol 构建 Block非常适合把预训练模型用作特征提取器feature extractor。以 model zoo 中的 AlexNet 为例alexnet mx.gluon.model_zoo.vision.alexnet(pretrainedTrue) out alexnet(inputs) internals out.get_internals() outputs [internals[model_dense0_relu_fwd_output]] feat_model gluon.SymbolBlock(outputs, inputs, paramsalexnet.collect_params())该能力是 PyTorch 没有的相关 API 见 gluon/symbol_block.rst。PyTorch 优化器与 Gluon Trainer大多数情况下无需调用 zero_gradzero_gradPyTorch optimizer 或 Gluon Trainer 中用于清空所有参数的梯度。在 Gluon 中只要grad_reqwrite默认值每个 batch 都无需手动清梯度。功能PyTorchMXNet Gluon清梯度optm torch.optim.SGD(model.parameters(), lr0.1)optm.zero_grad()loss_fn(model(input), target).backward()optm.step()trainer gluon.Trainer(net.collect_params(), sgd, {learning_rate: 0.1})with autograd.record():loss loss_fn(net(data), label)loss.backward()trainer.step(batch_size)多 GPU 训练功能PyTorchMXNet Gluon数据并行net torch.nn.DataParallel(model, device_ids[0, 1, 2])output net(data)ctx [mx.gpu(i) for i in range(3)]data gluon.utils.split_and_load(data, ctx)label gluon.utils.split_and_load(label, ctx)with autograd.record():losses [loss(net(X), Y) for X, Y in zip(data, label)]for l in losses:l.backward()分布式训练功能PyTorchMXNet Gluon分布式数据并行torch.distributed.init_process_group(...)model torch.nn.parallel.DistributedDataParallel(model, ...)store kv.create(dist)trainer gluon.Trainer(net.collect_params(), ..., kvstorestore)MXNet 侧通过 kvstore键值存储抽象分布式通信kv.create(dist)创建分布式 store 并传入 Trainer 即可相关实现可参考 src/kvstore 目录含kvstore_dist.h、kvstore_nccl.h等后端及 kvstore/index.rst。训练监控预定义评估指标Gluon 提供了多个预定义指标可以在训练过程中在线评估模型表现mx.metric.Accuracy()等见 gluon/metric/index.rst。功能PyTorchMXNet Gluon指标不可用metric mx.metric.Accuracy()with autograd.record():output net(data)L loss(output, label)L.backward()trainer.step(batch_size)metric.update(label, output)数据可视化PyTorch 常用 TensorboardXMXNet 对应工具为 MXBoard均可用于可视化网络结构并绘制图执行的定量指标曲线PyTorchMXNetsw tensorboardX.SummaryWriter()sw mxboard.SummaryWriter()......for name, param in model.named_parameters():for name, param in net.collect_params():grad param.clone().cpu().data.numpy()grad param.grad.asnumpy().flatten()sw.add_histogram(name, grad, n_iter)sw.add_histogram(tagstr(param),...valuesgrad,sw.close()bins200,global_stepi)...sw.close()I/O 与部署数据加载Dataset与DataLoader是加载数据的基本组件。类PyTorchMXNet Gluon持有数组的数据集torch.utils.data.TensorDataset(data_tensor, label_tensor)gluon.data.ArrayDataset(data_array, label_array)数据加载器torch.utils.data.DataLoader(dataset, batch_size1, shuffleFalse, samplerNone, batch_samplerNone, num_workers0, collate_fnfunction default_collate, drop_lastFalse)gluon.data.DataLoader(dataset, batch_sizeNone, shuffleFalse, samplerNone, last_batchkeep, batch_samplerNone, batchify_fnNone, num_workers0)顺序采样器torch.utils.data.sampler.SequentialSampler(data_source)gluon.data.SequentialSampler(length)随机采样器torch.utils.data.sampler.RandomSampler(data_source)gluon.data.RandomSampler(length)注意默认参数差异PyTorch 的 DataLoader 默认drop_lastFalse、batch_size1而 Gluon 的 DataLoader 默认last_batchkeep、batch_sizeNone后者由数据集长度与批次划分决定。此外 Gluon 常用视觉数据集位于mx.gluon.data.vision包内类PyTorchMXNet GluonMNIST 手写数字torchvision.datasets.MNISTmx.gluon.data.vision.MNISTCIFAR10torchvision.datasets.CIFAR10mx.gluon.data.vision.CIFAR10CIFAR100torchvision.datasets.CIFAR100mx.gluon.data.vision.CIFAR100按文件夹组织的通用图像加载器torchvision.datasets.ImageFolder(root, transformNone, target_transformNone, loaderfunction default_loader)mx.gluon.data.vision.ImageFolderDataset(root, flag, transformNone)序列化模型的序列化与反序列化通过save_parameters与load_parameters完成类PyTorchMXNet Gluon保存模型参数torch.save(the_model.state_dict(), filename)model.save_parameters(filename)加载参数the_model.load_state_dict(torch.load(PATH))model.load_parameters(filename, ctx, allow_missingFalse, ignore_extraFalse)load_parameters的可选参数如allow_missing、ignore_extra可用于迁移学习场景下部分层参数缺失或新增层的兼容加载更完整的用法可参考 保存与加载参数教程。迁移路径与下一步综合以上对照从 PyTorch 迁移到 MXNet 的推荐路径是先照搬结构把torch.nn.Sequential/nn.Module翻译为gluon.nn.Sequential/gluon.Block把torch.optim.*翻译为gluon.Trainer把 DataLoader 翻译为gluon.data.DataLoader理解三个关键差异NDArray 形状传元组、autograd.record()作用域包住前向计算、Trainer.step(batch_size)必须带 batch size按需进阶需要更高性能时用hybridize()需要预训练特征提取时用SymbolBlock需要多卡/分布式时用split_and_load kvstore持续对照官方维护的 PyTorch to MXNet 对比指南 和 Gluon 官方教程 是迁移过程中的权威参考本仓库的 example/gluon 目录还提供了 MNIST、房价预测等可运行的 Gluon 示例代码可直接对照阅读。【免费下载链接】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/mx/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 从 PyTorch 迁移到 Apache MXNet：Gluon API 逐模块对照实战指南
从 PyTorch 迁移到 Apache MXNetGluon API 逐模块对照实战指南【免费下载链接】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/mx/mxnetApache MXNet 的 Gluon API 为开发者提供了与 PyTorch 同样简单、直观的命令式imperative编程体验同时额外支持通过 hybridize 将网络编译为符号图以获取性能优化。本文以 MNIST 手写数字识别MLP 模型为贯穿案例从安装、数据操作、模型构建、训练循环到多 GPU/分布式训练、监控与模型部署逐模块对照 PyTorch 与 MXNet 的写法差异并深入 Gluon 的源码级实现细节帮助有 PyTorch 经验的读者在 10 分钟内掌握迁移要点。读完本文你将能够独立把 PyTorch 代码改写成等价或更简的 Gluon 代码理解 NDArray 与 Tensor、autograd 作用域与 zero_grad、Trainer 与 optimizer 等核心概念的对应关系并利用官方 对比指南目录 持续深化迁移技能。为什么值得从 PyTorch 迁移到 MXNetPyTorch 以易于理解的 API 和完全命令式的编程范式广受欢迎Apache MXNet 的 Gluon API 既保留了这种简单性与灵活性又允许你通过 hybridize 将命令式网络转换为符号图symbolic graph从而利用符号执行的性能优化算子融合、静态内存规划等。需要说明的是官方文档中引用的 NVIDIA 训练/推理性能基准数据ResNet-50 上 MXNet 每秒训练 10,925 张图像比 PyTorch 的 6,175 张快约 77%是 2019 年 4 月的快照硬件、软件版本与今日环境差异较大实际性能需以自身环境实测为准。对迁移决策而言更可靠的参考是 API 易用性与功能覆盖度这正是本文的主体。安装pip 一条命令搞定PyTorch 默认使用 conda 安装例如# !conda install pytorch-cpu -c pytorch, torchvision而 MXNet 使用 pip 安装 CPU 版本# !pip install mxnet如果需要 GPU 支持需要指定 CUDA 版本例如安装 CUDA 10.2 版本# !pip install mxnet-cu102从当前仓库的发布脚本tools/pip和 python/setup.py 可以看出MXNet 的 pip 包按 CUDA 版本拆分为mxnet-cuXXX系列安装时务必使 CUDA 版本与本地驱动/运行时匹配。也可以直接基于本仓库源码构建参见 config 目录下的 cmake 配置如 linux_cpu.cmake。数据操作Tensor 与 NDArray 的对照两个框架都以多维矩阵作为数据载体PyTorch 沿用 Torch 的命名习惯称之为 TensorMXNet 则遵循 NumPy 惯例称之为 NDArray。以下代码创建一个 5×3、元素全为 1 的二维矩阵并将每个元素加 1 后输出PyTorchimport torch x torch.ones(5, 3) y x 1 yMXNetfrom mxnet import np x np.ones((5, 3)) y x 1 y两者最主要的差异在于MXNet 的形状shape参数需要像 NumPy 一样以元组形式传入即np.ones((5, 3))而不是torch.ones(5, 3)。当前仓库的 python/mxnet/numpy 模块正是对 NumPy 风格的完整实现其 NDArray 与 PyTorch Tensor 在绝大多数张量操作上保持一一对应详见后文附录的映射表。模型训练MNIST 上的 MLP 四步走下面用多层感知机MLP在 MNIST 数据集上做图像分类逐步对比两个框架的写法。整个任务分为 4 步。第 1 步读取数据先从网络下载 MNIST 数据集并加载到内存以便按批次读取。PyTorchfrom torchvision import datasets, transforms trans transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.13,), (0.31,))]) pt_train_data torch.utils.data.DataLoader(datasets.MNIST( root., trainTrue, downloadTrue, transformtrans), batch_size128, shuffleTrue, num_workers4)MXNetfrom mxnet import gluon from mxnet.gluon.data.vision import datasets, transforms trans transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.13, 0.31)]) mx_train_data gluon.data.DataLoader( datasets.MNIST(trainTrue).transform_first(trans), batch_size128, shuffleTrue, num_workers4)两个框架都支持从各自数据源自动下载 MNIST并只取训练集部分。核心差异在于 MXNet 使用了transform_first方法——它的作用是把数据变换只应用到每个数据批次中的第一个元素即 MNIST 图像而不作用于第二个元素标签。这一点可以从 python/mxnet/gluon/data/dataset.py 的实现得到印证transform_first(fn, lazyTrue)与transform(fn)的区别正是对样本各字段的变换范围控制。API 文档见 gluon/data/index.rst。第 2 步创建模型定义一个带单个隐藏层256 个单元和 10 个输出单元的 MLP。PyTorchimport torch.nn as pt_nn pt_net pt_nn.Sequential( pt_nn.Linear(28 * 28, 256), pt_nn.ReLU(), pt_nn.Linear(256, 10))MXNetimport mxnet.gluon.nn as mx_nn mx_net mx_nn.Sequential() mx_net.add(mx_nn.Dense(256, activationrelu), mx_nn.Dense(10)) mx_net.initialize()两者都用 Sequential 容器逐层堆叠网络但 MXNet 有三个显著差异输入尺寸自动推断PyTorch 必须在Linear的第一个参数中显式给出输入大小MXNet 的Dense只需声明输出单元数输入尺寸会在第一次前向传播后自动推断网络结构更灵活。这一机制在 Gluon 的 Block 实现中表现为惰性参数分配deferred parameter allocation。激活函数内联MXNet 可以在全连接层、卷积层中直接指定激活函数如activationrelu无需单独插入 ReLU 层。显式初始化模型结构定义完成后MXNet 要求显式调用mx_net.initialize()完成参数初始化。Sequential 的层按顺序依次执行若需要自定义执行流程PyTorch 通过继承nn.Module并重写.forward()实现MXNet 则通过继承 gluon.Block 达到同样效果__init__中声明子层、forward中定义前向计算逻辑。第 3 步损失函数与优化算法选用交叉熵损失与随机梯度下降SGD优化器。PyTorchpt_loss_fn pt_nn.CrossEntropyLoss() pt_trainer torch.optim.SGD(pt_net.parameters(), lr0.1)MXNetmx_loss_fn gluon.loss.SoftmaxCrossEntropyLoss() mx_trainer gluon.Trainer(mx_net.collect_params(), sgd, {learning_rate: 0.1})两者代码差异很小。MXNet 使用 Trainer 类优化算法作为构造参数传入字符串sgd并通过 .collect_params() 收集网络全部参数交给 Trainer 管理。SoftmaxCrossEntropyLoss定义于 python/mxnet/gluon/loss.py它将 softmax 与交叉熵合并为单一算子数值上更稳定。第 4 步训练循环实现完整的训练算法。注意由于权重初始化值不同、shuffle 导致的数据读取顺序不同每次运行的结果会略有差异。PyTorchimport time for epoch in range(5): total_loss .0 tic time.time() for X, y in pt_train_data: pt_trainer.zero_grad() loss pt_loss_fn(pt_net(X.view(-1, 28 * 28)), y) loss.backward() pt_trainer.step() total_loss loss.mean() print(epoch %d, avg loss %.4f, time %.2f % ( epoch, total_loss / len(pt_train_data), time.time() - tic))MXNetfrom mxnet import autograd for epoch in range(5): total_loss .0 tic time.time() for X, y in mx_train_data: with autograd.record(): loss mx_loss_fn(mx_net(X), y) loss.backward() mx_trainer.step(batch_size128) total_loss loss.mean().item() print(epoch %d, avg loss %.4f, time %.2f % ( epoch, total_loss / len(mx_train_data), time.time() - tic))MXNet 侧需要留意的差异点无需手动展平把 4 维输入喂给前向传播时不需要手动view(-1, 28*28)展平成 2 维Dense会自动展平输入。autograd 作用域计算必须在 autograd.record() 上下文内进行这样反向传播时才能自动求导。不需要 zero_grad默认grad_reqwrite时新梯度是覆盖写入而非累加所以无需像 PyTorch 那样每步调用trainer.zero_grad()。step 必须指定更新步长调用 step() 时需要传入 batch sizebatch_size128用于梯度归一化。标量转换需要调用 .item() 把多维数组转成 Python 标量才能参与total_loss 这类运算。官方文档提到在该示例中 MXNet 比 PyTorch 快一倍但这类玩具级对比需要谨慎看待不能作为普遍性能结论。总结两者的核心差异从上面的完整流程可以看到Gluon API 与 PyTorch 的相似度非常高迁移成本很低。主要差异集中在两点术语不同Tensor 对应 NDArraynn.Module对应gluon.Blockoptimizer对应Trainer梯度行为不同PyTorch 中梯度是累加的需要手动zero_grad()MXNet 中默认是覆盖写入的grad_reqwrite通常无需清梯度。其余代码几乎可以一一对照平移。进一步加速Hybridize 与自定义层Gluon API 与 PyTorch 相似但还提供了一些让代码更快的额外功能阅读 Hybridize 教程学习如何编写既能以命令式运行、又能通过net.hybridize()转换为符号图执行的代码HybridBlock/HybridSequential阅读 自定义层教程学习如何继承gluon.Block/HybridBlock扩展自己的网络组件。附录PyTorch 与 Gluon API 逐函数对照表以下对照表覆盖了张量操作、GPU 使用、跨设备拷贝、自动求导、模块定义、优化器、监控、数据加载与序列化等全部常用场景。张量操作对照以下列出 PyTorch Tensor 与 MXNet NDArray 中函数名不同的部分功能PyTorchMXNet Gluon逐元素反余弦x.acos()或torch.acos(x)nd.arccos(x)批量矩阵乘并累加torch.addbmm(M, batch1, batch2)nd.linalg_gemm(M, batch1, batch2)前 n-2 维被约减按元素除法后乘 v 再加 ttorch.addcdiv(t, v, t1, t2)t v*(t1/t2)矩阵乘并累加torch.addmm(M, mat1, mat2)nd.linalg_gemm(M, mat1, mat2)两向量外积加矩阵m.addr(vec1, vec2)不可用逐元素套用函数x.apply_(callable)不可用可用nd.custom(x, op)代替逐元素反正弦x.asin()或torch.asin(x)nd.arcsin(x)逐元素反正切x.atan()或torch.atan(x)nd.arctan(x)两张量的反正切x.atan2(y)或torch.atan2(x, y)不可用批量矩阵乘x.bmm(y)或torch.bmm(x, x)nd.linalg_gemm2(x, y)从伯努利分布采样x.bernoulli()不可用用柯西分布随机数填充x.cauchy_()不可用沿给定维度切分x.chunk(num_of_chunk)nd.split(x, num_outputsnum_of_chunk)数值裁剪到 [min, max]x.clamp(min, max)nd.clip(x, min, max)返回张量副本x.clone()x.copy()叉积x.cross(y)不可用沿轴累积求积x.cumprod(1)不可用沿轴累积求和x.cumsum(1)不可用首元素地址x.data_ptr()不可用创建对角张量x.diag()不可用计算张量范数x.dist()nd.norm(x)仅 L2 范数高斯误差函数x.erf()不可用广播/扩展到新形状x.expand(3, 4)x.broadcast_to([3, 4])指数分布随机填充x.exponential_()nd.random_exponential()逐元素取模x.fmod(3)nd.module(x, 3)取小数部分x.frac()x - nd.trunc(x)按索引收集元素torch.gather(x, 1, torch.LongTensor([[0,0],[1,0]]))nd.gather_nd(x, nd.array([[[0,0],[1,1]],[[0,0],[1,0]]]))最小二乘/最小范数求解B.gels(A)不可用几何分布采样x.geometric_(p)不可用张量所在设备print(x)会打印所在设备x.context重复张量x.repeat(4, 2)x.tile(4, 2)张量数据类型x.type()x.dtype散点写入torch.zeros(2, 4).scatter_(1, torch.LongTensor([[2], [3]]), 1.23)nd.scatter_nd(nd.array([1.23, 1.23]), nd.array([[0,1],[2,3]]), (2,4))返回形状x.size()x.shape元素个数x.numel()x.size转为 NumPy 数组x.numpy()x.asnumpy()对称矩阵特征分解e, v a.symeig()v, e nd.linalg.syevd(a)转置x.t()x.T均匀分布采样torch.uniform_()nd.sample_uniform()插入新维度x.unsqueeze()nd.expand_dims(x)重塑x.view(16)x.reshape((16,))按指定张量形状重塑x.view_as(y)x.reshape_like(y)类型转换x.type(type)x.astype(dtype)拷贝值到另一张量dst.copy_(src)src.copyto(dst)创建全零张量x torch.zeros(2, 3)x nd.zeros((2, 3))创建全一张量x torch.ones(2, 3)x nd.ones((2, 3))与输入同形状的全一张量y torch.ones_like(x)y nd.ones_like(x)GPU 使用与 Tensor 类似MXNet 的 NDArray 也可以通过指定 context 拷贝到 GPU 上并执行运算功能PyTorchMXNet Gluon拷贝到 GPUy torch.FloatTensor(1).cuda()y mx.nd.ones((1,), ctxmx.gpu(0))转成 NumPy 数组x y.cpu().numpy()x y.asnumpy()上下文作用域with torch.cuda.device(1):y torch.cuda.FloatTensor(1)with mx.gpu(1):y mx.nd.ones((3, 5))跨设备拷贝NDArray 可以跨多块 GPU 拷贝功能PyTorchMXNet GluonGPU 0 拷贝到 GPU 1x torch.cuda.FloatTensor(1)y x.cuda(1)x mx.nd.ones((1,), ctxmx.gpu(0))y x.as_in_context(mx.gpu(1))不同 GPU 间拷贝数据y.copy_(x)x.copyto(y)自动求导Variable 包装 vs autograd 作用域PyTorch 与 MXNet 的 autograd 包都实现了张量的自动微分。功能PyTorchMXNet Gluon记录计算图x Variable(torch.FloatTensor(1), requires_gradTrue)y x * 2y.backward()x mx.nd.ones((1,))x.attach_grad()with mx.autograd.record():y x * 2y.backward()作用域覆盖pause、train_mode、predict_mode部分算子Dropout、BatchNorm 等在训练与推理阶段行为不同MXNet 通过train_mode与predict_mode作用域控制pause作用域用于不需要计算梯度的代码段。功能PyTorchMXNet Gluon作用域覆盖不可用x mx.nd.ones((1,))with autograd.train_mode():y mx.nd.Dropout(x)with autograd.predict_mode():z mx.nd.Dropout(y)w mx.nd.ones((1,))w.attach_grad()with autograd.record():y x * wy.backward()with autograd.pause():w w.grad批末同步MXNet 采用惰性求值lazy evaluation以获得更优性能Python 线程只把操作推入后端引擎队列便立即返回。因此在训练阶段每个 batch 结束需要显式同步例如调用asnumpy()、wait_to_read()、metric.update(...)等。功能PyTorchMXNet Gluon批末同步不可用for (data, label) in train_data:with autograd.record():output net(data)L loss(output, label)L.backward()trainer.step(data.shape[0])metric.update([label], [output])PyTorch 模块与 Gluon Block自定义 Block写法几乎一致功能PyTorchMXNet Gluon自定义块class Net(torch.nn.Module):def __init__(self, D_in, D_out):super(Net, self).__init__()self.linear torch.nn.Linear(D_in, D_out)def forward(self, x):return self.linear(x)class Net(mx.gluon.Block):def __init__(self, D_in, D_out):super(Net, self).__init__()self.dense mx.gluon.nn.Dense(D_out, in_unitsD_in)def forward(self, x):return self.dense(x)参数与初始化器PyTorch 创建层时无需指定参数初始化器不同层有各自的默认初始化方式Gluon 中创建层时可以指定初始化器也可以留空。调用net.initialize(init method)后所有参数完成初始化——除显式指定了初始化器的层外其余层统一使用initialize传入的初始化方法。相关 API 见 initializer/index.rst如mx.init.Xavier()、mx.init.Normal(0.01)。功能PyTorchMXNet Gluon获取全部参数net.parameters()net.collect_params()初始化网络不可用net.initialize(mx.init.Xavier())指定层初始化器layer torch.nn.Linear(20, 10)torch.nn.init.normal(layer.weight, 0, 0.01)layer mx.gluon.nn.Dense(10, weight_initializermx.init.Normal(0.01))已有 Block 的调用方式一致功能PyTorchMXNet Gluon调用已有块y net(x)y net(x)HybridBlock可混合化、支持部分形状信息HybridBlock 同时支持以 Symbol 和 NDArray 前向执行。混合化hybridize后HybridBlock 会创建并缓存代表前向计算的符号图。仓库中绝大多数内置块Dense、Conv2D、MaxPool2D、BatchNorm 等都是 HybridBlock。无需显式声明层的输入数量只需给出输出数量形状会在网络收到输入时自动推断。功能PyTorchMXNet Gluon部分形状 混合化不可用net mx.gluon.nn.HybridSequential()net.add(mx.gluon.nn.Dense(10))net.hybridize()SymbolBlock从符号图构建模型SymbolBlock 可以从 Symbol 构建 Block非常适合把预训练模型用作特征提取器feature extractor。以 model zoo 中的 AlexNet 为例alexnet mx.gluon.model_zoo.vision.alexnet(pretrainedTrue) out alexnet(inputs) internals out.get_internals() outputs [internals[model_dense0_relu_fwd_output]] feat_model gluon.SymbolBlock(outputs, inputs, paramsalexnet.collect_params())该能力是 PyTorch 没有的相关 API 见 gluon/symbol_block.rst。PyTorch 优化器与 Gluon Trainer大多数情况下无需调用 zero_gradzero_gradPyTorch optimizer 或 Gluon Trainer 中用于清空所有参数的梯度。在 Gluon 中只要grad_reqwrite默认值每个 batch 都无需手动清梯度。功能PyTorchMXNet Gluon清梯度optm torch.optim.SGD(model.parameters(), lr0.1)optm.zero_grad()loss_fn(model(input), target).backward()optm.step()trainer gluon.Trainer(net.collect_params(), sgd, {learning_rate: 0.1})with autograd.record():loss loss_fn(net(data), label)loss.backward()trainer.step(batch_size)多 GPU 训练功能PyTorchMXNet Gluon数据并行net torch.nn.DataParallel(model, device_ids[0, 1, 2])output net(data)ctx [mx.gpu(i) for i in range(3)]data gluon.utils.split_and_load(data, ctx)label gluon.utils.split_and_load(label, ctx)with autograd.record():losses [loss(net(X), Y) for X, Y in zip(data, label)]for l in losses:l.backward()分布式训练功能PyTorchMXNet Gluon分布式数据并行torch.distributed.init_process_group(...)model torch.nn.parallel.DistributedDataParallel(model, ...)store kv.create(dist)trainer gluon.Trainer(net.collect_params(), ..., kvstorestore)MXNet 侧通过 kvstore键值存储抽象分布式通信kv.create(dist)创建分布式 store 并传入 Trainer 即可相关实现可参考 src/kvstore 目录含kvstore_dist.h、kvstore_nccl.h等后端及 kvstore/index.rst。训练监控预定义评估指标Gluon 提供了多个预定义指标可以在训练过程中在线评估模型表现mx.metric.Accuracy()等见 gluon/metric/index.rst。功能PyTorchMXNet Gluon指标不可用metric mx.metric.Accuracy()with autograd.record():output net(data)L loss(output, label)L.backward()trainer.step(batch_size)metric.update(label, output)数据可视化PyTorch 常用 TensorboardXMXNet 对应工具为 MXBoard均可用于可视化网络结构并绘制图执行的定量指标曲线PyTorchMXNetsw tensorboardX.SummaryWriter()sw mxboard.SummaryWriter()......for name, param in model.named_parameters():for name, param in net.collect_params():grad param.clone().cpu().data.numpy()grad param.grad.asnumpy().flatten()sw.add_histogram(name, grad, n_iter)sw.add_histogram(tagstr(param),...valuesgrad,sw.close()bins200,global_stepi)...sw.close()I/O 与部署数据加载Dataset与DataLoader是加载数据的基本组件。类PyTorchMXNet Gluon持有数组的数据集torch.utils.data.TensorDataset(data_tensor, label_tensor)gluon.data.ArrayDataset(data_array, label_array)数据加载器torch.utils.data.DataLoader(dataset, batch_size1, shuffleFalse, samplerNone, batch_samplerNone, num_workers0, collate_fnfunction default_collate, drop_lastFalse)gluon.data.DataLoader(dataset, batch_sizeNone, shuffleFalse, samplerNone, last_batchkeep, batch_samplerNone, batchify_fnNone, num_workers0)顺序采样器torch.utils.data.sampler.SequentialSampler(data_source)gluon.data.SequentialSampler(length)随机采样器torch.utils.data.sampler.RandomSampler(data_source)gluon.data.RandomSampler(length)注意默认参数差异PyTorch 的 DataLoader 默认drop_lastFalse、batch_size1而 Gluon 的 DataLoader 默认last_batchkeep、batch_sizeNone后者由数据集长度与批次划分决定。此外 Gluon 常用视觉数据集位于mx.gluon.data.vision包内类PyTorchMXNet GluonMNIST 手写数字torchvision.datasets.MNISTmx.gluon.data.vision.MNISTCIFAR10torchvision.datasets.CIFAR10mx.gluon.data.vision.CIFAR10CIFAR100torchvision.datasets.CIFAR100mx.gluon.data.vision.CIFAR100按文件夹组织的通用图像加载器torchvision.datasets.ImageFolder(root, transformNone, target_transformNone, loaderfunction default_loader)mx.gluon.data.vision.ImageFolderDataset(root, flag, transformNone)序列化模型的序列化与反序列化通过save_parameters与load_parameters完成类PyTorchMXNet Gluon保存模型参数torch.save(the_model.state_dict(), filename)model.save_parameters(filename)加载参数the_model.load_state_dict(torch.load(PATH))model.load_parameters(filename, ctx, allow_missingFalse, ignore_extraFalse)load_parameters的可选参数如allow_missing、ignore_extra可用于迁移学习场景下部分层参数缺失或新增层的兼容加载更完整的用法可参考 保存与加载参数教程。迁移路径与下一步综合以上对照从 PyTorch 迁移到 MXNet 的推荐路径是先照搬结构把torch.nn.Sequential/nn.Module翻译为gluon.nn.Sequential/gluon.Block把torch.optim.*翻译为gluon.Trainer把 DataLoader 翻译为gluon.data.DataLoader理解三个关键差异NDArray 形状传元组、autograd.record()作用域包住前向计算、Trainer.step(batch_size)必须带 batch size按需进阶需要更高性能时用hybridize()需要预训练特征提取时用SymbolBlock需要多卡/分布式时用split_and_load kvstore持续对照官方维护的 PyTorch to MXNet 对比指南 和 Gluon 官方教程 是迁移过程中的权威参考本仓库的 example/gluon 目录还提供了 MNIST、房价预测等可运行的 Gluon 示例代码可直接对照阅读。【免费下载链接】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/mx/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考