FNetGraph: an ergonomic alternative to NetGraph​
​by Alec Graves
NetGraph, the primary function in WL for constructing neural networks, is quite lengthy. For those of us who are used to the world of TensorFlow / Keras / PyTorch where idiomatic Python can be easily used to construct networks, the string association required to build custom networks in WL is (in my opinion) sub-optimally elegant.
​
But thankfully, WL is a functional language with nearly limitless flexibility thanks to rewrite rules, so we can easily construct a more ergonomic interface to NetGraph.

Other neural network libraries

Keras has one of the simplest interfaces for neural network construction. It is unfortunate that Keras’s TensorFlow backend takes over a second to import into a python program; otherwise, it would be a fine experience.
Here is an example of Keras' "Functional API":
〉
from tensorflow import keras
from tensorflow.keras import layers
​
# Example of two branches and an add:
netin = keras.Input(shape=(784,))
a = layers.Dense(64, activation="relu", name="dense_relu_1")(netin)
a = layers.Dense(64, activation="relu", name="dense_relu_2")(a)
b = layers.Dense(64, activation="relu", name="dense_relu_3")(netin)
netout = layers.Add(name="plus")([a, b])
​
model = tf.keras.models.Model(inputs= netin, outputs=netout)
model.summary()
The functional style yields a very minimal and fairly repetition-free representation compared to other neural network libraries (such as PyTorch with its OOP style leveraging Python classes).
Unfortunately, to build such a network in WL, NetGraph with its string association format must be used (or you can do rules of numbers where you have to look up what each variable is). Here is the same example from above in WL:

Current state of custom nets in WL

In[]:=
NetGraph[​​<|​​ "dense1" -> LinearLayer[64],​​ "dense2" -> LinearLayer[64],​​ "dense3" -> LinearLayer[64],​​ "act1" -> ElementwiseLayer["ReLU"],​​ "act2" -> ElementwiseLayer["ReLU"],​​ "act3" -> ElementwiseLayer["ReLU"],​​ "add" -> Plus​​|>,​​{{"dense1"->"act1"->"dense2"->"act2", "dense3"->"act3"}->"add"}​​]
Out[]=
NetGraph
uninitial
ized
Inputport:
array
Outputport:
vector(size: 64)

Notice the need to name everything in one place and then use those string names a second time in a list of rules to describe the layout of the architecture? This separation makes writing networks (in my opinion, for complicated networks) repetitive, as it is easy to mistype a string without autocompletion. It is also it is needed to go back to the name association every time you want to add something to the network.
So what would a better interface look like? WL rewrite rules allow us to basically do anything we want, so here, I’ll show what I arrived at after a couple hours of playing around.

Towards a better API for net building in WL

FNetGraph is a wrapper around NetGraph that makes a lot of implementation easier.
In[]:=
FNetGraph[{​​ LinearLayer[64]->ElementwiseLayer["ReLU"]->LinearLayer[64]->ElementwiseLayer["ReLU"],​​ LinearLayer[64]->ElementwiseLayer["ReLU"]}->Plus​​]
Out[]=
NetGraph
uninitial
ized
Inputport:
array
Outputport:
vector(size: 64)

Notice the automatic names? That makes me happy.
We can also add custom names to our layers pretty easily using the “F” function as a wrapper:
In[]:=
FNetGraph[{​​ F[LinearLayer[64], "dense1"]->F[ElementwiseLayer["ReLU"], "act1"]​​ ->F[LinearLayer[64], "dense2"]->F[ElementwiseLayer["ReLU"], "act2"],​​ F[LinearLayer[64], "dense3"]->F[ElementwiseLayer["ReLU"], "act3"]}->F[Plus, "add"]​​]
Out[]=
NetGraph
uninitial
ized
Inputport:
array
Outputport:
vector(size: 64)

We can also greatly simplify construction in this example:
In[]:=
With[{denseRelu = LinearLayer[64]->ElementwiseLayer["ReLU"]},​​ FNetGraph[{denseRelu->denseRelu, denseRelu}->Plus]]
Out[]=
NetGraph
uninitial
ized
Inputport:
array
Outputport:
vector(size: 64)

We also have the option to reference named layers just by using strings:
In[]:=
denseReluBn[n_Integer]:= NetChain[{​​ LinearLayer@n, ElementwiseLayer@Ramp, BatchNormalizationLayer[]}]​​​​FNetGraph[{​​ {​​ F[denseReluBn[64], "dense1"],​​ F[denseReluBn[64], "dense2"],​​ denseReluBn[128] -> denseReluBn[64],​​ denseReluBn[128] -> denseReluBn[64],​​ NetPort["pinecone"]​​ } -> Plus -> {​​ denseReluBn[64],​​ denseReluBn[64]​​ }, "dense1", "dense2", NetPort@"pinecone"} -> Plus]
Out[]=
NetGraph
uninitial
ized
Number of inputs:
2
Outputport:
vector(size: 64)

