Tip: Limiting TensorFlow to one GPU

Part 2 (2017)
1

###Goal

Limit a TensorFlow session to one GPU within a Jupyter notebook.

###Method

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # so the IDs match nvidia-smi
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # "0, 1" for multiple

Discussion

This code hides any GPUs you haven’t listed. This will allow you to train a model on one GPU while experimenting on another.

###Reference

2 Likes
2

Good tip. Personally, I prefix my jupyter notebook command with CUDA_VISIBLE_DEVICES=0, which also works well. That way, I can run multiple notebook servers on different ports, each one using a different GPU.

Probably best to put export CUDA_DEVICE_ORDER=PCI_BUS_ID in your .bashrc.

2 Likes
3

e.g.

CUDA_VISIBLE_DEVICES=0 jupyter notebook --port=8888
CUDA_VISIBLE_DEVICES=1 jupyter notebook --port=8889

2 Likes
4

bash script

#!/usr/bin/env bash

echo "DEVICE: ${1-0}"
echo "PORT:   ${2-8888}"
CUDA_VISIBLE_DEVICES=${1-3} jupyter notebook --port=${2-8888}