使用GPU在anaconda提示符下运行python脚本



我正在尝试从anaconda提示符运行以下python脚本:

python object_tracker.py --video test.mp4 --model yolov4 --dont_show

这直接来自AI Guys yolov4深度排序存储库(https://github.com/theAIGuysCode/yolov4-deepsort)。这是一个对象跟踪器的代码,所以它的计算量很大,在更长的视频上运行需要几个小时。我可以使用一台带有Nvidia显卡的电脑,我想使用GPU来加快进程,但我不确定如何。

Python OpenCV使用NumPy进行计算,NumPy在CPU上运行。您可以将NumPy数组转换为Pytorch张量,并可以在GPU上运行代码。一个简单的想法是


N = 8000
np.random.seed(42)
nA = np.random.rand(N,N).astype(np.float32)
nB = np.random.rand(N,N).astype(np.float32)
nC = nA.dot(nB) # numpy dot product runs on CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'PyTorch set to {device}') # cuda when available
tA = torch.from_numpy(nA).float().to(device)
tB = torch.from_numpy(nB).to(device)
tC = torch.mm(tA,tB) # torch matrix multiplication aka dot product runs on GPU

PS:

刚刚发现这个

https://medium.com/swlh/understanding-torchvision-functionalities-for-pytorch-391273299dc9

看来你可以用torchvision做得更好。

最新更新