如何在 torch.nn.Conv2d() function 中使用 GPU

Question

我不确定这个计算卷积的 function 是否使用 GPU，如果没有，如何修改它以使用 GPU？

import torch
import time

cuda = torch.device('cuda')
x = torch.randn(1,16,4,4)
t0 = time.time()
conv = torch.nn.Conv2d(16,1,(3,3))
res = conv(x) 
torch.cuda.synchronize()
t1 = time.time()

print(t0, t1, t1-t0,"sec")
print(res.shape) 

Answer 1

您可以在代码开头键入此命令：

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

如果它打印“Using cuda device”表示它正在使用 GPU