PyTorch CUDA GPU高占用测试

0x00 问题描述

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

0x01 代码设计

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

python 复制代码
## 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 x

net = 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()

0x02 实验结果

笔者使用的3060 Laptop GPU 占用在95%以上,代码效果显著,说明PyTorch、CUDA环境安装成功。

0x03 后记

  • No Pains, No Gains.
相关推荐
Wilber的技术分享5 分钟前
【大模型实战笔记 1】Prompt-Tuning方法
人工智能·笔记·机器学习·大模型·llm·prompt
YUELEI1186 分钟前
langchain 提示模版 PromptTemplate
python·langchain
liliangcsdn9 分钟前
结合prompt分析NodeRAG的build过程
java·服务器·人工智能·数据分析·知识图谱
东方不败之鸭梨的测试笔记9 分钟前
LangChain: Models, Prompts 模型和提示词
人工智能·python·langchain
初恋叫萱萱11 分钟前
AI驱动开发实战:基于飞算JavaAI的在线考试系统设计与实现
人工智能
yzx99101317 分钟前
图像去雾:从暗通道先验到可学习融合——一份可跑的 PyTorch 教程
人工智能·pytorch·学习
博大世界27 分钟前
解剖智驾“大脑”:一文读懂自动驾驶系统软件架构
人工智能·机器学习·自动驾驶
大熊猫侯佩31 分钟前
苹果 AI 探秘:代号 “AFM” —— “温柔的反叛者”
人工智能·sft·ai 大模型·apple 本地大模型·foundationmodel·苹果智能·applebot
AI Echoes1 小时前
别再手工缝合API了!开源LLMOps神器LMForge,让你像搭积木一样玩转AI智能体!
人工智能·python·langchain·开源·agent
AI Echoes1 小时前
从零构建企业级LLMOps平台:LMForge——支持多模型、可视化编排、知识库与安全审核的全栈解决方案
人工智能·python·langchain·开源·agent