Current and Future Backends for the Neural Net Framework
Matteo Salvarezza, Wolfram Research

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

ABSTRACT & OUTLINE

This talk will explore upcoming changes in the Neural Net Framework’s backend. We’ll discuss enhanced inference capabilities introduced in Version 14.1, which improve model performance on various devices via just-in-time ONNX compilation. Attendees will see a demonstration of the new backend and learn about future developments.

Backends & the Neural Net framework

Backends & NetExternalObject

Backends & accelerated hardware

Future backends

BACKENDS & THE NEURAL NET FRAMEWORK

The neural network framework relies on external code (backends) to manage net evaluation and training

The backend is NOT USED for operation involving creating and manipulating net expressions and to run NetEncoder / NetDecoders:

Net construction:

net=NetInitialize@NetGraph[
{LinearLayer[2],LinearLayer[4],CatenateLayer[]},
{NetPort["Input1"]->1,NetPort["Input2"]->2,{1,2}->3},
"Input1"->{5},"Input2"->{6}
]

Extraction:

NetExtract[net,1]

Surgery:

NetAppend[net,ElementwiseLayer[Sin]]

NetEncoder/NetDecoder creation and evaluation:

enc=NetEncoder[{"Image",{4,4}}]

enc@RandomImage[]

The backend is USED for net evaluation and training

Forward pass evaluation

net@<|"Input1"->{1,2,3,4,5},"Input2"->{1,2,3,4,5,6}|>

Backward pass evaluation

net[<|"Input1"->{1,2,3,4,5},"Input2"->{1,2,3,4,5,6}|>,NetPortGradient[All]]

Training

NetTrain[net,<|"Input1"->RandomReal[1,{10,5}],"Input2"->RandomReal[1,{10,6}],"Output"->RandomReal[1,{10,6}]|>,MaxTrainingRounds->10000]

When the backend needs to be used the internal expression of the net is translated to a form that the backend can understand and passed to it so that it can be evaluated

NeuralNetworks`NData[net]

BACKENDS & NET EXTERNAL OBJECT

NetExternalObject is a thin wrapper over a net which is managed directly by the backend

They can only be constructed from a file of a format associated to the backend

The backend is also declared explicitly in the object

Import["ExampleData/onnx_example.onnx","NetExternalObject"]

Import["ExampleData/mxnet_example1.json",{"MXNet","NetExternalObject"},InputPorts-><|"Input"->{4,1}|>]

External objects are treated like black boxes so there’s no access to its internal structure for surgery, extraction, etc

Running them used the associated backend:

net=Import["ExampleData/onnx_example.onnx","NetExternalObject"]

net[{{1,2,3},{4,5,6}}]

NetExternalObject doesn’t support training and gradient computation, it’s meant only to use external pre-trained nets

In this case there is no translation from our representation to the backend, the object only exists in the backend format

BACKENDS & ACCELERATED HARDWARE

Two backends are in use at the moment, MXNet and ONNXRuntime, and both the NN framework and NetExternalObject can use them

But ONNXRuntime can’t be used for training ad gradient computation

NetExternalObject is explicitly tied to a backend, while the NN framework selects which backend to use depending on the TargetDevice setting:

On Apple Silicon systems, TargetDevice -> "CoreML" uses ONNXRuntime backend to access the Apple Neural Engine

On Windows systems, TargetDevice -> "DirectML" uses ONNXRuntime backend to access the integrated GPU or an Nvidia GPU

The Nvidia GPU can be used both via CUDA and DirectML, although CUDA is recommended

In all other cases, MXNet is used

In addition, TargetDevice -> "GPU" is now a meta-setting that translates to either "CUDA"/"DirectML"/"CoreML" depending on the platform and available hardware

See the documentation page of TargetDevice for detailed information

This picture will get simpler and more transparent in the future as new backends will be added

FUTURE BACKENDS

Adding backends to NetExternalObject is easy

Just pass the imported model file to the backend and ask it to run the net

Possible backends added in the future include llama.cpp (LLM evaluation) and Tensorflow

Adding backends to the NN framework is hard

Need to translate every layer to the backend, manage gradient computation, optimizers for training, ...

We are forced to add a new backend to the NN framework because MXNet was retired and ONNXRuntime is not (yet) a fully-fledged and established deep learning framework

MXNet will eventually be phased out and removed from the system, and the NN framework will be migrated to use one or more different backends

The main candidates are obviously PyTorch and Tensorflow

The first goal of introducing a new backend is feature-parity with what the framework can do now, aiming at a seamless transition for the end user

After this step, new backends will probably enable new functionality to be added

PYTORCHLINK

A link to the PyTorch is currently in development

At the moment, the code shown in this demo is only available internally

It will not constitute official documented functionality, the end goal is to connect the user-exposed NN framework to it

Despite the “Py” in PyTorch, this project does not involve any Python. We are directly hooking up to the C++ library responsible for all the computations performed by PyTorch

PacletDirectoryLoad["~/Git/Paclets/pytorchlink/PyTorchLink"];
Get["PyTorchLink`"]