Do you see how easy it was to make that graph? This simple rewrite-rule wrapper is able to decrease the boilerplate naming of things and unlock the functional programing potential of WL for neural network construction.

Looking forward

Ideally, the ability to name layers would be integrated into the API for those layers as it is in Keras — so you could write something like LinearLayer[64, “dense1”]. Then FNetGraph would not need the “F” function as a wrapper just to add names.
Anyway, WRI, please add something like this to your language! Maybe this would be a good temporary addition to your function repository, but unorthodox usage of the function “F” as a wrapper to add names to layers is kinda unfortunate. I am open to suggestions.

Appendix: The implementation of FNetGraph

ClearAll[autoLayerName, matchingStrings]​​(* matchingStrings wraps StringCases to provide less annoying output *)​​(* matchingStrings[{"asdf123", "asdf2345", "bsdf4543"},"asdf"~~NumberString] *)​​(* out[]= {"asdf123","asdf2345"} *)​​matchingStrings[xx_List, p_]:=Map[First, Map[(#/.{}->Nothing)&,StringCases[xx, p]]];​​​​(* autoLayerName returns the first unused numbered name in labels (e.g convolutionlayer1) *)​​autoLayerName[x_, labels_]:= Module[{name=ToLowerCase@ToString@Head@x, matchingNums},​​ matchingNums = Map[FromDigits,​​ matchingStrings[Keys@labels, name~~NumberString]​​ // matchingStrings[#, NumberString]&​​ ] // Sort;​​ name<>ToString[If[Length[matchingNums]==0, 1, matchingNums[[-1]]+1]]​​];​​​​ClearAll[FNetGraph]​​(* An elegent wrapper for writing NetGraphs *)​​FNetGraph[fnet_, opts___Rule]:= Module[{labels=<||>, rules,​​ (* pure WL Functions to save to the labels list: *)​​ ThreadingLayerFns={Plus,Subtract,Times,Divide,Power,Surd,Min,Max,Clip,Mod,ArcTan,Erf}​​ },​​ rules=fnet /. {​​ (* Named Layer: *)​​ F[x_, name_String] :> (AssociateTo[labels, name -> x]; name),​​ (* Automatically named layer: *)​​ x:Alternatives@@Join[Blank/@NeuralNetworks`$NetHeads, ThreadingLayerFns, ​​ {_NetGraph, _NetChain (*Anything to add to labels <||> goes here*)}]​​ :> With[{name = autoLayerName[x, labels]}, AssociateTo[labels, name -> x]; name]​​ }; NetGraph[labels, {rules}, opts]​​];

Appendix: Other fun things you can do with graphs and rewrite rules

In[]:=
ClearAll[FGraph];​​(* Ergonomically creates labels and a rule list from ​​ functional-style inputs using rewrite rules *)​​FGraph[fRules_] := Module[{labels=<||>},​​{{fRules /. {​​ (* Named Layer: *)​​ F[x_, name_String] :> (AssociateTo[labels, name -> x]; name),​​ (* Reference named layer: *)​​ F[x_String] :> x, ​​ (* Automatically named layer: *)​​ F[x_] :> With[{name = autoLayerName[x, labels]}, ​​ AssociateTo[labels, name -> x]; name] ​​} //. {​​ {x_} :> x, (* prevents rule matching List[x_] *)​​ x_[y_] :> y -> x (* function application to rule *)​​}}, labels}//Reverse]​​​​ClearAll[FNetGraph1]​​(* a more 'functional' alternative to FNetGraph ​​ that requires everything be wrapped in F *)​​FNetGraph1[graph_, opts___Rule]:= Module[{labels, rules},​​ {labels, rules}=FGraph[graph];​​ NetGraph[labels, rules, opts]]​​​​ClearAll[RuleFlatten]​​(* Flattens a list of rules e.g. {x, y}->z to {x->z, y->z} *)​​RuleFlatten[rules_]:= (rules //. {​​ ((x_ ->{ y__}) :> (#[[1]]->#[[2]]& /@ Tuples[{{x}, {y}}])),​​ (({x__}->y_) :> (#[[1]]->#[[2]]& /@ Tuples[{{x}, {y}}])),​​ Rule[Rule[x__, y__], z__] :> {Rule[x, y], Rule[y, z]},​​ (x_ -> y_ -> z_) :> ({x->y, y->z})​​} // Flatten // DeleteDuplicates)​​(* TODO write the inverse of this function? RuleCollect? *)​​​​{labels, graph} = FGraph[​​{ {F[input, "MyInput"], F@input1} // {F@d // {F@a, F@b}},​​ F@"MyInput"​​} // {F[d], F[e] // F@xml} // F@output]​​​​grules = RuleFlatten@graph;​​Graph[grules, VertexLabels->Thread[VertexList[grules]->(VertexList[grules] /. labels)]]