MXNet 使用 GPU 训练与推理实战：设备分配、多卡数据并行与底层原理解析

MXNet 使用 GPU 训练与推理实战设备分配、多卡数据并行与底层原理解析【免费下载链接】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本文是 MXNet 快速入门系列Crash Course的收官教程完整讲解如何在 MXNet 中启用 GPU 加速从安装 GPU 版 MXNet 与 CUDA 环境、查询可用 GPU 数量到把数据与神经网络参数分配到指定 GPU、在多张 GPU 之间复制数据、在 GPU 上执行运算再到通过数据并行Data Parallelism用多张 GPU 联合训练神经网络。读者学完后将掌握npx.gpu()、npx.num_gpus()、ndarray.device、copyto、gluon.utils.split_and_load、net.initialize(device...)等核心 API 的完整用法并能理解这些 API 在 MXNet 源码中的底层实现从而在自己的训练任务中正确、高效地利用 GPU 资源。1. 前置条件GPU 硬件、CUDA 与 GPU 版 MXNet在使用 GPU 训练或部署神经网络之前需要确认三件事机器上至少有一块 NVIDIA GPU。MXNet 的 GPU 加速基于 CUDA 实现因此只支持 NVIDIA GPUAMD 与 Intel 的 GPU 不受支持。CUDA 已正确安装且驱动版本与 CUDA 工具包版本匹配。安装了 GPU 版CUDA 构建的 MXNet。CPU 版 MXNet 无法访问 GPU 设备需要按对应系统的安装说明安装 GPU 版本。在代码层面GPU 相关功能统一由npxNumPy 兼容接口暴露导入方式如下from mxnet import np, npx, gluon, autograd from mxnet.gluon import nn import time npx.set_np() # 切换到 NumPy 兼容的算子语义2. 查询可用 GPU 数量npx.num_gpus()运行下面的命令即可查看 MXNet 当前可以访问的 GPU 数量npx.num_gpus() # 返回 MXNet 可访问的 GPU 数量源码层面的实现npx.num_gpus()最终调用 python/mxnet/device.py#L231-L246 中定义的num_gpus()函数该函数通过ctypes调用底层 C APIMXGetGPUCount定义于 src/c_api/c_api.cc#L1933向 CUDA 查询 GPU 数量若 CUDA 调用出错会直接抛出异常。同文件中还提供了gpu_memory_info(device_id0)python/mxnet/device.py#L249-L269可通过 C APIMXGetGPUMemoryInformation64查询指定 GPU 的可用与总显存字节数便于训练前评估显存是否充足。3. 分配数据到 GPUdevice 属性与 npx.gpu()MXNet 的ndarray与 NumPy 的数组非常相似但有一个关键区别每个 ndarray 都有一个device属性指明该数组存放在哪个设备上。默认情况下数组存放在npx.cpu()CPU 主内存上。要把数组放到第一块 GPU 上使用npx.gpu()或npx.gpu(0)两者等价都表示 0 号 GPUgpu npx.gpu() if npx.num_gpus() 0 else npx.cpu() x np.ones((3, 4), devicegpu) x几点说明npx.gpu(device_id0)的device_id缺省为 0其实现见 python/mxnet/device.py#L199-L228本质是Device(gpu, device_id)的快捷方式。上面这段代码是无 GPU 也能安全运行的写法当机器没有 GPU 时自动回退到 CPU避免在纯 CPU 环境下抛错。当使用 CPU 时MXNet 把数据分配在主内存上并会尽量利用所有可用的 CPU 核心当使用多块 GPU 时MXNet 会在打印 ndarray 时明确显示该数组被分配在哪块 GPU 上例如gpu(0)。4. 在多块 GPU 之间移动数据copyto 与隐式拷贝假设机器上至少有 2 块 GPU。你可以把x复制到第二块 GPU1 号 GPU上gpu_1 npx.gpu(1) if npx.num_gpus() 1 else npx.cpu() x.copyto(gpu_1)注意如果机器只有 1 块 GPU上面的代码会因访问不存在的 1 号设备而报错。因此示例中同样加入了npx.num_gpus() 1的保护判断。MXNet 要求用户显式地在设备之间移动数据。copyto的语义见 python/mxnet/ndarray/ndarray.py#L2716-L2760当参数是Device对象时它会在目标设备上新建一个同 shape、同 dtype 的 ndarray并通过内部算子把数据复制过去返回新数组当参数是另一个NDArray时则把值复制到该数组中要求形状一致。不过也有例外——若干算子会隐式地把数据搬回主内存例如print打印 ndarray 时需要读取其内容asnumpy()把 MXNet ndarray 转成numpy.ndarray。其实现见 python/mxnet/ndarray/ndarray.py#L2635-L2659内部通过MXNDArraySyncCopyToCPU做一次同步的 GPU→CPU 拷贝bfloat16 数据会先转成 float32 再拷贝。这种隐式拷贝是即用即取不会改变原数组所在的设备。5. 选择 GPU ID0 起始编号当机器有多块 GPU 时MXNet 通过0 起始索引0-indexing访问每一块 GPU第一块 GPUnpx.gpu(0)第二块 GPUnpx.gpu(1)……第 K 块 GPUnpx.gpu(K-1)。例如机器有 8 块 GPU最后一块就用npx.gpu(7)访问。这种机制让你能够精确地指定某个操作或某次训练用哪几块 GPU在多卡训练中尤其有用——例如把不同卡分配给不同任务或把数据并行训练限定在指定的卡集合上。6. 在 GPU 上运行运算要在某块 GPU 上执行一个运算只需要保证该运算的输入数据已经在目标 GPU 上输出会自动分配在同一块 GPU 上。np与npx模块中几乎所有的算子都支持在 GPU 上运行y np.random.uniform(size(3, 4), devicegpu) x y这里x与y都创建在gpu0 号 GPU上因此加法运算在 GPU 上执行结果也驻留在 GPU 上。重要约束如果一个运算的多个输入不在同一块 GPU 上会直接报错。例如x在 gpu(0)、y在 gpu(1) 时执行x y是不允许的——必须先通过copyto把其中一个输入显式搬移到另一块卡上。这也再次体现了设备间数据移动必须显式的设计原则。7. 在 GPU 上运行神经网络在 GPU 上运行一个神经网络本质上只需要把输入数据和网络参数都放到 GPU 上。这里复用快速入门系列 Step 6训练神经网络 中定义的LeafNetwork一个基于卷积块与全连接块构建的HybridBlock# 卷积块包含卷积层、最大池化层和批归一化层 def conv_block(filters, kernel_size2, stride2, batch_normTrue): conv_block nn.HybridSequential() conv_block.add(nn.Conv2D(channelsfilters, kernel_sizekernel_size, activationrelu), nn.MaxPool2D(pool_size4, stridesstride)) if batch_norm: conv_block.add(nn.BatchNorm()) return conv_block # 全连接块包含全连接层和 Dropout 层 def dense_block(neurons, activationrelu, dropout0.2): dense_block nn.HybridSequential() dense_block.add(nn.Dense(neurons, activationactivation)) if dropout: dense_block.add(nn.Dropout(dropout)) return dense_block # 用上述两种块组合出完整网络 class LeafNetwork(nn.HybridBlock): def __init__(self): super(LeafNetwork, self).__init__() self.conv1 conv_block(32) self.conv2 conv_block(64) self.conv3 conv_block(128) self.flatten nn.Flatten() self.dense1 dense_block(100) self.dense2 dense_block(10) self.dense3 nn.Dense(2) def forward(self, batch): batch self.conv1(batch) batch self.conv2(batch) batch self.conv3(batch) batch self.flatten(batch) batch self.dense1(batch) batch self.dense2(batch) batch self.dense3(batch) return batch7.1 把已保存的参数直接加载到 GPU第 6 步中模型通过model.save_parameters(leaf_models.params)保存了训练好的权重。加载时可以直接把参数载入到 0 号 GPU 上net LeafNetwork() net.load_parameters(leaf_models.params, devicegpu)load_parameters的device参数见 python/mxnet/gluon/block.py#L381-L434既可以是单个Device也可以是Device列表用于把参数直接初始化到指定设备上未指定时默认使用当前设备device.current_device()。除了加载时指定设备你还可以在之后随时用net.collect_params().reset_device(gpu)把全部参数重新分配到新的设备实现见 python/mxnet/gluon/block.py#L758-L769当传入Device列表时会在每个设备上保留一份参数副本。7.2 在 GPU 上执行前向推理在 0 号 GPU 上创建输入数据前向函数就会在 0 号 GPU 上执行x np.random.uniform(size(1, 3, 128, 128), devicegpu) net(x)输入张量的形状(1, 3, 128, 128)对应LeafNetwork期望的(batch, channels, height, width)布局与第 6 步中model.summary使用的(4, 3, 128, 128)一致。推理结束后如需在 CPU 侧进一步处理结果例如用matplotlib绘图直接调用asnumpy()即可隐式把结果拷回主内存。8. 使用多块 GPU 联合训练数据并行训练是 GPU 加速收益最大的场景。MXNet 支持通过数据并行Data Parallelism让多块 GPU 联合训练一个网络假设有n块 GPU就把每一个数据批次batch沿批次维平均切分成n份每块 GPU 负责其中一份的前向forward与反向backward计算。8.1 数据准备与 DataLoader首先复用第 6 步的数据定义与变换函数用gluon.data.vision.transforms构造训练/验证/测试三套变换流水线并用ImageFolderDatasetgluon.data.DataLoader创建数据加载器from mxnet.gluon.data.vision import transforms jitter_param 0.05 # 归一化用的 mean 与 std图像取值在 (0,1) 范围内 mean [0.485, 0.456, 0.406] std [0.229, 0.224, 0.225] training_transformer transforms.Compose([ transforms.Resize(size224, keep_ratioTrue), transforms.CenterCrop(128), transforms.RandomFlipLeftRight(), transforms.RandomColorJitter(contrastjitter_param), transforms.ToTensor(), transforms.Normalize(mean, std) ]) validation_transformer transforms.Compose([ transforms.Resize(size224, keep_ratioTrue), transforms.CenterCrop(128), transforms.ToTensor(), transforms.Normalize(mean, std) ]) # 用 ImageFolderDataset 从目录结构创建 Dataset train_dataset gluon.data.vision.ImageFolderDataset(./datasets/train) val_dataset gluon.data.vision.ImageFolderDataset(./datasets/validation) test_dataset gluon.data.vision.ImageFolderDataset(./datasets/test) # 创建数据加载器 batch_size 4 train_loader gluon.data.DataLoader(train_dataset.transform_first(training_transformer), batch_sizebatch_size, shuffleTrue, try_nopythonTrue) validation_loader gluon.data.DataLoader(val_dataset.transform_first(validation_transformer), batch_sizebatch_size, try_nopythonTrue) test_loader gluon.data.DataLoader(test_dataset.transform_first(validation_transformer), batch_sizebatch_size, try_nopythonTrue)注意batch_size 4与2 块 GPU的配合每个批次会被平均切成 2 份每块 GPU 拿到 2 个样本。这与split_and_load的约束直接相关见下文 8.2 的源码说明。8.2 定义多设备评测辅助函数test函数与第 6 步基本相同区别在于使用了gluon.utils.split_and_load把每个批次切分并加载到多块设备上def test(val_data, devices): acc gluon.metric.Accuracy() for batch in val_data: data, label batch[0], batch[1] data_list gluon.utils.split_and_load(data, devices) label_list gluon.utils.split_and_load(label, devices) outputs [net(X) for X in data_list] acc.update(label_list, outputs) _, accuracy acc.get() return accuracysplit_and_load的底层语义该函数定义于 python/mxnet/gluon/utils.py#L86-L113核心流程是若输入不是 ndarray先把它转成ctx_list[0]上的 ndarray若ctx_list长度为 1直接as_in_context放到该设备并返回否则沿batch_axis默认 0即批次维把数据切分成len(ctx_list)份再逐份as_in_context到对应设备返回一个列表。其前置校验逻辑python/mxnet/gluon/utils.py#L64-L69值得注意当even_splitTrue默认值时批次大小必须能被 GPU 数量整除否则抛出ValueError提示使用 GPU 数量的整数倍的 batch size或设置even_splitFalse允许不均匀切分。这正是教程中batch_size4与 2 卡搭配的原因。若采用不均匀切分源码会通过divmod把多余样本逐个分配给靠前的切片python/mxnet/gluon/utils.py#L71-L73但会牺牲负载均衡。8.3 多卡训练主循环训练循环与单卡版本大体相同关键差异用注释Diff 1~Diff 5标出# Diff 1: 使用两块 GPU 训练 available_gpus [npx.gpu(i) for i in range(npx.num_gpus())] num_gpus 2 devices available_gpus[:num_gpus] print(Using {} GPUs.format(len(devices))) # Diff 2: 重新初始化参数并放到多块 GPU 上 net.initialize(force_reinitTrue, devicedevices) # 损失函数与优化器与之前一致 loss_fn gluon.loss.SoftmaxCrossEntropyLoss() optimizer sgd optimizer_params {learning_rate: 0.001} trainer gluon.Trainer(net.collect_params(), optimizer, optimizer_params) epochs 2 accuracy gluon.metric.Accuracy() log_interval 5 for epoch in range(epochs): train_loss 0. tic time.time() btic time.time() accuracy.reset() for idx, batch in enumerate(train_loader): data, label batch[0], batch[1] # Diff 3: 把批次切分并加载到对应设备 data_list gluon.utils.split_and_load(data, devices) label_list gluon.utils.split_and_load(label, devices) # Diff 4: 在每个设备上分别做前向与反向 # MXNet 会自动并行执行这些计算 with autograd.record(): outputs [net(X) for X in data_list] losses [loss_fn(output, label) for output, label in zip(outputs, label_list)] for l in losses: l.backward() trainer.step(batch_size) # Diff 5: 汇总所有设备上的损失float() 会把数据拷回 CPU train_loss sum([float(l.sum()) for l in losses]) accuracy.update(label_list, outputs) if log_interval and (idx 1) % log_interval 0: _, acc accuracy.get() print(fEpoch[{epoch 1}] Batch[{idx 1}] Speed: {batch_size / (time.time() - btic)} samples/sec \ batch loss {train_loss} | accuracy {acc}) btic time.time() _, acc accuracy.get() acc_val test(validation_loader, devices) print(f[Epoch {epoch 1}] training: accuracy{acc}) print(f[Epoch {epoch 1}] time cost: {time.time() - tic}) print(f[Epoch {epoch 1}] validation: validation accuracy{acc_val})逐条解读这 5 处差异Diff 1先列出机器上全部 GPU[npx.gpu(i) for i in range(npx.num_gpus())]再取前 2 块组成devices列表。这样做的好处是代码对 GPU 数量具有自适应能力同时便于扩展到更多卡。Diff 2net.initialize(force_reinitTrue, devicedevices)。由于此前load_parameters已经把参数放到单块 GPU 上这里用force_reinitTrue强制重新初始化device传入设备列表后initializepython/mxnet/gluon/block.py#L557-L577会在每一块设备上各保留一份参数副本供各卡独立计算梯度。Diff 3split_and_load(data, devices)把 batch 沿批次维平均切成 2 份并分别as_in_context到两块 GPU返回data_listlabel 同理。Diff 4在autograd.record()上下文中对每个切块分别做前向与损失计算再逐个backward()。由于各卡数据相互独立MXNet 的引擎会自动把不同设备上的计算并行调度执行这也是 MXNet 依赖调度引擎设计带来的并发收益。Diff 5trainer.step(batch_size)使用完整批次大小而非单卡切块大小更新参数保证学习率语义与单卡一致float(l.sum())会把 GPU 上的损失值同步拷回 CPU 主存再求和这正是前文提到的asnumpy/同步拷贝机制在训练中的典型应用随后统一更新Accuracy指标。训练结束后test(validation_loader, devices)用同一套多设备切分逻辑在验证集上评测你也可以把test_loader传入test函数评估测试集表现。9. 使用 GPU 的常见注意事项综合上述教程与源码使用 GPU 时有几个容易踩坑的点场景注意事项设备一致性参与同一运算的多个输入必须位于同一 GPU跨设备运算会报错需用copyto显式搬运数据并行切分split_and_load默认要求 batch size 能被 GPU 数量整除even_splitTrue否则抛ValueError可设even_splitFalse允许不均匀切分但会破坏负载均衡参数初始化多卡训练时net.initialize(devicedevices)会为每块卡保留参数副本重复初始化需force_reinitTrue设备移动开销asnumpy()、print等操作会触发 GPU→CPU 的同步拷贝训练热循环中应避免频繁调用以免阻塞流水线显存查询训练前可用npx.gpu_memory_info(device_id)见 python/mxnet/device.py#L249-L269检查各卡可用显存避免 OOM无卡回退用npx.gpu() if npx.num_gpus() 0 else npx.cpu()这类写法可让同一份代码在 CPU/GPU 环境都安全运行10. 总结与后续学习至此你已经完成了 MXNet 快速入门系列的全部内容从创建网络、自动求导、数据加载到单卡训练再到本教程的 GPU 加速与多卡数据并行训练。核心要点可以归纳为一句话在 MXNet 中启用 GPU 加速的关键就是设备管理——用device参数指定数据与参数的位置用copyto显式完成跨设备迁移用split_and_load 多设备initialize实现数据并行训练。这套设备抽象贯穿 ndarraypython/mxnet/ndarray/ndarray.py、Gluon 块python/mxnet/gluon/block.py与工具函数python/mxnet/gluon/utils.py三层接口理解它之后无论是单卡推理还是多卡大规模训练你都能准确地控制计算发生在哪一块硬件上。如果希望继续深入可以探索 MXNet 生态中的更上层工具如 GluonCV、GluonNLP、GluonTS、AutoGluon 等在各自领域对 GPU 与分布式训练的进一步封装也可以在 快速入门系列目录 中回顾前序步骤巩固整套工作流。【免费下载链接】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），仅供参考