帮我编写一个测试GPU 推理的python脚本
时间: 2024-10-23 07:05:43 浏览: 20
tensorflow-gpu下载、安装、测试
在Python中,我们可以使用PyTorch或TensorFlow等深度学习库来进行GPU推理的测试。这里是一个简单的例子,我们将演示如何使用PyTorch:
```python
# 首先确保已经安装了torch和cuda相关的库
import torch
from torch import nn, cuda
# 检查是否支持CUDA
if not torch.cuda.is_available():
print("当前系统不支持CUDA")
else:
print("已检测到可用的CUDA设备")
# 定义一个简单的模型,假设这是一个卷积神经网络
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.conv = nn.Conv2d(1, 1, kernel_size=3)
def forward(self, x):
return self.conv(x)
# 创建模型并移至GPU
model = SimpleModel().to('cuda')
# 准备一个随机数据张量(假定输入是4D Tensor)
x = torch.randn((1, 1, 8, 8)).to('cuda')
# 执行前向传播
output = model(x)
print("Output on GPU:", output.shape)
# 计算时间(为了演示性能差异)
start_time = torch.cuda.Event(enable_timing=True)
end_time = torch.cuda.Event(enable_timing=True)
start_time.record()
output = model(x)
end_time.record()
# 等待事件计时结束
torch.cuda.synchronize()
time_taken = start_time.elapsed_time(end_time) / 1e3
print(f"GPU推理时间 (ms): {time_taken:.6f}")
阅读全文