显存占用 显存测试

目录

显存测试

显存占用示例

一个模型多卡占用

---

显存测试

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显存

import timeimport 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_btime.sleep(0.01)index+=1print(index)

一个模型多卡占用

import timeimport 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 += 1print(index)