How do I apply a NetChain element-wise during training?

Question

I need to train a Neural network that models a function that takes a single real number, and outputs a real number.

x -> y = f[x]

However, in the training data, each data point doesn't link a single real number to their output, instead, I have several list (of variable size) of inputs, and the sum of their outputs. The value of the output of each individual input in the training data is unknown.

{x1, x2, x3, ... , xn} -> y1 + y2 + y3 + ... + yn

So, my idea is making a neural network that takes a list of real numbers, then to each real number apply the trainable NetChain independently (each with a real as input and a real as output), and then add them together. After training, I could extract the NetChain from the network, which is what I need.

The problem is that I haven't found a way to do this, ElementwiseLayers only seem to support some few functions, and doesn't support having a neural network acting as a function.

Answer 1

Here is an example of how to use NetMapOperator to do this.

Define a test function and generate training data from it:

testFun = Function[x, Cos[x] + Sqrt[x]];
trainingData = Table[With[{
     x = RandomReal[{0, 10}, {RandomInteger[{5, 20}], 1}]
     },
    x -> Total[testFun /@ x]
    ],
   100
   ];
Dimensions /@ First[trainingData]

{8, 1} -> {1}

As you can see, I'm mapping a n by 1 vector (with n chosen randomly) to a length-1 vector. This is easier to work with inside of the neural networks.

Next, define a simple learnable mapping x -> y using NetChain with Ramp non-linearities. This effectively gives a piecewise linear model for testFun:

regressionNet = NetChain[{5, Ramp, 5, Ramp, 1}]
trainingNet =  NetChain[{
    NetMapOperator[regressionNet], 
    AggregationLayer[Total, 1]
   },
   "Input" -> {"Varying", 1}
]

Train the net and visualize the learned version of testFun:

trainedNet = NetTrain[trainingNet, trainingData];
With[{
  trainedFun = NetExtract[trainedNet, {1, "Net"}]
  },
 Plot[{testFun[x], trainedFun[x]}, {x, 0, 10}]
]

You may not always get this result, so it feel free to fiddle around with the training methods etc.