Pytorch在cuda、AMD DirectML和AMD CPU下性能比较

一、测试环境

CUDA环境: i7-8550u + 16G DDR4 2133MHz + nVidia MX150 2GB

AMD DirectML环境: Ryzen 5 5600G + 32G DDR4 3200MHz + Vega7 4GB

AMD 纯CPU环境：Ryzen 5 5600G + 32G DDR4 3200MHz 

其他硬件配置的硬盘、电源均一致。Pytorch版本为2.0.0，Python环境为3.7.11，Win10 LTSC。

二、测试代码

拟合一个100万点数的函数，并计算从神经网络被传入内存/显存开始，到计算结果出来，所耗费的时间。不含前面准备时间、出图时间。计算三次手动记录平均值。代码如下：

CUDA测试代码

# -*- coding: utf-8 -*-
# @Time    : 19/12/9 16:38
# @Author  : JL
# @File    : pytorchTest.py
# @Software: PyCharmimport matplotlib.pyplot as plt
import torch
import timex = torch.unsqueeze(torch.linspace(-1, 1, 1000000), dim=1).cuda()
y = x.pow(2) + 0.3 * torch.rand(x.size()).cuda()net1 = torch.nn.Sequential(torch.nn.Linear(1, 10),torch.nn.ReLU(),torch.nn.Linear(10, 1)
)
optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)
lossFunc = torch.nn.MSELoss()device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("当前使用的设备是：" + str(torch.cuda.get_device_name(torch.cuda.current_device())))
print("当前CUDA、CUDNN版本号分别为："+str(torch.version.cuda)+"、"+str(torch.backends.cudnn.version()))
print("当前Pytorch版本号为："+str(torch.__version__))startTime = time.perf_counter()net1.to(device)for t in range(100):prediction = net1(x)loss = lossFunc(prediction, y)optimizer.zero_grad()loss.backward()optimizer.step()print(loss.data.cpu().numpy())endTime = time.perf_counter()
delta = endTime-startTimeprint("Treat a net in %0.2f s." % delta)plt.scatter(x.data.cpu().numpy(), y.data.cpu().numpy())
plt.show()

DirectML、AMD CPU测试代码：

# -*- coding: utf-8 -*-
# @Time    : 19/12/9 16:38
# @Author  : JL
# @File    : pytorchTest.py
# @Software: PyCharmimport matplotlib.pyplot as plt
import torch
import torch_directml
import timedml = torch_directml.device()  # 如果使用DirectML，则分配到dml上
cpuML = torch.device("cpu")  # 如果仅使用CPU，则选择分配到cupML上# 注意修改dml或cpuML
x = torch.unsqueeze(torch.linspace(-1, 1, 1000000), dim=1).to(dml)
y = x.pow(2) + 0.3 * torch.rand(x.size()).to(dml)net1 = torch.nn.Sequential(torch.nn.Linear(1, 10),torch.nn.ReLU(),torch.nn.Linear(10, 1)
)
optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)
lossFunc = torch.nn.MSELoss()print("当前Pytorch版本号为：" + str(torch.__version__))startTime = time.perf_counter()net1.to(dml)  # 注意修改dml或cpuMLfor t in range(100):prediction = net1(x)loss = lossFunc(prediction, y)optimizer.zero_grad()loss.backward()optimizer.step()print(loss.data.cpu().numpy())endTime = time.perf_counter()
delta = endTime - startTimeprint("Treat a net in %0.2f s." % delta)plt.scatter(x.data.cpu().numpy(), y.data.cpu().numpy())
plt.show()

三、测试结论

测试类型	耗费时间（秒，越小越好）
CUDA	3.57
DirectML	4.48
纯CPU	5.31

看起来DirectML有点加速效果，但是还是和CUDA有差距，更何况这个是笔记本上最弱的MX150显卡的CUDA。微软要加油了。另外AMD的CPU，还是安心打游戏好了。