Python函数内部与函数外部执行相同语句的显存区别

一只积极向上的小咸鱼2024-12-01 5:06

执行代码

python 复制代码 
mport torch
import torch.cuda
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 参数设置
B = 64  # batch size
L = 32  # sequence length
C = 512  # embedding dimension
H = 8  # number of heads
D = C // H  # head dimension
# 创建随机张量
q = torch.randn(B, H, L, D).to(device)
k = torch.randn(B, H, D, L).to(device)
v = torch.randn(B, H, L, D).to(device)
x = torch.randn(B, L, C).to(device)
# 记录当前显存使用
def fa( x):
    print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
    prev_memory = torch.cuda.memory_allocated() 
    # 执行矩阵乘法
    attn = (q @ k) * 0.125  # 假设 self.scale = 0.125
    current_memory = torch.cuda.memory_allocated()
    memory_change = (current_memory - prev_memory) / 1024**2
    print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
    prev_memory = current_memory  # 更新 prev_memory
    if True:
        # 执行最终的矩阵乘法和重新整形
        x = (attn @ v).transpose(1, 2).reshape(B, L, C)
        current_memory = torch.cuda.memory_allocated()
        memory_change = (current_memory - prev_memory) / 1024**2
        print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
fa(x)
current_memory = torch.cuda.memory_allocated()
print(f"final : {current_memory / 1024**2:.2f} MB")

结果为

bash 复制代码 
Initial Memory: 16.00 MB
After matmul: 18.00 MB, Change: 2.00 MB
After final matmul and reshape: 22.00 MB, Change: 4.00 MB
final : 16.00 MB

但是执行代码

python 复制代码 
import torch
import torch.cuda
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 参数设置
B = 64  # batch size
L = 32  # sequence length
C = 512  # embedding dimension
H = 8  # number of heads
D = C // H  # head dimension
# 创建随机张量
q = torch.randn(B, H, L, D).to(device)
k = torch.randn(B, H, D, L).to(device)
v = torch.randn(B, H, L, D).to(device)
x = torch.randn(B, L, C).to(device)
print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
prev_memory = torch.cuda.memory_allocated() 
# 执行矩阵乘法
attn = (q @ k) * 0.125  # 假设 self.scale = 0.125
current_memory = torch.cuda.memory_allocated()
memory_change = (current_memory - prev_memory) / 1024**2
print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
prev_memory = current_memory  # 更新 prev_memory
if True:
    # 执行最终的矩阵乘法和重新整形
    x = (attn @ v).transpose(1, 2).reshape(B, L, C)
    current_memory = torch.cuda.memory_allocated()
    memory_change = (current_memory - prev_memory) / 1024**2
    print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")

结果为

bash 复制代码 
Initial Memory: 16.00 MB
After matmul: 18.00 MB, Change: 2.00 MB
After final matmul and reshape: 18.00 MB, Change: 0.00 MB

主要涉及 PyTorch 的显存管理机制 和 Python 的作用域规则。

1. 作用域与变量生命周期：

在 Python 中，变量的生命周期受到作用域的影响：

  • 不在函数中 时：
    • 当执行 x = (attn @ v).transpose(1, 2).reshape(B, L, C) 时，旧的 x 会被立即覆盖，旧的 x 的引用计数变为 0，显存会被立即释放或加入缓存。
  • 在函数中 时：
    • 局部变量（如 attn, v 等）在函数执行完之前不会被释放，即使 x 被重新赋值，中间张量（如 attn @ v）依然在显存中占用空间。
    • 这些局部变量直到函数返回后才会被 Python 的垃圾回收机制回收，导致显存未及时释放，从而造成额外的显存占用。
2. PyTorch 的显存缓存机制：

PyTorch 使用 CUDA 显存缓存机制来优化显存分配，但有以下行为特征：

  • 不在函数中时 ：
    • PyTorch 可能更积极地释放未引用的张量。
  • 在函数中时 ：
    • 中间结果的显存使用可能会被缓存，直到函数执行完毕，导致显存使用看起来增加了。
3. 临时张量未释放：

在函数中，PyTorch 可能会生成一些 临时张量，这些临时张量通常在 PyTorch 后台管理，直到函数执行完毕后才会释放。因此：

  • 不在函数中时，这些张量更早地被回收。
  • 在函数中时，临时张量的显存延迟释放，导致显存增加。
5 函数中内存管理
python 复制代码 
def fa( x):
    print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
    prev_memory = torch.cuda.memory_allocated() 
    x + 1
    current_memory = torch.cuda.memory_allocated()
    memory_change = (current_memory - prev_memory) / 1024**2
    print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
    prev_memory = current_memory
    # 执行矩阵乘法
    attn = (q @ k) * 0.125  # 假设 self.scale = 0.125
    current_memory = torch.cuda.memory_allocated()
    memory_change = (current_memory - prev_memory) / 1024**2
    print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
    prev_memory = current_memory  # 更新 prev_memory
    if True:
        # 执行最终的矩阵乘法和重新整形
        x = (attn @ v).transpose(1, 2).reshape(B, L, C)
        current_memory = torch.cuda.memory_allocated()
        memory_change = (current_memory - prev_memory) / 1024**2
        print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
fa(x)
current_memory = torch.cuda.memory_allocated()
print(f"final : {current_memory / 1024**2:.2f} MB")

