Export BatchNormalizationLayer as MXNet

Question

Bug introduced in 11.2 and remains in 12.0.0

I just tried to import a net from mathemtica into pythons using MXNet. It works fine with the following manual.

However, as soon as I try to add a BatchNormalizationLayer to my net MXNets bind() function fails because the auxiliary states of the BatchNormlayer are not provided (MovingMean,MovingVar).

Mathematica code:

Export["~/batchnorm.json", foo = NetInitialize@BatchNormalizationLayer
      ["Input"-> {1, 2, 2}], "MXNet"]

Python code:

import mxnet as mx
import numpy as np
sym=mx.symbol.load('batchnorm.json')
nd=mx.nd.load('batchnorm.params')
#parse Input
inputND = mx.nd.array(np.array([[[[1,1],[2,2]]]]));
nd["Input"] = inputND
#bind NN
e = sym.bind(mx.cpu(), nd)

Above code generates the following error:

ValueError: Length of aux_states does not match the number of arguments

There appears to be no trace of the auxiliary layers in the params file and hence not in my nd variable. Am I missing something?

Best, Max

EDIT: The only workaround seems to be manual exporting of the auxiliary parameters. Which works. It would be nice if Mathematica would generate this file automatically.

Answer 1

This is fixed in Mathematica 12.0:

In[287]:= Export["~/batchnorm.json", 
foo = NetInitialize@BatchNormalizationLayer["Input" -> {1, 2, 2}], "MXNet"];
MXNetLink`NDArrayImport@"~/batchnorm.params"

Out[288]=
Association["Scaling" -> NumericArray[{1.}, "Real32"], 
 "Biases" -> NumericArray[{0.}, "Real32"], 
 "MovingMean" -> NumericArray[{0.}, "Real32"], 
 "MovingVariance" -> NumericArray[{1.}, "Real32"]]

Answer 2

I found the following workaround for NetChains with several Batchnorm layers by manually exporting aux states:

Mathematica

GetAux[layer_BatchNormalizationLayer] := {NetExtract[layer, 
      "MovingMean"], NetExtract[layer, "MovingVariance"]}
net = NetInitialize@NetChain@
     {BatchNormalizationLayer["Input" -> {2, 2, 2}], BatchNormalizationLayer["Input" -> {2, 2, 2}]}
aux = Flatten[GetAux /@ NetExtract[net, All], 1]
Export[FileNameJoin@{NotebookDirectory[],"batchnorm.json"}, net, "MXNet"]
Export[FileNameJoin@{NotebookDirectory[], "auxparam.txt"},
      ExportString[aux, "PythonExpression"]]

Then you can call the following lines of python code to import the aux states and execute the net:

import mxnet as mx
import numpy as np
sym=mx.symbol.load('batchnorm.json')
nd=mx.nd.load('batchnorm.params')

#get aux states from file
foo=np.loadtxt('auxparam.txt',dtype='str')
batchaux=eval(foo.item())
aux = {}
naux=sym.list_auxiliary_states()
for i in range(len(naux)):
    aux[naux[i]] = mx.nd.array(np.array(batchaux[i]))

#parse Input
inputND = mx.nd.array(np.array([[[[1,1],[2,2]]]]));
nd["Input"] = inputND

#run the net on input
e = sym.bind(mx.cpu(), nd, aux_states=aux);
out = e.forward()