pythongputheanolasagne

Select GPU during execution in Theano


I am training neural nets with theano and lasagne on a 4 GPU machine. My .theanorc contains the following lines:

[global]
device = gpu0

So when in python I execute import theano, I get Using gpu device 0: GRID K520

What if, after importing theano, I chose to use say gpu1? I'd like to do this dynamically, that is, without editing .theanorc is it possible? Or even to choose it at runtime?


Solution

  • I'm afraid it's not possible to change the execution device after Theano has been imported. From the documentation:

    config.device

    String value: either 'cpu', 'gpu', 'gpu0', 'gpu1', 'gpu2', or 'gpu3'

    [...]

    This flag’s value cannot be modified during the program execution.

    Bonus: however, let's say you wanted to have two Python processes each running on a separate GPU (is that what you want?), then you could do something like:

    import os
    os.system("THEANO_FLAGS='device=gpu0' python myscript.py")
    os.system("THEANO_FLAGS='device=gpu1' python myscript.py")
    

    or hack into/extend Python's multiprocessing module (which works by spawning subprocesses) to ensure the flag is set before a child process is spawned.