Under the Hood: A PyTorch Engine for Neural Networks in Wolfram Language

Matteo Salvarezza, Wolfram Research

#WolframTechConf

Abstract

The Wolfram Neural Net Repository and neural net framework are in the process of being migrated to a PyTorch back end, bringing expanded hardware support and performance gains for both training and inference. This session takes a technical look at how PyTorch is being integrated with Wolfram Language’s symbolic neural network definitions, which will enable hybrid execution and smoother workflows for researchers.

Setup steps

To follow along, download and install the experimental paclet. Run these commands from a fresh kernel:

nnURL="https://amoeba.wolfram.com/index.php/s/dw6BwP7tm9EPLHm/download";
PacletInstall@URLDownload[nnURL,FileNameJoin@{$TemporaryDirectory,"NeuralNetworks-15.0.1.paclet"}]

Mac (Apple Silicon) users will also need to install an additional paclet for GPU support:

If[$SystemID==="MacOSX-ARM64",
ptlURL="https://amoeba.wolfram.com/index.php/s/MNi887NNQ4HNx9N/download";
PacletInstall@URLDownload[ptlURL,FileNameJoin@{$TemporaryDirectory,"PyTorchLink-Mac-15.0.1.paclet"}]
]

The paclets are only compatible with the WL 15.0 beta. They can be uninstalled later using these commands:

PacletUninstall["NeuralNetworks"->"15.0.1"]
If[$SystemID==="MacOSX-ARM64",
PacletUninstall["PyTorchLink"->"15.0.1"]
]

Define a GPU device string depending on your platform:

gpuDevice=Switch[$SystemID,
"MacOSX-ARM64","MPS",
"Linux-x86-64"|"Windows-x86-64","CUDA",
_,None
]

Outline

PyTorchLink progress & overview

PyTorchLink as a high-performance numerical utility library

PyTorchLink in the Neural Networks framework

Backends overview in the Neural Networks framework

PyTorchLink progress & overview

PyTorchLink has enjoyed a lot of progress in the past year, with new features including:

900+ operators added

CUDA support for Nvidia GPUs

Metal support for Apple GPUs

Higher order gradients and general multi-output backward pass

Complex number support

Note: despite its name, PyTorchLink doesn’t involve any Python. It is a direct link between WL and libtorch, the C++ library used in the PyTorch Python package.

Note: PyTorchLink is not user-facing functionality, it’s an internal paclet. It is not mentioned in the official documentation and its functions can change without notice across WL versions.

Load PyTorchLink and check its documented symbols:

Get["PyTorchLink`"]

?PyTorchLink`*

The terminology and feel adopted in PyTorchLink are similar to its Python interface.

Tensors

?PTTensorCreate

Create a Torch tensor and read data back from it:

tensor=PTTensorCreate[{{-0.62,0.38,0.67},{0.75,-0.34,-0.89}}]

PTTensorGet[tensor]
PTTensorGet[tensor,"ReturnedObject"->List]

Create a Torch tensor specifying the data type:

PTTensorCreate[{{-0.62,0.38,0.67},{0.75,-0.34,-0.89}},"Real16"]

Create a Torch tensor on the GPU:

gpuTensor=PTTensorCreate[{{-0.62,0.38,0.67},{0.75,-0.34,-0.89}},"Real16",gpuDevice]

PyTorchLink progress & overview

Operators

Most operators are accessed via the syntax PTOp[opName][arguments]. The available operators and their arguments are given by the Association $PTOperatorArguments:

Length@$PTOperatorsArguments

Keys@$PTOperatorsArguments//Short

$PTOperatorsArguments["relu"]//GeneralUtilities`PrettyForm
$PTOperatorsArguments["linear"]//GeneralUtilities`PrettyForm

$PTOperatorArguments only contains the signature of operators but not their documentation. Operator names in PyTorchLink are almost always the same as the names in the Python interface, so one can refer to the Python documentation.

Run the Linear operator:

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

