Skip to content

Instantly share code, notes, and snippets.

@lfoppiano
Last active September 10, 2021 08:05
Show Gist options
  • Save lfoppiano/849fc076d437a95f389d4b61065238c0 to your computer and use it in GitHub Desktop.
Save lfoppiano/849fc076d437a95f389d4b61065238c0 to your computer and use it in GitHub Desktop.
NVIDIA benchmark
# Credits to https://marmelab.com/blog/2018/03/21/using-nvidia-gpu-within-docker-container.html
# Run with
# [CPU] docker run --runtime=nvidia --rm -ti -v "${PWD}:/app" tensorflow/tensorflow:1.15.5-gpu python /app/nvidia-benchmark.py cpu 10000
# [GPU] docker run --runtime=nvidia --rm -ti -v "${PWD}:/app" tensorflow/tensorflow:1.15.5-gpu python /app/nvidia-benchmark.py gpu 10000
import sys
import numpy as np
import tensorflow as tf
from datetime import datetime
device_name = sys.argv[1] # Choose device from cmd line. Options: gpu or cpu
shape = (int(sys.argv[2]), int(sys.argv[2]))
if device_name == "gpu":
device_name = "/gpu:0"
else:
device_name = "/cpu: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:
for x in range(0,10):
result = session.run(sum_operation)
print(result)
# It can be hard to see the results on the terminal with lots of output -- add some newlines to improve readability.
print("\n" * 5)
print("Shape:", shape, "Device:", device_name)
print("Time taken:", str(datetime.now() - startTime))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment