这是我的代码(模拟前馈神经网络):
import torch
import time
print(torch.cuda.is_available()) # True
device = torch.device('cuda:0' )
a = torch.tensor([1,2,3,4,5,6]).float().reshape(-1,1)
w1 = torch.rand(120,6)
w2 = torch.rand(1,120)
b1 = torch.rand(120,1)
b2 = torch.rand(1,1).reshape(1,1)
start = time.time()
for _ in range(100000):
ans = torch.mm(w2, torch.mm(w1,a)+b1)+b2
end = time.time()
print(end-start) # 1.2725720405578613 seconds
a = a.to(device)
w1 = w1.to(device)
w2 = w2.to(device)
b1 = b1.to(device)
b2 = b2.to(device)
start = time.time()
for _ in range(100000):
ans = torch.mm(w2, torch.mm(w1,a)+b1)+b2
end = time.time()
print(end-start) # 5.6569812297821045 seconds
我不确定我是不是做错了什么,如何更改代码以显示GPU在矩阵乘法上比CPU更快?
torch.cuda.synchronize()
。 - Hossein