pytorch测试代码

pytorch测试代码

---

使用测试pytorch是否正常调用CUDA

1.测试代码

代码如下（示例）：

# Summary: 检测当前Pytorch和设备是否支持CUDA和cudnn

import torch
 
if __name__ == '__main__':
	print("Support CUDA ?: ", torch.cuda.is_available())
	x = torch.Tensor([1.0])
	xx = x.cuda()
	print(xx)
 
	y = torch.randn(2, 3)
	yy = y.cuda()
	print(yy)
 
	zz = xx + yy
	print(zz)
 
	# CUDNN TEST
	from torch.backends import cudnn
	print("Support cudnn ?: ",cudnn.is_acceptable(xx))

2.输出如下结果

结果：

Support CUDA ?:  True
tensor([1.], device='cuda:0')
tensor([[ 1.3227,  0.6754, -0.3934],
        [ 0.7802, -0.3730,  1.8654]], device='cuda:0')
tensor([[2.3227, 1.6754, 0.6066],
        [1.7802, 0.6270, 2.8654]], device='cuda:0')
Support cudnn ?:  True

---

总结

若输出上面所示结果，则说明pytorch成功调用cuda。