资讯动态

MXNet GPU 使用指南:在 GPU 上加载、运行与多卡并行训练神经网络

发布时间:2026/9/20 19:28:46 来源:尧图企业网站定制
深度学习人工智能机器学习分布式训练【免费下载链接】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点击查看免费下载本教程是 MXNet 快速入门系列crash course的第 7 步主题是如何在 NVIDIA GPU 上加载并运行神经网络从检查可用 GPU、把数据显式分配到指定 GPU到把已保存的网络参数加载到 GPU 上做推理再到利用数据并行data parallelism在多个 GPU 上联合训练网络。读完本文你将掌握npx.gpu()/npx.num_gpus()的设备管理 API、device参数与copyto的显式数据搬运机制以及基于gluon.utils.split_and_load的多卡训练完整套路。前置条件在开始之前请确认你的机器满足以下条件至少一块 NVIDIA GPUMXNet 目前只支持 NVIDIA GPUAMD 与 Intel 的 GPU 不受支持正确安装 CUDA需要与 MXNet 编译时使用的 CUDA 版本匹配的驱动与运行时环境安装 GPU 版 MXNet必须安装启用了 CUDA 支持的 MXNet 构建版本。若需要从源码自行编译 GPU 版本可以参考仓库 config/distribution 目录下的 CUDA 编译配置例如linux_cu92.cmake、linux_cu100.cmake、linux_cu110.cmake等它们分别对应不同 CUDA 版本的 CMake 构建选项。满足上述条件后即可在 Python 中开始使用 GPU。首先导入相关模块并开启 NumPy 语义from mxnet import np, npx, gluon, autograd from mxnet.gluon import nn import time npx.set_np()npx.set_np()将 MXNet 切换到与 NumPy 兼容的算子与数组语义其实现位于 python/mxnet/util.py这也是本教程所有示例使用np/npx模块的前提。检查 MXNet 可访问的 GPU 数量训练或推理前先确认 MXNet 能看到多少块 GPUnpx.num_gpus() # 返回 MXNet 可以访问的 GPU 数量该命令直接查询 CUDA 运行时其底层实现在 python/mxnet/device.py 的num_gpus()函数中它通过 CUDA 查询当前机器上的 GPU 数量任何 CUDA 错误都会抛出异常。仓库的 GPU 测试代码也大量依赖此函数做环境探测与用例跳过判断例如 tests/python/gpu/test_device.py 会用mx.device.num_gpus()探测可用 GPU 并限制测试规模tests/python/unittest/test_operator.py 中的test_device_num_gpus则验证该接口在 CPU/GPU 环境下均能正常返回非负值。一个稳妥的惯例是优先使用 GPU没有 GPU 时回退到 CPUgpu npx.gpu() if npx.num_gpus() 0 else npx.cpu()把数据分配到 GPUMXNet 的 ndarray 与 NumPy 数组非常相似但有一个关键区别MXNet ndarray 带有device属性标明该数组存放在哪个设备上。默认情况下数组被创建在npx.cpu()上。要创建一块位于第一块 GPU 上的数组只需在创建时显式传入device参数gpu npx.gpu() if npx.num_gpus() 0 else npx.cpu() x np.ones((3,4), devicegpu) x其中npx.gpu()与npx.gpu(0)等价都表示第一块 GPU。如果使用 CPUMXNet 会把数据分配在主内存中并尽可能利用所有 CPU 核心进行计算如果机器上有多块 GPUMXNet 会明确告诉你每个 ndarray 分配在哪块 GPU 上。多 GPU 场景下的显式搬运假设机器上至少有两块 GPU你可以把x复制到第二块 GPUnpx.gpu(1)gpu_1 npx.gpu(1) if npx.num_gpus() 1 else npx.cpu() x.copyto(gpu_1)注意如果机器上只有一块 GPU运行这段代码会直接报错。MXNet 要求用户显式地在设备之间移动数据——不存在自动跨设备搬运。不过少数算子如print、asnumpy会隐式地把数据搬回主内存方便调试与查看。gpu()与cpu()都是 python/mxnet/device.py 中定义的便捷工厂函数内部统一构造Device(gpu, device_id)/Device(cpu, device_id)。Device类见 python/mxnet/device.py维护了设备类型与 ID 的映射{1: cpu, 2: gpu, 3: cpu_pinned, 5: cpu_shared}device_id默认值为 0。除了gpu/cpuMXNet 还提供cpu_pinned()页锁定内存向 GPU 拷贝更快等设备类型可查看同一文件了解。选择 GPU ID如果机器上有多块 GPUMXNet 通过0 起始索引0-indexing访问每一块npx.gpu(0)是第一块npx.gpu(1)是第二块以此类推。如果机器有 8 块 GPU最后一块是npx.gpu(7)。精确指定 GPU ID 意味着你可以自主决定让哪些运算和训练跑在哪块卡上——这在希望利用多卡做并行训练时尤其有用。同时python/mxnet/device.py 还支持with mx.gpu(2):这样的上下文管理器用法可以在代码块内临时切换默认设备。在 GPU 上运行运算要在某块 GPU 上执行运算只需要保证该运算的输入已经位于这块 GPU 上输出就会自动分配在同一块 GPU 上。np与npx模块中的几乎所有算子都支持在 GPU 上运行y np.random.uniform(size(3,4), devicegpu) x y这里x与y都创建在gpu第一块 GPU上因此加法运算直接在 GPU 上完成。如果参与运算的输入不在同一块 GPU 上你会得到报错——这也是 MXNet 显式设备管理风格的体现。在 GPU 上运行神经网络要在 GPU 上运行一个神经网络核心只有一步把输入数据和网络参数都搬到 GPU 上。下面复用快速入门第 6 步 Training Neural Networks 中定义的LeafNetwork卷积块 全连接块的混合序贯网络展示完整流程# 卷积块一个卷积层 一个最大池化层 一个批归一化层 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 batch实例化网络后把之前保存的参数例如leaf_models.params直接加载到 GPU 0上net LeafNetwork() net.load_parameters(leaf_models.params, devicegpu)load_parameters的完整签名见 python/mxnet/gluon/block.pydevice参数默认是cpu()可传单个Device或Device列表它还支持allow_missing是否静默跳过文件中缺失的参数、ignore_extra是否忽略文件中多余的参数、cast_dtype是否按当前 dtype 强转等选项用于模型迁移与精调场景。除了加载参数时指定device你还可以在参数已初始化后通过net.collect_params().reset_device(gpu)把整个网络的参数重新分配到 GPU 上。reset_device的实现见 python/mxnet/gluon/block.py它收集网络全部参数并逐个重新分配设备若传入Device列表则会在每块设备上都保留一份参数拷贝——这正是多卡训练所需的参数布局。最后在 GPU 0 上创建输入数据并前向传播整个forward就会在 GPU 0 上执行x np.random.uniform(size(1, 3, 128, 128), devicegpu) net(x)多 GPU 数据并行训练最后一部分如何使用多块 GPU 以**数据并行data parallelism**的方式联合训练网络。数据并行的思想很直观假设有n块 GPU就把每个数据批次按 batch 维度切成n份每块 GPU 负责其中一份的前向与反向传播从而摊薄单个批次的训练时间。准备数据先从第 6 步教程 Training Neural Networks 中复制数据定义与变换函数# 导入 transforms把一系列图像变换组合起来 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)定义验证辅助函数这与第 6 步中定义的test函数一致注意它通过gluon.utils.split_and_load把每个 batch 切分并加载到多个设备上# 返回验证集与测试集准确率的函数 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它先把数据沿batch_axis默认 0按设备数量切分even_splitTrue时要求 batch 大小能被设备数整除否则抛出提示性错误再把每个切片通过as_in_context搬运到对应设备上当设备列表只有一个元素时直接原地迁移而不切分。多 GPU 训练循环训练循环与单卡版本大体相同差异集中在以下 5 处代码中已用Diff标注# 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: 切分 batch 并加载到对应的设备 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 处关键差异选择设备用列表推导式[npx.gpu(i) for i in range(npx.num_gpus())]枚举所有 GPU再取前num_gpus此处为 2块作为训练设备参数初始化到多设备net.initialize(force_reinitTrue, devicedevices)中device传入设备列表MXNet 会在每块 GPU 上各保留一份参数副本。initialize的签名见 python/mxnet/gluon/block.pyforce_reinitTrue用于强制重新初始化覆盖此前已加载或初始化过的参数device同样支持单设备或设备列表切分数据gluon.utils.split_and_load(data, devices)把一个 batch 沿 batch 维切成与设备数相同的份数并逐份搬运到对应 GPU并行前反向对每份数据各自执行net(X)与loss_fnautograd.record()记录计算图随后逐个l.backward()。各 GPU 上的计算由 MXNet 引擎自动调度并行执行无需手工管理线程损失汇总与日志float(l.sum())会把 GPU 上的张量拷贝回 CPU 得到标量这正是前文提到的隐式搬运算子累加得到整批损失accuracy.update(label_list, outputs)汇总各设备的预测结果计算整体准确率。每轮 epoch 结束后用test(validation_loader, devices)在多设备上完成验证并输出训练准确率、验证准确率与耗时。这套「split_and_load 参数复制 各设备独立前反向 汇总损失」的模式在仓库测试中被广泛验证例如 tests/python/gpu/test_gluon_gpu.py 直接在 GPU 设备列表上使用split_and_load验证多卡数据切分与加载的正确性tests/python/unittest/test_gluon.py 也在设备列表上对split_and_load的切分行为做了断言。实际应用时请确保batch_size能被 GPU 数量整除或在必要时设置even_splitFalse否则split_and_load会抛出「数据无法均匀切分」的 ValueError。小结与延伸阅读至此你已走完 MXNet 快速入门系列的全部 7 步掌握了在 GPU 上训练与推理神经网络的核心技能用npx.num_gpus()探测 GPU 数量用npx.gpu(i)按 0 起始索引选择设备通过device参数或copyto在设备间显式搬运 ndarray并牢记「输入在哪个设备运算就在哪个设备」用net.load_parameters(..., devicegpu)或net.collect_params().reset_device(gpu)把网络参数放到 GPU 上用split_and_loadinitialize(force_reinitTrue, devicedevices)实现多 GPU 数据并行训练并由 MXNet 引擎自动并行调度各设备上的前向与反向计算。如果你想继续深入学习可以阅读本系列前面的步骤例如 Training Neural Networks以及仓库中 python/mxnet/device.py 与 python/mxnet/gluon/utils.py 的相关实现进一步理解设备管理与多卡数据流背后的机制。赞分享深度学习人工智能机器学习分布式训练【免费下载链接】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点击查看免费下载相关推荐MXNet GPU 使用指南从 NDArray 分配到多 GPU 数据并行训练MXNet GPU 使用指南从 NDArray 分配到多 GPU 数据并行训练 本教程面向希望用 GPU 加速 MXNet 训练与推理的开发者系统讲解 MX深度学习机器学习人工智能突破图神经网络训练瓶颈PyTorch Geometric多GPU并行计算实战指南突破图神经网络训练瓶颈PyTorch Geometric多GPU并行计算实战指南 在处理大规模图数据时单GPU训练常受限于内存容量和计算速度。本文将系统介绍人工智能机器学习深度学习图计算使用 MShadow 在 CPU/GPU 上训练神经网络MNIST 全连接网络、卷积网络与多 GPU 参数服务器实战指南使用 MShadow 在 CPU/GPU 上训练神经网络MNIST 全连接网络、卷积网络与多 GPU 参数服务器实战指南 导读 本文基于 MShadow 官方深度学习机器学习人工智能创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价