gpu_benchmark.py

# --
import subprocess
from datetime import datetime

time_start = datetime.now()
hostname = subprocess.run("printf `uname -n`", shell = True, capture_output=True, text=True).stdout

print(f"Running PyTorch benchmark on: {hostname}")

# --
import torch

cpu_iterations = 1
gpu_iterations = 6

if torch.cuda.is_available():
    cuda_device_found="Yes"
    device = torch.device("cuda")

    #
    matrix_size = 128 * 128
    x = torch.randn(matrix_size, matrix_size)
    y = torch.randn(matrix_size, matrix_size)

    # CPU Benchmark
    print("")
    print("- CPU Benchmarks -", cpu_iterations, "iterations")
    for i in range(cpu_iterations):
        start = datetime.now()
        cpu_result = torch.mm(x, y)
        print("    Using device:", cpu_result.device, "- took:", datetime.now() - start)

    # GPU Benchmark
    print("- GPU Benchmarks -", gpu_iterations, "iterations")
    x_gpu = x.to(device)
    y_gpu = y.to(device)
    for i in range(gpu_iterations):
        start = datetime.now()
        gpu_result = torch.mm(x_gpu, y_gpu)
        print("    Using device:", gpu_result.device, "- took:", datetime.now() - start)
else:
    cuda_device_found="No"

# --
time_stop = datetime.now()

# --

print('')
print("- REPORT:")
print("    Hostname:", hostname)
print("    Nvidia driver version:", subprocess.run("printf `nvidia-smi --query-gpu=driver_version --format=noheader,csv`", shell = True, capture_output=True, text=True).stdout)
print("    CUDA device found:", cuda_device_found)
print("    CUDA version:", torch.version.cuda)
print("    PyTorch version:", torch.__version__)
print('')
print('    Started at:  ', time_start.strftime("%Y-%m-%d %H:%M:%S"))
print('    Finished at: ', time_stop.strftime("%Y-%m-%d %H:%M:%S"))
print('    Duration:    ', time_stop - time_start)
print('_____________________________')