Last active
December 5, 2018 12:19
-
-
Save YiqinZhao/68a15426845efb2100cada9169b2f345 to your computer and use it in GitHub Desktop.
GPU Bastard, control the resource belong to you!
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import time | |
import argparse | |
import GPUtil | |
from multiprocessing import Process | |
parser = argparse.ArgumentParser(description='GPU Bastard, control the resource belong to you!') | |
parser.add_argument('--index', | |
help='GPU Index', | |
default='0') | |
parser.add_argument('--fraction', | |
help='Fraction of target GPU memory needed.\n Default is 0, means all memory', | |
default=0) | |
args = parser.parse_args() | |
args = vars(args) | |
os.environ['CUDA_VISIBLE_DEVICES'] = args['index'] | |
target_index = int(args['index']) | |
fraction_needed = float(args['fraction']) | |
if not (0 <= fraction_needed < 1): | |
print('--fraction should be a float number between 0 and 1.') | |
exit(1) | |
print('\nGPU Bastard Activated! Target GPU:%s\n' % str(target_index)) | |
def take_control(): | |
print('\n==> [%s] Available GPU Found, taking control....\n' % time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime())) | |
import tensorflow as tf | |
if not fraction_needed == 0: | |
config = tf.ConfigProto() | |
config.gpu_options.per_process_gpu_memory_fraction = fraction_needed | |
tf.Session(config=config) | |
tf.enable_eager_execution() | |
tf.add(1, 2) | |
print('\n==> [%s] Available GPU taken\n' % time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime())) | |
# This almostly equals to forever... | |
time.sleep(60 * 60 * 24) | |
while True: | |
available_list = GPUtil.getAvailable(maxLoad=1, maxMemory=min(0.9, 1 - fraction_needed)) | |
if target_index in available_list: | |
p = Process(target=take_control) | |
p.start() | |
time.sleep(3) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment