How to export an MXNet?

Question

I was hoping there was some way to generate the .params and .json file needed to define an MXNet model, from a network trained using NetTrain[] in Mathematica. I was hopeful because I found these functions in the NeuralNetworks` package:

NeuralNetworks`ToMXNetSymbol
NeuralNetworks`ToMXNetJSON

You can use them on a net link this:

<< NeuralNetworks`;
net = NetInitialize@NetChain[{
     ConvolutionLayer[20, {5, 5}],
     ElementwiseLayer[Ramp],
     PoolingLayer[{2, 2}, {2, 2}],
     FlattenLayer[],
     DotPlusLayer[500],
     ElementwiseLayer[Ramp]
     }, "Input" -> NetEncoder[{"Image", {32, 32}}]];
ImportString[Normal@NeuralNetworks`ToMXNetJSON[net][[1]], "RawJSON"]
NeuralNetworks`ToMXNetSymbol[net]

Now, ToMXNetJSON[] returns a tuple, and it looks like the first element is the JSON for the symbol file. But I don't know what the second element is, and I don't have a clue as to what ToMXNetSymbol[] is returning.

Motivation: A solution to this would enable one to take any net from MMA and run it in C or Python with GPU inference!

Answer 1

It seems the model file in MXNet (checkpoint) is defined by two files: a ".json" file and a ".params" file. The json file contains the definition of the network, and the params file contains the actual weight and bias of each neuron. The params file is in the binary format of MXNet's NDArray representation.

Thus, to export a network in Mathematica to MXNet, we need to generate these two files. The json file can be generated easily with the NeuralNetworks``ToMXNetJSON. The param file can be generated using the MXNetLink``NDArrayExport. Here is an example of this process using MNIST example in the documentation.

We first load the packages

<< MXNetLink`;
<< NeuralNetworks`;
<< GeneralUtilities`;

and define the network

net = NetChain[{
   ConvolutionLayer[20, {5, 5}],
   ElementwiseLayer[Ramp],
   PoolingLayer[{2, 2}, {2, 2}],
   ConvolutionLayer[50, {5, 5}],
   ElementwiseLayer[Ramp],
   PoolingLayer[{2, 2}, {2, 2}],
   FlattenLayer[],
   DotPlusLayer[500],
   ElementwiseLayer[Ramp],
   DotPlusLayer[10],
   SoftmaxLayer[]},
  "Output" -> NetDecoder[{"Class", Range[0, 9]}],
  "Input" -> NetEncoder[{"Image", {28, 28}, "Grayscale"}]
  ]

Then we train the network

resource = ResourceObject["MNIST"];
trainingData = ResourceData[resource, "TrainingData"];
testData = ResourceData[resource, "TestData"];
trained = 
 NetTrain[net, trainingData, ValidationSet -> testData, 
  MaxTrainingRounds -> 3]

We will now export the trained network into the MXNet's model files.

The json file can be exported using ToMXNetJSON

jsonPath = "~/Downloads/MNIST-symbol.json";
Export[jsonPath, ToMXNetJSON[trained][[1]], "String"]
(* "~/Downloads/MNIST-symbol.json" *)

The second part of ToMXNetJSON[trained] contains the weights of our network. However, the weights are in Mathematica's RawArray format, so we need to convert those into MXNet's NDArray format. Also we will drop the encoder layer, and change the names of the layers to comply with the convention

paraPath = "~/Downloads/MNIST-0000.params";
ass = KeyDrop[ToMXNetJSON[trained][[2]], ".Inputs.Input"];

f[str_] := 
 If[StringFreeQ[str, "Arrays"], str, 
  StringReplace[
   StringSplit[str, ".Arrays."] /. {a_, b_} :> 
     StringJoin[{"arg:", a, "_", b}], {"Weights" -> "weight", 
    "Biases" -> "bias"}]]

NDArrayExport[paraPath, NDArrayCreate /@ KeyMap[f, ass]]
(* "~/Downloads/MNIST-0000.params" *)

Now the two files "MNIST-symbol.json" and "MNIST-0000.params" can be used to load the network in MXNet.

To verify that the files are correct, we can use ImportMXNetModel to import MXNet model files we just generated.

trained2 = 
 NetGraph[{ImportMXNetModel[jsonPath, paraPath]}, {}, 
  "Input" -> NetEncoder[{"Image", {28, 28}, "Grayscale"}], 
  "Output" -> NetDecoder[{"Class", Range[0, 9]}]]

And we see that the network from the MXNet's model files produces the same predictions as our original network:

testsample = RandomSample[testData, 100];

(trained[Keys[#], "Probabilities"] & /@ 
   testsample) == (trained2[Keys[#], "Probabilities"] & /@ testsample)
(* True *)

Edit for version 11.1

It seems that the structure of the implementation of neural network has been updated in 11.1. The trained weight is no longer in the NeuralNetworks``ToMXNetJSON. The trained weight can be accessed using NeuralNetworks``ToNetPlan.

plan = ToNetPlan[trained]

So to export the weight, we can do

NDArrayExport[paraPath, NDArrayCreate /@ KeyMap[f, plan["WeightArrays"]]]