gpu_check

gpu_check.py

def gpu_check():
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '/device/gpu:0'
    import tensorflow as tf

    from datetime import datetime

    shape = (10, 10)
    device_name = "/gpu:0"
    with tf.device(device_name):
        random_matrix = tf.random_uniform(shape=shape, minval=0, maxval=1)
        dot_operation = tf.matmul(random_matrix, tf.transpose(random_matrix))
        sum_operation = tf.reduce_sum(dot_operation)

    startTime = datetime.now()
    with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as session:
        result = session.run(sum_operation)
        print(result)
    print("Time taken:", datetime.now() - startTime)
    
def tf_check():
    import tensorflow.contrib.keras as K, numpy as np

    resnet = K.applications.resnet50
    model_settings = {'include_top': False, 'weights': 'imagenet', 'pooling': 'max'}
    model = resnet.ResNet50(**model_settings)
    a = np.random.rand(224, 224, 3)
    b = np.expand_dims(a, axis=0)
    ready = resnet.preprocess_input(b)
    model.predict(ready)