{inputT,wT,bT}=PTTensorCreate/@{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

Run the Linear operator on GPU:

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

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

Normal@PTTensorGet[output]

All the tensor inputs to operators must live on the same device, otherwise an error is thrown:

PTOp["linear"][PTTensorCreate[{{1,2},{3,4}},"Real32","CPU"],PTTensorCreate[{{1,2},{3,4}},"Real32",gpuDevice]]

Several native operators are supported outside of the PTOp wrapper. Arithmetic and in-place arithmetics:

sum=PTTensorCreate[{1,2,3}]+PTTensorCreate[{1,2,3}]
pow=PTTensorCreate[{1,2,3}]^PTTensorCreate[{1,2,3}]

Normal@PTTensorGet[sum]
Normal@PTTensorGet[pow]

t=PTTensorCreate[{1,2,3}];
t/=4.5
Normal@PTTensorGet[t]

Part extraction with Torch and WL semantics (in torch semantics, positive values are 0-indexed and Span end positions are not included):

data=RandomReal[1,{3,4,5}];
t=PTTensorCreate[data]

(* Torch semantics *)
torchPartResult=t[[0,2,1;;4]]
(* WL semantics *)
wlPartResult=t[[1,3,2;;4,"WLPart"->True]]

Normal@PTTensorGet@torchPartResult
Normal@PTTensorGet@wlPartResult

data[[1,3,2;;4]]

PyTorchLink progress & overview

Gradients

Create a tensor with "RequiresGrad”->True, feed it to an operator and perform the backward pass:

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

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

output=PTOp["sum"]@PTOp["sin"][input]

PTTensorBackward[output]

Inspect the computed gradient and compare with the expected result

grad=PTTensorGradient[input]

Normal@PTTensorGet[grad]

Cos[data]

Note: running the backward pass twice in a row triggers a libtorch error:

input=PTTensorCreate[data,"RequiresGrad"->True];
output=PTOp["sum"]@PTOp["sin"][input];
PTTensorBackward[output]
PTTensorBackward[output]

Higher order gradients are computed by calling PTTensorBackward repeatedly in conjunction with “CreateGraph” -> True:

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

output=PTOp["sum"]@PTOp["tanh"][input]

order=4;
Do[
PTTensorBackward[output,Automatic,"CreateGraph"->True];
output=PTTensorGradient[input];
PTTensorResetGradient[input],
order-1
];
PTTensorBackward[output,Automatic];
grad=PTTensorGradient[input]

Normal@PTTensorGet[grad]
D[Tanh[x],{x,4}]/.x->data

PTTensorBackward is completely general and can be used with any number of non-scalar outputs:

?PTTensorBackward

PyTorchLink progress & overview

LeNet Evaluation

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

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

Create a function that evaluates LeNet using Torch 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"}]]

PyTorchLink progress & overview

LeNet Training

Set a device to use for training (“CPU” or gpuDevice)

trainingDevice="CPU";

Randomly initialize the weights and biases and attach a Torch 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",trainingDevice,"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",trainingDevice];
targetBatch=PTTensorCreate[data[[2,span]],"Integer64",trainingDevice];

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",trainingDevice];
targetBatch=PTTensorCreate[data[[2,span]],"Integer64",trainingDevice];

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",trainingDevice];
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",trainingDevice]

probs=Normal@PTTensorGet[output]

Ordering[probs,-1]-1

PyTorchLink as a high-performance numerical utility library

With its rich set of optimized operators, ease of use and device support, PyTorchLink is a great candidate to augment several system functions, especially in the image and audio processing domains.

Define a pyBlur function which blurs an image by performing a convolution with a 2D Gaussian filter in PyTorchLink:

pad1Ch=ArrayReshape[#,Join[If[ArrayDepth[#]==2,{1,1},{1}],Dimensions[#]]]&;

imageToBatch[img_Image,type_,device_]:=PTTensorCreate[pad1Ch@Image`InternalImageData[img,type,Interleaving->False],type,device]
batchToImage[tensor_PTTensor]:=Image[First@PTTensorGet[tensor],Interleaving->False]

pyBlur[img_,{r_,σ_},type_:"Real64",device_:"CPU"]:=Module[{batch,ch,weights,biases},
batch=imageToBatch[ImagePad[img,r,"Fixed"],type,device];
ch=ImageChannels[img];
weights=PTTensorCreate[ConstantArray[GaussianMatrix[{r,σ},Method->"Gaussian"],{ch,1}],type,device];
biases=PTTensorCreate[ConstantArray[0.,ch],type,device];
batch=PTOp["conv2d"][batch, weights, biases, {1,1}(*stride*), {0,0}(*padding*), {1,1}(*dilation*), ch(*groups*)];
batchToImage[batch]
]

Define an optimized pyBlurSeparable, which separates the 2D convolution in two successive 1D convolutions:

pyBlurSeparable[img_,{r_,σ_},type_:"Real64",device_:"CPU"]:=Module[{batch,ch,kernel,weightsH,weightsV,biases},
batch=imageToBatch[ImagePad[img,r,"Fixed"],type,device];
ch=ImageChannels[img];
(*Horizontal and vertical 1D kernels*)
kernel=ConstantArray[GaussianMatrix[{{r},σ},Method->"Gaussian"],{ch,1,1}];
weightsH=PTTensorCreate[kernel,type,device];
weightsV=PTTensorCreate[Transpose[kernel,{1,2,4,3}],type,device];
biases=PTTensorCreate[ConstantArray[0.,ch],type,device];
(*First horizontal blur*)
batch=PTOp["conv2d"][batch,weightsH,biases,{1,1},{0,0},{1,1},ch];
(*Then vertical blur*)
batch=PTOp["conv2d"][batch,weightsV,biases,{1,1},{0,0},{1,1},ch];
batchToImage[batch]
]

Test the performance of the functions vs GaussianFilter for a blur radius of 50 pixels, and also check that the results are the same in all functions:

img=;

r=50;

With 64 bit precision:

(res1=pyBlur[img,r{1,1/2}])//RepeatedTiming
(res2=pyBlurSeparable[img,r{1,1/2}])//RepeatedTiming
(res3=GaussianFilter[img,r{1,1/2}])//RepeatedTiming

DistanceMatrix[{res1,res2,res3},DistanceFunction->Mean@*Mean@*Abs@*Subtract]//MatrixForm

With 32 bit precision:

(res1=pyBlur[img,r{1,1/2},"Real32"])//RepeatedTiming
(res2=pyBlurSeparable[img,r{1,1/2},"Real32"])//RepeatedTiming
(res3=GaussianFilter[img,r{1,1/2}])//RepeatedTiming

DistanceMatrix[{res1,res2,res3},DistanceFunction->Mean@*Mean@*Abs@*Subtract]//MatrixForm

With 32 bit precision on GPU:

(res1=pyBlur[img,r{1,1/2},"Real32",gpuDevice])//RepeatedTiming
(res2=pyBlurSeparable[img,r{1,1/2},"Real32",gpuDevice])//RepeatedTiming
(res3=GaussianFilter[img,r{1,1/2}])//RepeatedTiming

DistanceMatrix[{res1,res2,res3},DistanceFunction->Mean@*Mean@*Abs@*Subtract]//MatrixForm

For audio processing there are several Fourier Transform operators, whose name starts with fft:

Select[Keys@$PTOperatorsArguments,StringStartsQ[#,"fft"]&]

Again, refer to the Python documentation for information about these operators.

PyTorchLink in the Neural Networks framework

PyTorchLink is now partially hooked up to the top-level framework:

Supported features:

Net evaluation (NetChain, NetGraph, ~50% layer support)

Evaluation gradients (NetPortGradient)

GPU evaluation (TargetDevice)

Unsupported features:

Variable-length inputs (currently WIP)

NetEncoders/NetDecoders (easy)

NetTrain

NetMeasurements

PyTorchLink evaluation is enabled by setting EvaluationBackend -> “Torch” (undocumented):

net=NetModel["LeNet Trained on MNIST Data"]

{enc,dec}=NetExtract[net,{{"Input"},{"Output"}}]

imgs=RandomChoice[ResourceData["MNIST"],3][[All,1]]

dec@net[enc@imgs,EvaluationBackend->"Torch"]

Gradient/GPU evaluation:

net2=NetTake[net,{1,-2}]

out=net2[enc@imgs,{NetPort["Output"],NetPortGradient[All]},EvaluationBackend->"Torch",TargetDevice->"GPU"]

TargetDevice -> “GPU” will route to either Cuda or Metal depending on the platform (more in the next slide)

Torch evaluation is experimental and not tested thoroughly yet, so expect problems. Variable length inputs and layers that were not hooked up yet can fail in weird ways:

l=ElementwiseLayer[Ramp,"Input"->"Varying"];
l[{1,2,3},EvaluationBackend->"Torch"]

l=AttentionLayer["Dot","Key"->{1,2},"Value"->{1,3},"Query"->{1,2}]
l[<|"Key"->{{1,2}},"Value"->{{1,2,3}},"Query"->{{1,2}}|>,EvaluationBackend->"Torch"]

Check the supported layers:

supported=Select[NeuralNetworks`$NetHeads,Lookup[NeuralNetworks`$LayerData[NeuralNetworks`Private`$SymbolToType[#]],"TorchForward"]=!=None&]

unsupported=Complement[NeuralNetworks`$NetHeads,supported]

Backends overview in the Neural Networks framework

The following picture illustrates the available backends and devices for net evaluation only. NetTrain and NetMeasurements only support MXNet still.

Backends and devices can now be selected independently via options EvaluationBackend and TargetDevice

The default value for EvaluationBackend is Automatic, which will try to select the best backend given the platform and the provided value for TargetDevice (default “CPU”), but it’s better to set both options to control explicitly which path the framework routes to

Allowed values for EvaluationBackend are Automatic, “MXNet”, “ONNXRuntime” and “Torch” (undocumented)

EvaluationBackend -> Automatic never resolves to “Torch”

EvaluationBackend is not documented yet but it will be in the released version of WL 15.0

NeuralNetworks`$EvaluationBackend can be set to any of the backends to force the framework to use it by default

Passing the EvaluationBackend option explicitly overrides this default

Not all combinations of backend and device are valid, sometimes depending on the platform

ElementwiseLayer[Ramp][1,EvaluationBackend->"MXNet",TargetDevice->"DirectML"]

ElementwiseLayer[Ramp][1,EvaluationBackend->"ONNXRuntime",TargetDevice->"DirectML"]

On macOS, when EvaluationBackend -> “Torch” is passed, TargetDevice -> “Metal” becomes available to use the GPU

TargetDevice -> ”GPU” is a meta-setting that translates to either “CUDA”/”DirectML”/”CoreML”/”Metal” depending on the platform, the evaluation backend and available hardware

NetExternalObject is explicitly tied to a backend so there’s no EvaluationBackend option for it