结果

bash 复制代码 
Initial Memory: 16.00 MB
After matmul: 16.00 MB, Change: 0.00 MB
After matmul: 18.00 MB, Change: 2.00 MB
After final matmul and reshape: 22.00 MB, Change: 4.00 MB
final : 16.00 MB
python 复制代码 
def fa( x):
    print(f"Initial Memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
    prev_memory = torch.cuda.memory_allocated() 
    x = x + 1
    current_memory = torch.cuda.memory_allocated()
    memory_change = (current_memory - prev_memory) / 1024**2
    print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
    prev_memory = current_memory
    # 执行矩阵乘法
    attn = (q @ k) * 0.125  # 假设 self.scale = 0.125
    current_memory = torch.cuda.memory_allocated()
    memory_change = (current_memory - prev_memory) / 1024**2
    print(f"After matmul: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
    prev_memory = current_memory  # 更新 prev_memory
    if True:
        # 执行最终的矩阵乘法和重新整形
        x = (attn @ v).transpose(1, 2).reshape(B, L, C)
        current_memory = torch.cuda.memory_allocated()
        memory_change = (current_memory - prev_memory) / 1024**2
        print(f"After final matmul and reshape: {current_memory / 1024**2:.2f} MB, Change: {memory_change:.2f} MB")
fa(x)
current_memory = torch.cuda.memory_allocated()
print(f"final : {current_memory / 1024**2:.2f} MB")
bash 复制代码 
Initial Memory: 16.00 MB
After matmul: 20.00 MB, Change: 4.00 MB
After matmul: 22.00 MB, Change: 2.00 MB
After final matmul and reshape: 22.00 MB, Change: 0.00 MB
final : 16.00 MB
1. 操作与赋值的差异：

首先，理解 x = x + 1 和 x + 1 在 PyTorch 中的区别很重要。
x + 1：不会增加显存

  • x + 1 是一个普通的张量运算，它会 创建一个新的张量 作为结果，但不会修改原张量 x。
  • 然而，PyTorch 在执行这种操作时，通常会 复用已有的显存。它会将中间计算结果存放在新的内存位置，并且使用现有的显存池优化内存分配。
  • 在这种情况下，如果没有显式的赋值给 x，就不会创建额外的内存开销。
    -x = x + 1：会增加显存
    x = x + 1：会增加显存
  • x = x + 1 这一操作实际上会执行 "计算+赋值" 。这一过程中：
    1. 中间结果 （即 x + 1 计算出的新张量）会被存储到 新内存位置。
    2. 原来存储 x 的显存会被新值 覆盖 ，但是由于这是一个"计算+赋值"操作，PyTorch 为了避免覆盖数据，通常会先分配新的内存空间来存储计算结果。
    3. 这意味着，新张量 和原始张量会在一段时间内 共享显存 ，直到 Python 的垃圾回收机制清理旧张量的内存。
      在函数调用期间，新的张量（x + 1）会占用新的显存，而原来的 x 张量的内存要等到函数结束或者垃圾回收时才能释放。因此，暂时增加了显存占用。
相关推荐
IT古董18 分钟前
第四章:大模型（LLM)】06.langchain原理-(3)LangChain Prompt 用法
java·人工智能·python
fantasy_arch5 小时前
pytorch例子计算两张图相似度
人工智能·pytorch·python
AndrewHZ6 小时前
【3D重建技术】如何基于遥感图像和DEM等数据进行城市级高精度三维重建？
图像处理·人工智能·深度学习·3d·dem·遥感图像·3d重建
WBluuue7 小时前
数学建模：智能优化算法
python·机器学习·数学建模·爬山算法·启发式算法·聚类·模拟退火算法
赴3357 小时前
矿物分类案列 （一）六种方法对数据的填充
人工智能·python·机器学习·分类·数据挖掘·sklearn·矿物分类
大模型真好玩7 小时前
一文深度解析OpenAI近期发布系列大模型：意欲一统大模型江湖？
人工智能·python·mcp
RPA+AI十二工作室7 小时前
亚马逊店铺绩效巡检_影刀RPA源码解读
chrome·python·rpa·影刀
nonono8 小时前
深度学习——常见的神经网络
人工智能·深度学习·神经网络
小艳加油8 小时前
Python机器学习与深度学习；Transformer模型/注意力机制/目标检测/语义分割/图神经网络/强化学习/生成式模型/自监督学习/物理信息神经网络等
python·深度学习·机器学习·transformer
钢铁男儿9 小时前
如何构建一个神经网络？从零开始搭建你的第一个深度学习模型
人工智能·深度学习·神经网络
热门推荐
01UV安装并设置国内源02Qwen3-Coder 快速上手教程 | Qwen Code + Claude Code03【2025.08.06最新版】Android Studio下载、安装及配置记录（自动下载sdk）04KGG转MP3工具|非KGM文件|解密音频052025最新国内服务器可用docker源仓库地址大全（2025年8月更新）06蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接07NVIDIA显卡驱动、CUDA、cuDNN 和 TensorRT 版本匹配指南08TRAE 规则（Rules）配置指南：个人习惯、团队规范与最佳实践09阿里开源首个图像生成基础模型——Qwen-Image本地部署教程，超强中文渲染能力刷新SOTA！10TRAE Rules 实践：为项目配置 6A 工作流