帮我编写一个测试GPU 推理的python脚本

在Python中，我们可以使用PyTorch或TensorFlow等深度学习库来进行GPU推理的测试。这里是一个简单的例子，我们将演示如何使用PyTorch：

# 首先确保已经安装了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}")