-
Notifications
You must be signed in to change notification settings - Fork 0
/
cpu_vs_gpu.py
30 lines (26 loc) · 848 Bytes
/
cpu_vs_gpu.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import torch
import numpy as np
import time
if torch.backends.mps.is_available():
mps_device = torch.device("mps")
print("MPS device found.")
else:
print("MPS device not found.")
exit()
torch_rand1 = torch.rand(100, 100, 100, 100).to(mps_device)
torch_rand2 = torch.rand(100, 100, 100, 100).to(mps_device)
np_rand1 = torch.rand(100, 100, 100, 100)
np_rand2 = torch.rand(100, 100, 100, 100)
# GPU (Metal).
start_time = time.time()
# @ symbol is the matrix multiplication operator in pytorch.
rand = (torch_rand1 @ torch_rand2)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"GPU elapsed time: {elapsed_time:.10f} seconds")
# CPU.
start_time = time.time()
rand = np.multiply(np_rand1, np_rand2)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"CPU elapsed time: {elapsed_time:.10f} seconds")