Batches of inputs can be given to a net and it will automatically thread over each input. This will usually be much faster than using Map
to apply the net to each input on CPU or GPU.
Example:
net = NetInitialize[ NetChain[{1}, "Input" -> 3] ];
net[RandomReal[1, {5, 3}], TargetDevice->"GPU"]
(*{{-0.0646134}, {-0.0683393}, {-0.548216}, {-0.312385}, {-0.220518}}*)
If you have multiple CPUs and GPUs, here is a simple function to distribute the computations in batches evenly across all these devices:
ParallelInference[net_, inputs_List, devices_List] :=
Module[
{parts = Append[Length@inputs]@
Range[0, Length@inputs-1,
Length@inputs/Length@devices // Ceiling]
},
Flatten[
ParallelTable[
net[Take[inputs, {parts[[i]]+1, parts[[i+1]]}]
, TargetDevice -> devices[[i]]]
, {i, Length@parts - 1}]
, 1]
]
It can be used like:
ParallelInference[net,
RandomReal[1, {10, 3}], {"CPU", "CPU", "GPU", {"GPU", 2}}]
This function can be tweaked to allow different percentage of inputs to be distributed to each device.
Note: The first time something is evaluated on a new GPU, it might take a bit longer since the net has to be copied into that GPU.