Neural network evaluate on multiple gpus simultaneously

Question

For instance, I have 4 gpus, now I want to calculate neural network on each gpu simultaneously, or make use of all of my gpus to calculate the data.

The code might like

data = RandomReal[{0, 1}, {10000, 10}];
net1 = NetInitialize@NetChain[{LinearLayer[{}], LogisticSigmoid}, "Input" -> {10}];
net2 = NetInitialize@NetChain[{LinearLayer[{}], LogisticSigmoid}, "Input" -> {10}];
net3 = NetInitialize@NetChain[{LinearLayer[{}], LogisticSigmoid}, "Input" -> {10}];
net4 = NetInitialize@NetChain[{LinearLayer[{}], LogisticSigmoid}, "Input" -> {10}];

{net1[data, TargetDevice -> {"GPU", 1}], 
 net2[data, TargetDevice -> {"GPU", 2}], 
 net3[data, TargetDevice -> {"GPU", 3}], 
 net4[data, TargetDevice -> {"GPU", 4}]}

Is this code make the neural network calculate simultaneously? Is there any way to realize this idea?

I have tried net1[data,TargetDevice->{"GPU",All}], but it returned the message NetChain::nomultisupp: Net evaluation is not currently supported with TargetDevice -> {GPU, All}. In NetTrain we can set TargetDeivce->{"GPU",All}, why there cannot?

---

Update

Edmund's method is a right way to distribute jobs to multi-GPUs, I have tested it on my linux system with four GPUs. Here's the test record.

First, launch 4 local kernels and prepare the data

LaunchKernels[4];
data = RandomReal[{0, 1}, {10000000, 10}];

Second, initialize the network

nets = NetInitialize@{
 NetChain[{LinearLayer[{}], LogisticSigmoid},"Input" -> {10}],
 NetChain[{LinearLayer[{}], LogisticSigmoid},"Input" -> {10}],
 NetChain[{LinearLayer[{}], LogisticSigmoid},"Input" -> {10}],
 NetChain[{LinearLayer[{}], LogisticSigmoid},"Input" -> {10}]};

Then use the parallel functions, I prefer to use ParallelEvaluate, and Edmund's method is okay. In my test (4 1080Ti), the time cost roughly 16s.

l = ParallelEvaluate[
 nets[[$KernelID]][data,TargetDevice ->{"GPU",$KernelID}]
];//AbsoluteTiming

You can compare with sequentially calculate

{nets[[1]][data, TargetDevice -> {"GPU", 1}], 
 nets[[2]][data, TargetDevice -> {"GPU", 2}], 
 nets[[3]][data, TargetDevice -> {"GPU", 3}], 
 nets[[4]][data, TargetDevice -> {"GPU", 4}]}//AbsoluteTiming

and use the shell command nvidia-smi to see whether four GPUs are working simultaneously. The time cost is roughly 42s, you can see it almost 2.5 times slower than use 4 GPUs together.

Last I found that it's essential to launch kernels before initialize the networks, if you exchange the order, there's a bug report like

OMP: Error #13: Assertion failure at z_Linux_util.cpp(2338).
OMP: Hint Please submit a bug report with this message, compile and run commands used, and machine configuration info including native compiler and operating system versions. Faster response will be obtained by including all program sources. For information on submitting this issue, please see http://www.intel.com/software/products/support/.
SubKernels`SubKernels::timekernels: Timeout for subkernels. Received only 0 of 4 connections.

So, anyone who knows why?

Answer 1

You may use ParallelSubmit and WaitAll.

With data and net*'s as defined in OP.

Create a set of parallel jobs with ParallelSubmit. I only have 1 GPU so I set the others to "CPU".

jobs =
 ReleaseHold[
  Map[
   ParallelSubmit
   , Hold[{
     net1[data, TargetDevice -> {"GPU", 1}]
     , net2[data, TargetDevice -> "CPU"]
     , net3[data, TargetDevice -> "CPU"]
     , net4[data, TargetDevice -> "CPU"]
     }]
   , {2}
   ]]

The Hold is necessary to prevent the nets from evaluating in serial in Map. ReleaseHold removes Hold; ParallelSubmit has HoldAllComplete so the nets will not evaluate within them.

Then execute with WaitAll. You will see the job objects update their status as they run.

res = WaitAll[jobs];

Hope this helps.