python - Tensorflow and Multiprocessing: Passing Sessions

25
votes

I have recently been working on a project that uses a neural network for virtual robot control. I used tensorflow to code it up and it runs smoothly. So far, I used sequential simulations to evaluate how good the neural network is, however, I want to run several simulations in parallel to reduce the amount of time it takes to get data.

To do this I am importing python's multiprocessing package. Initially I was passing the sess variable (sess=tf.Session()) to a function that would run the simulation. However, once I get to any statement that uses this sess variable, the process quits without a warning. After searching around for a bit I found these two posts: Tensorflow: Passing a session to a python multiprocess and Running multiple tensorflow sessions concurrently

While they are highly related I haven't been able to figure out how to make it work. I tried creating a session for each individual process and assigning the weights of the neural net to its trainable parameters without success. I've also tried saving the session into a file and then loading it within a process, but no luck there either.

Has someone been able to pass a session (or clones of sessions) to several processes?

Thanks.

2
pythonparallel-processingmultiprocessingtensorflowreinforcement-learning
I suspect you can't pass clones of Sessions between processes because there's state in C address space that Python doesn't know how to copy. But creating brand new sessions in each new process should work. I haven't used multiprocessing, but I often have couple of processes open in parallel that keep their own TensorFlow sessions - Yaroslav Bulatov
The second link I provided runs several processes in parallel, but the issue is that I need the neural network to be the same for ALL the processes. - MrRed
You may be able to work-around it by using TensorFlow distributed -- ie, have a local worker and ps and create multiple sessions in parallel like tf.Session("grpc://localhost:2222") github.com/tensorflow/tensorflow/blob/master/tensorflow/g3doc/… - Yaroslav Bulatov
@YaroslavBulatov , any chance you could elaborate a little bit? I am implementing parallel TRPO (Viva RL) but now working with the minimal example [ gist.github.com/dd210/e8bad8eadc19f44cafcdc5313a39a53f ]. Basically I would like to launch Session in main process, do some calculations, then transfer network to another X parallel processes and do some calculations in parallel there. Then repeat Y times. But I got stuck with this minimal example, because it seems that I need to work in one session all the time. Official tutorial is a <bit> vague. Appreciate any help and advice. - dd210
Multiprocessing seems to be semi-broken with TensorFlow. Here's an example of using distributed tensorflow -- gist.github.com/yaroslavvb/ea1b1bae0a75c4aae593df7eca72d9ca - Yaroslav Bulatov

2 Answers

14
votes

You can't use Python multiprocessing to pass a TensorFlow Session into a multiprocessing.Pool in the straightfoward way because the Session object can't be pickled (it's fundamentally not serializable because it may manage GPU memory and state like that).

I'd suggest parallelizing the code using actors, which are essentially the parallel computing analog of "objects" and use used to manage state in the distributed setting.

Ray is a good framework for doing this. You can define a Python class which manages the TensorFlow Session and exposes a method for running your simulation.

import ray
import tensorflow as tf

ray.init()

@ray.remote
class Simulator(object):
    def __init__(self):
        self.sess = tf.Session()
        self.simple_model = tf.constant([1.0])

    def simulate(self):
        return self.sess.run(self.simple_model)

# Create two actors.
simulators = [Simulator.remote() for _ in range(2)]

# Run two simulations in parallel.
results = ray.get([s.simulate.remote() for s in simulators])

Here are a few more examples of parallelizing TensorFlow with Ray.

See the Ray documentation. Note that I'm one of the Ray developers.

10
votes

I use keras as a wrapper with tensorflow as a backed, but the same general principal should apply.

If you try something like this:

import keras
from functools import partial
from multiprocessing import Pool

def ModelFunc(i,SomeData):
    YourModel = Here
    return(ModelScore)

pool = Pool(processes = 4)
for i,Score in enumerate(pool.imap(partial(ModelFunc,SomeData),range(4))):
    print(Score)

It will fail. However, if you try something like this: 

from functools import partial
from multiprocessing import Pool

def ModelFunc(i,SomeData):
    import keras
    YourModel = Here
    return(ModelScore)

pool = Pool(processes = 4)
for i,Score in enumerate(pool.imap(partial(ModelFunc,SomeData),range(4))):
    print(Score)

It should work. Try calling tensorflow separately for each process.