I know that Mathematica supports GPU inference, and can be told which GPU to run models on by specifying TargetDevice->{"GPU",2} or such.

My questions are as follows: 1) By using the above statement with multiple GPU numbers and some parallelization, can Mathematica address multiple GPUs within the same session for performance gains in inference? 2) Can inference be run in batches such as to better utilize GPU memory and processing power?


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.


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] :=
   {parts = Append[Length@inputs]@
      Range[0, Length@inputs-1, 
         Length@inputs/Length@devices // Ceiling]
         net[Take[inputs, {parts[[i]]+1, parts[[i+1]]}]
         , TargetDevice -> devices[[i]]]
      , {i, Length@parts - 1}]
   , 1]

It can be used like:

  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.

  • $\begingroup$ Thank you, I'll test this as soon as possible. If I want to batch inputs for a network with multiple input ports, the syntax would be <|"Port1"->{port1input1,port1input2},"Port2"->{port2input1,port2input2}|>, where input1 and input2 are discrete examples to be inferred? $\endgroup$ – Yss Aug 15 '17 at 22:22
  • $\begingroup$ Yes. That should work $\endgroup$ – dan7geo Aug 15 '17 at 22:26
  • $\begingroup$ I tested the batched inference portion of this and got about a 2x speed increase over using an external Map statement. Thanks again. $\endgroup$ – Yss Aug 16 '17 at 14:41

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.