Export BatchNormalizationLayer as MXNet
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()
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"]]