MXNet 深度学习框架

FreeGuideOnline 7阅读 2026-07-10

bash pip install mxnet


若需要 GPU 支持(CUDA 11.x),请执行:

```bash
pip install mxnet-cu112

2.2 验证安装

在 Python 中导入并查看版本:

import mxnet as mx
print(mx.__version__)

若输出形如 1.9.1 则表示安装成功。对于 GPU 版本,可进一步检查是否识别到 GPU:

mx.context.num_gpus()  # 返回可用的 GPU 数量

3. 基础数据操作:NDArray

NDArray 是 MXNet 的核心数据结构,类似于 NumPy 的 ndarray,但支持自动微分和 GPU 计算。

from mxnet import nd

# 创建 NDArray
a = nd.array([1, 2, 3])
b = nd.ones((2, 3))
c = nd.random.normal(0, 1, shape=(3, 4))

# 基本运算
sum_a = a.sum()
prod_b = nd.dot(b, c.T)

# 索引与切片
print(b[0, 1])       # 取第一行第二列
print(b[:, 1:3])     # 切片操作

# 上下文切换:将数据放到 GPU(0)
a_gpu = a.as_in_context(mx.gpu(0))

4. 自动微分:autograd

MXNet 通过 autograd 包自动计算梯度,是训练神经网络的基础。

from mxnet import autograd, nd

x = nd.array([1, 2, 3])
x.attach_grad()               # 申请存储梯度的内存

with autograd.record():
    y = 2 * x.dot(x)          # y = 2 * (x·x)

y.backward()                  # 反向传播计算梯度
print(x.grad)                 # 输出 [4. 8. 12.] (即 dy/dx = 4x)
  • attach_grad():为变量附加梯度缓冲区
  • autograd.record():记录前向计算,以便后续 backward() 求导
  • 默认情况下,record 作用域内的 NDArray 会追踪计算图

5. 使用 Gluon 构建神经网络

Gluon 是 MXNet 的高级 API,提供类似 PyTorch 的动态图体验,同时支持轻松导出为符号图。

5.1 定义模型

from mxnet.gluon import nn

net = nn.Sequential()
net.add(nn.Dense(128, activation='relu'),
        nn.Dense(64, activation='relu'),
        nn.Dense(10))        # 10 分类输出

也可以使用更灵活的子类化方式:

class MLP(nn.Block):
    def __init__(self, **kwargs):
        super(MLP, self).__init__(**kwargs)
        self.dense1 = nn.Dense(64, activation='relu')
        self.dense2 = nn.Dense(10)

    def forward(self, x):
        return self.dense2(self.dense1(x))

net = MLP()

5.2 初始化模型参数

net.initialize(init=init.Xavier())   # Xavier 初始化器

初次执行前向传播时,形状会自动推断:

# 给定一个批量大小为 4,特征维度为 20 的输入
input_data = nd.random.uniform(shape=(4, 20))
net(input_data)

6. 模型训练基本流程

6.1 准备数据

使用 gluon.data 加载数据集,此处以虚拟数据为例:

from mxnet.gluon import data as gdata

# 生成随机数据集
num_samples, num_features = 1000, 20
X = nd.random.normal(shape=(num_samples, num_features))
y = nd.random.randint(0, 10, shape=(num_samples,))

dataset = gdata.ArrayDataset(X, y)
dataloader = gdata.DataLoader(dataset, batch_size=32, shuffle=True)

6.2 定义损失函数和优化器

from mxnet.gluon import loss as gloss, Trainer
from mxnet import init

loss_fn = gloss.SoftmaxCELoss()                # 交叉熵损失
trainer = Trainer(net.collect_params(), 'adam', {'learning_rate': 0.01})

6.3 训练循环

epochs = 5
for epoch in range(epochs):
    total_loss = 0
    for batch_x, batch_y in dataloader:
        with autograd.record():
            outputs = net(batch_x)
            l = loss_fn(outputs, batch_y)
        l.backward()
        trainer.step(batch_size=32)            # 更新参数
        total_loss += l.sum().asscalar()
    print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataset):.4f}')

7. 模型保存与加载

MXNet 可以灵活地保存参数或整个模型(仅 Gluon 图)。推荐保存参数,因为占用空间小且便于迁移。

# 保存参数
net.save_parameters('mlp.params')

# 加载参数(需重新定义相同结构的网络)
new_net = MLP()
new_net.load_parameters('mlp.params')

8. 使用 GPU 加速

MXNet 的上下文(context)机制让切换设备非常容易。

# 将模型参数放置到 GPU
net.initialize(ctx=mx.gpu(0))

# 将数据放置到 GPU
batch_x_gpu = batch_x.as_in_context(mx.gpu(0))
# 前向传播自动在 GPU 进行
output = net(batch_x_gpu)