当前位置：首页 > news >正文

PyTorch CUDA GPU高占用测试

news 2025/6/29 6:08:00

0x00 问题描述

安装完成PyTorch、CUDA后，验证PyTorch是否能够通过CUDA高占用GPU（占用>95%），特地使用以下代码测试。

0x01 代码设计

这个代码会持续执行神经网络的训练任务，每次循环都进行前向传播、反向传播和参数更新，以保持高强度的GPU占用。

## CUDA - GPU 占用测试
## 正确运行结果为：GPU占用显著提高（>95,NVIDIA 3060 LAPTOP）import torch
import torch.nn as nn
import torch.optim as optim# 检查CUDA是否可用
if torch.cuda.is_available():device = torch.device("cuda")print("CUDA is available. Using GPU.")
else:raise Exception("CUDA is not available. Please ensure you have a GPU.")# 创建一个简单的神经网络
class SimpleNet(nn.Module):def __init__(self):super(SimpleNet, self).__init__()self.fc1 = nn.Linear(10000, 10000)  # 大规模线性层，可以根据需要调整大小def forward(self, x):x = self.fc1(x)return xnet = SimpleNet().to(device)# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)# 创建一个大型随机输入张量
batch_size = 32
input_data = torch.randn(batch_size, 10000, device=device)# 持续执行神经网络训练任务以保持高占用率
try:while True:# 正向传播output = net(input_data)loss = criterion(output, input_data)# 反向传播和优化optimizer.zero_grad()loss.backward()optimizer.step()
except KeyboardInterrupt:print("Stopped by user.")# 释放GPU资源
net = None
torch.cuda.empty_cache()