TENSORS & OPERATORS

Create a PyTorch tensor and read data back from it

t=PTTensorCreate[{{-0.62,0.38,0.67},{0.75,-0.34,-0.89}},"Real32"]

PTTensorGet[t]

Run the ReLU operator in PyTorch and inspect its output

output=PTOp["relu"][t]

Normal@PTTensorGet[output]

Run the Linear operator

{input,w,b}=RandomReal[{-1,1},#]&/@{{2},{3,2},{3}}

{inputT,wT,bT}=PTTensorCreate[#,"Real32"]&/@{input,w,b}

output=PTOp["linear"][inputT,wT,bT]

Normal@PTTensorGet[output]

Compare with the result from LinearLayer

LinearLayer[3,"Input"->{2},"Weights"->w,"Biases"->b]@input

GRADIENTS

Create a tensor with "RequiresGrad”->True use it as an input to compute the backward pass  of a simple operator

data={{-0.62,0.38,0.67},{0.75,-0.34,-0.89}};

input=PTTensorCreate[data,"Real32","RequiresGrad"->True]

output=PTOp["sin"][input]

PTTensorBackward[output]

Inspect the computed gradient and compare with the expected result

grad=PTTensorGradient[input]

Normal@PTTensorGet[grad]

Cos[data]

LENET EVALUATION

Extract the arrays from the pre-trained LeNet model and convert them to PyTorch tensors

arrays=PTTensorCreate[#,"Real32"]&/@KeyMap[
StringRiffle[ToString/@#,"_"]&,
Information[NetModel["LeNet Trained on MNIST Data"],"Arrays"]
]

Create a function that evaluates LeNet using PyTorch operators

NetModel["LeNet Trained on MNIST Data"]

lenet=RightComposition[
PTOp["conv2d"][#,arrays["1_Weights"],arrays["1_Biases"],{1,1},{0,0},{1,1},1]&,
PTOp["relu"],
PTOp["max_pool2d"][#,{2,2},{2,2},{0,0},{1,1},False]&,
PTOp["conv2d"][#,arrays["4_Weights"],arrays["4_Biases"],{1,1},{0,0},{1,1},1]&,
PTOp["relu"],
PTOp["max_pool2d"][#,{2,2},{2,2},{0,0},{1,1},False]&,
PTOp["flatten"][#,-3,-1]&,
PTOp["linear"][#,arrays["8_Weights"],arrays["8_Biases"]]&,
PTOp["relu"],
PTOp["linear"][#,arrays["10_Weights"],arrays["10_Biases"]]&,
PTOp["softmax"][#,-1]&
];

Run it on a random MNIST image

inputImg=First@RandomChoice@ResourceData["MNIST"]

inputData={ImageData[inputImg]};

input=PTTensorCreate[inputData,"Real32"]
output=lenet[input]

probs=Normal@PTTensorGet[output]
Ordering[probs,-1]-1

Compute the scores for each class and compare the numerical results with the original

lenetLogits=RightComposition[
PTOp["conv2d"][#,arrays["1_Weights"],arrays["1_Biases"],{1,1},{0,0},{1,1},1]&,
PTOp["relu"],
PTOp["max_pool2d"][#,{2,2},{2,2},{0,0},{1,1},False]&,
PTOp["conv2d"][#,arrays["4_Weights"],arrays["4_Biases"],{1,1},{0,0},{1,1},1]&,
PTOp["relu"],
PTOp["max_pool2d"][#,{2,2},{2,2},{0,0},{1,1},False]&,
PTOp["flatten"][#,-3,-1]&,
PTOp["linear"][#,arrays["8_Weights"],arrays["8_Biases"]]&,
PTOp["relu"],
PTOp["linear"][#,arrays["10_Weights"],arrays["10_Biases"]]&
];

Normal@PTTensorGet@lenetLogits[input]

NetModel["LeNet Trained on MNIST Data"][inputImg,NetPort[{10,"Output"}]]

LENET TRAINING

Randomly initialize the weights and biases and attach a PyTorch optimizer to them

weightsInitializationSettings=<|
"1_Weights"->{{20,1,5,5},0.283},
"4_Weights"->{{50,20,5,5},0.063},
"8_Weights"->{{500,800},0.05},
"10_Weights"->{{10,500},0.048}
|>;
biasesInitializationSettings=<|"1_Biases"->20,"4_Biases"->50,"8_Biases"->500,"10_Biases"->10|>;
arrays=PTTensorCreate[#,"Real32","RequiresGrad"->True]&/@Join[
RandomVariate[NormalDistribution[0,#[[2]]],#[[1]]]&/@weightsInitializationSettings,
ConstantArray[0,#]&/@biasesInitializationSettings
]
opt=PTOptimizerCreate[Values@arrays,"SGD"]

Get the MNIST training and validation datasets

getImgData[dataset_]:=Transpose@MapApply[List]@RandomSample@MapAt[List@*ImageData,dataset,{All,1}];
trainData=getImgData@ResourceData["MNIST","TrainingData"];
valData=getImgData@ResourceData["MNIST","TestData"];

Define functions to perform a training round and to compute the validation loss

trainingRound[net_,data_,batchSize_,opt_,roundIdx_]:=Module[{span,inputBatch,targetBatch,netOutput,loss},
Do[
PTOptimizerZeroGrad[opt];

span=(batchIdx-1)*batchSize+1;;batchIdx*batchSize;
inputBatch=PTTensorCreate[data[[1,span]],"Real32"];
targetBatch=PTTensorCreate[data[[2,span]],"Integer64"];

netOutput=net[inputBatch];
loss=PTOp["cross_entropy_loss"][netOutput,targetBatch,lossWeights,1];

PTTensorBackward[loss];
PTOptimizerStep[opt];

If[batchIdx===1||Mod[batchIdx,30]===0,
loss=PTTensorGet[loss];
AppendTo[batchLossList,loss];
Print@TemplateApply["Round: `` Batch: `` Loss: ``",{roundIdx,batchIdx,loss}]
];
,
{batchIdx,1,Length[data[[1]]]/batchSize}
]
];

validationLoss[net_,data_,batchSize_,roundIdx_]:=Module[
{allLosses,allAcc,span,inputBatch,targetBatch,netOutput,batchLoss,pred,batchAcc,valLoss,valAcc},
allLosses=allAcc={};
Do[
span=(batchIdx-1)*batchSize+1;;batchIdx*batchSize;
inputBatch=PTTensorCreate[data[[1,span]],"Real32"];
targetBatch=PTTensorCreate[data[[2,span]],"Integer64"];

netOutput=net[inputBatch];
batchLoss=PTOp["cross_entropy_loss"][netOutput,targetBatch,lossWeights,1];
pred=PTOp["argmax"][netOutput,1,False];
batchAcc=1./batchSize*Count[
MapThread[SameQ,
{PTTensorGet[pred,"ReturnedObject"->List],data[[2,span]]}
],
True
];
AppendTo[allLosses,PTTensorGet[batchLoss]];
AppendTo[allAcc,batchAcc];
,
{batchIdx,1,Length[data[[1]]]/batchSize}
];
valLoss=Mean[allLosses];
valAcc=Mean[allAcc];
AppendTo[validationLossList,valLoss];
AppendTo[validationAccList,valAcc];
Print@TemplateApply["Validation Loss: `` Validation Accuracy: ``",{valLoss,valAcc}]
]

Train for 3 rounds with a batch size of 200:

nRounds=3;
batchSize=200;
batchLossList=validationLossList=validationAccList={};

lossWeights=PTTensorCreate[Table[1,10],"Real32"];
Do[
trainingRound[lenetLogits,trainData,batchSize,opt,roundIdx];
validationLoss[lenetLogits,valData,batchSize,roundIdx];
,
{roundIdx,1,nRounds}
]

Plot the losses and validation accuracy evolution

ListLogPlot[{Transpose@{Range[11,Length[batchLossList],11],validationLossList},batchLossList},Joined->True,PlotRange->All]
ListLogPlot[validationAccList,Joined->True,PlotRange->All]

Run the net on a random example of the validation set

randomExample=RandomChoice@valData[[1]];
Image[First@randomExample]

output=lenet@PTTensorCreate[randomExample,"Real32"]

probs=Normal@PTTensorGet[output]

Ordering[probs,-1]-1

POSSIBLE FUTURE DEVELOPMENTS

Add llama.cpp to NetExternalObject

Write a link to Tensorflow

Look into other options (such as Jax)

Compare the link to PyTorch and Tensorflow (and possible others) in order to asses their strengths and weaknesses as possible backends of the NN framework

Possibly decide to use multiple backends

Also look into adding them as backends for NetExternalObject

Use the new backend(s) in the main framework and phase out MXNet

Look into further developments unlocked by the new backend(s)