十分钟从 PyTorch 转 MXNet

十分钟从 PyTorch 转 MXNet 的关键点总结如下：

1. 安装方式

   PyTorch：默认使用 conda 安装，例如 conda install pytorch-cpu -c pytorch。

   MXNet：更常用 pip 安装，例如 pip install --pre mxnet。

2. 多维矩阵

   PyTorch 使用 tensor，MXNet 使用 ndarray。

   示例代码：

   PyTorch：

   import torchx = torch.ones(5,3)y = x + 1print(y)

   MXNet：

   from mxnet import ndx = nd.ones((5,3))y = x + 1print(y)

   主要区别：MXNet 的形状参数需要用括号括起来。

3. 模型训练

   读取数据：

   PyTorch：

   import torchfrom torchvision import datasets, transformstrain_data = torch.utils.data.DataLoader( datasets.MNIST(train=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.13,), (0.31,))])), batch_size=128, shuffle=True, num_workers=4)

   MXNet：

   from mxnet import gluonfrom mxnet.gluon.data.vision import datasets, transformstrain_data = gluon.data.DataLoader( datasets.MNIST(train=True).transform_first(transforms.Compose([ transforms.ToTensor(), transforms.Normalize(0.13, 0.31)])), batch_size=128, shuffle=True, num_workers=4)

   主要区别：MXNet 使用 transform_first 来表明数据变化是作用在读到的批量的第一个元素。

   定义模型：

   PyTorch：

   from torch import nnnet = nn.Sequential( nn.Linear(28*28, 256), nn.ReLU(), nn.Linear(256, 10))

   MXNet：

   from mxnet.gluon import nnnet = nn.Sequential()with net.name_scope(): net.add( nn.Dense(256, activation='relu'), nn.Dense(10) )net.initialize()

   主要区别：

   MXNet 不需要指定输入大小，系统会自动推理。

   MXNet 可以指定激活函数。

   MXNet 需要创建一个 name_scope 的域来给每一层附上一个独一无二的名字。

   MXNet 需要显示调用模型初始化函数。

   损失函数和优化算法：

   PyTorch：

   loss_fn = nn.CrossEntropyLoss()trainer = torch.optim.SGD(net.parameters(), lr=0.1)

   MXNet：

   loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.1})

   训练：

   PyTorch：

   from time import timefor epoch in range(5): total_loss = .0 tic = time() for X, y in train_data: X, y = torch.autograd.Variable(X), torch.autograd.Variable(y) trainer.zero_grad() loss = loss_fn(net(X.view(-1, 28*28)), y) loss.backward() trainer.step() total_loss += loss.mean() print('epoch %d, avg loss %.4f, time %.2f' % ( epoch, total_loss/len(train_data), time()-tic))

   MXNet：

   from time import timefor epoch in range(5): total_loss = .0 tic = time() for X, y in train_data: with mx.autograd.record(): loss = loss_fn(net(X.flatten()), y) loss.backward() trainer.step(batch_size=128) total_loss += loss.mean().asscalar() print('epoch %d, avg loss %.4f, time %.2f' % ( epoch, total_loss/len(train_data), time()-tic))

   主要区别：

   MXNet 不需要将输入放进 Variable，但需要将计算放在 mx.autograd.record() 里。

   MXNet 不需要每次梯度清 0，因为新梯度是写进去，而不是累加。

   MXNet 的 step 需要给定批量大小。

   MXNet 需要调用 asscalar() 来将多维数组变成标量。

4. 性能比较

   在这个样例中，MXNet 比 PyTorch 快两倍，但这样的比较需要谨慎对待。

5. 进一步学习

   更详细的 MXNet 教程欢迎留言。

总结：从 PyTorch 转到 MXNet 需要注意安装方式、多维矩阵的创建、数据读取、模型定义、损失函数和优化算法的使用以及训练过程中的一些细节差异。