显存占用 显存测试

目录

显存测试

显存占用示例

一个模型多卡占用


显存测试

python 复制代码
import torch

# 计算张量的大小(例如:每个 float 占用 4 字节)
# 40GB = 40 * 1024 * 1024 * 1024 字节
# 每个 float 4 字节,因此需要的 float 数量为 (40 * 1024 * 1024 * 1024) / 4
num_elements = (40 * 1024 * 1024 * 1024) // 4

# 创建一个在 GPU 上的张量
tensor = torch.empty(num_elements, dtype=torch.float32, device='cuda')

print(tensor)

显存占用示例

42G和62G显存

python 复制代码
import time

import torch

# 设置张量的大小
num_elements = (10 * 1024 * 1024 * 1024) // 4  # 42GB 大小
# num_elements = (15 * 1024 * 1024 * 1024) // 4  # 62GB 大小

# 创建两个随机数张量,存放在 GPU 上
tensor_a = torch.rand(num_elements, dtype=torch.float32, device='cuda:3')
tensor_b = torch.rand(num_elements, dtype=torch.float32, device='cuda:3')

# 创建一个用于存储结果的张量
# result_tensor = torch.empty(num_elements, dtype=torch.float32, device='cuda')
index=0
while True:
    result_tensor=tensor_a + tensor_b
    # tensor_a + tensor_b
    time.sleep(0.01)
    index+=1
    print(index)

一个模型多卡占用

python 复制代码
import time

import torch
import torch.nn as nn

# 设置张量的大小
num_elements = (6 * 1024 * 1024 * 1024) // 4  # 40GB 大小

# 确保有两个可用的 GPU
if torch.cuda.device_count() < 2:
    raise RuntimeError("至少需要两块 GPU")

# 创建两个随机数张量,存放在 GPU 上
tensor_a = torch.rand(num_elements , dtype=torch.float32, device='cuda:0')
tensor_b = torch.rand(num_elements , dtype=torch.float32, device='cuda:0')

# 创建一个用于存储结果的张量
result_tensor = torch.empty(num_elements , dtype=torch.float32, device='cuda:1')

class AddModel(nn.Module):
    def forward(self, tensor_a, tensor_b):
        return tensor_a + tensor_b

# 实例化模型并使用 DataParallel
model = AddModel().cuda()
model = nn.DataParallel(model)

index=0
# 不断相加的循环
while True:
    # 使用 DataParallel 进行加法
    result_tensor = model(tensor_a, tensor_b)

    # 将结果存储在第一个 GPU 上
    result_tensor = result_tensor.to('cuda:1')

    time.sleep(0.01)
    index += 1
    print(index)
相关推荐
Rock_yzh7 小时前
AI学习日记——参数的初始化
人工智能·python·深度学习·学习·机器学习
拆房老料8 小时前
Transformer推理优化全景:从模型架构到硬件底层的深度解析
深度学习·ai·自然语言处理·transformer
青衫客368 小时前
基于 Python 构建的安全 gRPC 服务——TLS、mTLS 与 Casbin 授权实战
python·安全·微服务
-dzk-9 小时前
【3DGS复现】Autodl服务器复现3DGS《简单快速》《一次成功》《新手练习复现必备》
运维·服务器·python·计算机视觉·3d·三维重建·三维
byzy10 小时前
【论文笔记】VisionPAD: A Vision-Centric Pre-training Paradigm for Autonomous Driving
论文阅读·深度学习·计算机视觉·自动驾驶
摩羯座-1856903059410 小时前
爬坑 10 年!京东店铺全量商品接口实战开发:从分页优化、SKU 关联到数据完整性闭环
linux·网络·数据库·windows·爬虫·python
ACERT33310 小时前
5.吴恩达机器学习—神经网络的基本使用
人工智能·python·神经网络·机器学习
韩立学长11 小时前
【开题答辩实录分享】以《基于python的奶茶店分布数据分析与可视化》为例进行答辩实录分享
开发语言·python·数据分析
2401_8315017311 小时前
Python学习之day03学习(文件和异常)
开发语言·python·学习
可触的未来,发芽的智生12 小时前
触摸未来2025.10.06:声之密语从生理构造到神经网络的声音智能革命
人工智能·python·神经网络·机器学习·架构