跟踪一个Pytorch Module在训练过程中的内存分配情况

2024-06-01 21:12
文章标签 训练 内存 过程 pytorch module 跟踪 分配情况

本文主要是介绍跟踪一个Pytorch Module在训练过程中的内存分配情况，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

跟踪一个Pytorch Module在训练过程中的内存分配情况

  • 代码
  • 输出

目的:跟踪一个Pytorch Module在训练过程中的内存分配情况
方法:
1.通过pre_hook module的来区分module的边界
2.通过__torch_dispatch__拦截所有的aten算子,计算在该算子中新创建tensor的总内存占用量
3.通过tensor.data_ptr()为tensor去重,表示一块独立的内存

代码


import numpy as np
import torch
from torch.nn import Module, Linear
import torch.nn as nn
from torch.optim import Adam,SGD
from torch.utils._python_dispatch import TorchDispatchMode
from dataclasses import dataclass
from typing import Any
import time@dataclass
class _ProfilerState:cls: Anyobject: Any = Nonecurrent_module=None
tesor_cache=set()def get_current_mem():global current_moduleprint(f'[INFO]{current_module["name"]}:{np.sum(current_module["size"])}')current_module=Noneclass InputDescriptor:def __init__(self) -> None:self.total_input_size=0def _save_var(self,v):class_name=v.__class__.__name__if class_name in ["Tensor","Parameter"]:global tesor_cachetensorid=v.data_ptr()if v.device.type!="cuda":return            if tensorid not in tesor_cache:tesor_cache.add(tensorid)sz=v.numel()*v.element_size()print(v.shape,v.dtype)self.total_input_size += szif class_name=="Parameter" and v.grad is not None:                tensorid=v.grad.data_ptr()if tensorid not in tesor_cache:tesor_cache.add(tensorid)sz=v.grad.numel()*v.grad.element_size()print("grad",v.grad.shape,v.grad.dtype)self.total_input_size += szelif class_name in ["list","tuple"]:for t in v:self._save_var(t)else:passdef save_vars(self,ret,*args,**kwargs):for arg in args:self._save_var(arg)        for k,v in kwargs.items():self._save_var(v)self._save_var(ret)global current_module        if current_module is None:current_module={"name":"Other","size":[]}current_module["size"].append(self.total_input_size)# 对象和类名缓存
object_cache = {}
class_name_count = {}def get_unique_name(class_name, obj_id):# 生成唯一的对象名称if class_name not in class_name_count:class_name_count[class_name] = 0uid = f"{class_name}_{obj_id}"if uid not in object_cache:class_name_count[class_name] += 1object_cache[uid] = {"idx": class_name_count[class_name]}return f'{class_name}-{object_cache[uid]["idx"]}'def initialize_module_attributes(module):# 初始化模块属性if not hasattr(module, 'uuid'):module.uuid = get_unique_name(module.__class__.__name__, id(module))if not hasattr(module, 'backward_mem'):module.backward_mem = []if not hasattr(module, 'forward_mem'):module.forward_mem = []def pre_backward_hook(module, grad_input):# 反向传播前的钩子函数initialize_module_attributes(module)global current_moduleif current_module is not None and np.sum(current_module["size"])>0:print(f'[INFO]{current_module["name"]}:{np.sum(current_module["size"])}')module.backward_mem.clear()current_module={"name":f"backward-{module.uuid}","size":module.backward_mem}def post_backward_hook(module, grad_input, grad_output):# 反向传播后的钩子函数initialize_module_attributes(module)def pre_forward_hook(module, input):# 前向传播前的钩子函数initialize_module_attributes(module)global current_moduleif current_module is not None and np.sum(current_module["size"])>0:print(f'[INFO]{current_module["name"]}:{np.sum(current_module["size"])}')module.forward_mem.clear()current_module={"name":f"forward-{module.uuid}","size":module.forward_mem}def post_forward_hook(module, input, output):# 前向传播后的钩子函数initialize_module_attributes(module)def register_forward_hooks(module):# 注册反向传播钩子module.register_forward_pre_hook(pre_forward_hook)module.register_forward_hook(post_forward_hook)def register_backward_hooks(module):# 注册反向传播钩子module.register_full_backward_pre_hook(pre_backward_hook)module.register_full_backward_hook(post_backward_hook)class HookModel(object):def __init__(self, model):output_dict = {}self.get_submodule_recrusicve(model, "", output_dict)for name, module in output_dict.items():if name.endswith("Sequential"):continueregister_forward_hooks(module)register_backward_hooks(module)def get_submodule_recrusicve(self,module, prefix, output_dict):prefix = prefix + "/" + type(module).__name__output_dict[prefix] = modulefor name, submodule in module.named_children():self.get_submodule_recrusicve(submodule, f"{prefix}[{name}]", output_dict)class TorchDumpDispatchMode(TorchDispatchMode):def __init__(self,parent):super().__init__()self.parent=parentdef __torch_dispatch__(self, func, types, args=(), kwargs=None):if kwargs is None:kwargs = {}  ret= func(*args, **kwargs)desc=InputDescriptor()desc.save_vars(ret,*args,**kwargs)if desc.total_input_size>0:print(f"{func.__name__}:{desc.total_input_size}")return retclass TorchDebugDumper:_CURRENT_Dumper = Nonedef __init__(self):self.p= _ProfilerState(TorchDumpDispatchMode)def __enter__(self):assert TorchDebugDumper._CURRENT_Dumper is NoneTorchDebugDumper._CURRENT_Dumper = selfif self.p.object is None:o = self.p.cls(self)o.__enter__()self.p.object = oelse:self.p.object.step()return selfdef __exit__(self, exc_type, exc_val, exc_tb):TorchDebugDumper._CURRENT_Dumper = Noneif self.p.object is not None:self.p.object.__exit__(exc_type, exc_val, exc_tb)del self.p.objectclass FeedForward(Module):def __init__(self,hidden_size,ffn_size):super().__init__()self.fc = nn.Sequential(Linear(in_features=hidden_size, out_features=ffn_size,bias=False),nn.ReLU(),Linear(in_features=ffn_size, out_features=ffn_size*2,bias=False),nn.Dropout(0.5),Linear(in_features=ffn_size*2, out_features=hidden_size,bias=False),)self.norm = nn.LayerNorm(normalized_shape=hidden_size, elementwise_affine=False)def forward(self, x):return x + self.fc(self.norm(x))def main():model=FeedForward(100,128)model=model.float().cuda()model.train()obj=HookModel(model)global current_modulewith TorchDebugDumper():opt=Adam(model.parameters(),lr=0.001)input=torch.randn(1,100).float().cuda()output=model(input)get_current_mem()loss=-torch.log(output.sum())opt.zero_grad()loss.backward()get_current_mem()current_module=Noneopt.step()    get_current_mem()num_model_params = sum(p.numel() for p in model.parameters())print(f"[INFO]Number of model parameters: {num_model_params}")
main()

输出

torch.Size([1, 100]) torch.float32
_to_copy.default:400
[INFO]Other:400
torch.Size([1, 100]) torch.float32
torch.Size([1, 1]) torch.float32
torch.Size([1, 1]) torch.float32
native_layer_norm.default:408
[INFO]forward-LayerNorm-1:408
torch.Size([128, 100]) torch.float32
t.default:51200
[INFO]forward-Linear-1:51200
torch.Size([256, 128]) torch.float32
t.default:131072
torch.Size([1, 256]) torch.float32
mm.default:1024
[INFO]forward-Linear-2:132096
torch.Size([1, 256]) torch.float32
native_dropout.default:1024
[INFO]forward-Dropout-1:1024
torch.Size([100, 256]) torch.float32
t.default:102400
torch.Size([1, 100]) torch.float32
add.Tensor:400
[INFO]forward-Linear-3:102800
torch.Size([]) torch.float32
log.default:4
torch.Size([]) torch.float32
neg.default:4
torch.Size([]) torch.float32
neg.default:4
torch.Size([]) torch.float32
div.Tensor:4
[INFO]Other:16
torch.Size([100, 256]) torch.float32
mm.default:102400
torch.Size([1, 256]) torch.float32
mm.default:1024
[INFO]backward-Linear-3:103424
torch.Size([128, 100]) torch.float32
mm.default:51200
[INFO]backward-Linear-1:51200
torch.Size([128, 100]) torch.float32
zeros_like.default:51200
torch.Size([128, 100]) torch.float32
zeros_like.default:51200
torch.Size([256, 128]) torch.float32
zeros_like.default:131072
torch.Size([256, 128]) torch.float32
zeros_like.default:131072
torch.Size([100, 256]) torch.float32
zeros_like.default:102400
torch.Size([100, 256]) torch.float32
zeros_like.default:102400
torch.Size([128, 100]) torch.float32
torch.Size([256, 128]) torch.float32
torch.Size([100, 256]) torch.float32
_foreach_sqrt.default:284672
[INFO]Other:854016
[INFO]Number of model parameters: 71168

这篇关于跟踪一个Pytorch Module在训练过程中的内存分配情况的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！


http://www.chinasem.cn/article/1022134。

相关文章

浅析Spring Security认证过程

浅析Spring Security认证过程

类图 为了方便理解Spring Security认证流程，特意画了如下的类图，包含相关的核心认证类 概述 核心验证器 AuthenticationManager 该对象提供了认证方法的入口，接收一个Authentiaton对象作为参数; public interface AuthenticationManager {Authentication authenticate(Authenti
阅读更多...
NameNode内存生产配置

NameNode内存生产配置

Hadoop2.x 系列，配置 NameNode 内存 NameNode 内存默认 2000m ，如果服务器内存 4G ， NameNode 内存可以配置 3g 。在 hadoop-env.sh 文件中配置如下。 HADOOP_NAMENODE_OPTS=-Xmx3072m Hadoop3.x 系列，配置 Nam
阅读更多...
作业提交过程之HDFSMapReduce

作业提交过程之HDFSMapReduce

作业提交全过程详解 （1）作业提交 第1步：Client调用job.waitForCompletion方法，向整个集群提交MapReduce作业。 第2步：Client向RM申请一个作业id。 第3步：RM给Client返回该job资源的提交路径和作业id。 第4步：Client提交jar包、切片信息和配置文件到指定的资源提交路径。 第5步：Client提交完资源后，向RM申请运行MrAp
阅读更多...
【机器学习】高斯过程的基本概念和应用领域以及在python中的实例

【机器学习】高斯过程的基本概念和应用领域以及在python中的实例

引言 高斯过程（Gaussian Process，简称GP）是一种概率模型，用于描述一组随机变量的联合概率分布，其中任何一个有限维度的子集都具有高斯分布 文章目录 引言一、高斯过程1.1 基本定义1.1.1 随机过程1.1.2 高斯分布 1.2 高斯过程的特性1.2.1 联合高斯性1.2.2 均值函数1.2.3 协方差函数（或核函数） 1.3 核函数1.4 高斯过程回归（Gauss
阅读更多...
MiniGPT-3D, 首个高效的3D点云大语言模型，仅需一张RTX3090显卡，训练一天时间，已开源

MiniGPT-3D, 首个高效的3D点云大语言模型，仅需一张RTX3090显卡，训练一天时间，已开源

项目主页：https://tangyuan96.github.io/minigpt_3d_project_page/ 代码：https://github.com/TangYuan96/MiniGPT-3D 论文：https://arxiv.org/pdf/2405.01413 MiniGPT-3D在多个任务上取得了SoTA，被ACM MM2024接收，只拥有47.8M的可训练参数，在一张RTX
阅读更多...
Solr 使用Facet分组过程中与分词的矛盾解决办法

Solr 使用Facet分组过程中与分词的矛盾解决办法

对于一般查询而言  ,  分词和存储都是必要的  .  比如  CPU  类型  ”Intel  酷睿  2  双核  P7570”,  拆分成  ”Intel”,”  酷睿  ”,”P7570”  这样一些关键字并分别索引  ,  可能提供更好的搜索体验  .  但是如果将  CPU  作为 Facet  字段  ,  最好不进行分词  .  这样就造成了矛盾  ,  解决方法
阅读更多...
JVM内存调优原则及几种JVM内存调优方法

JVM内存调优原则及几种JVM内存调优方法

JVM内存调优原则及几种JVM内存调优方法 1、堆大小设置。 2、回收器选择。 　　1、在对JVM内存调优的时候不能只看操作系统级别Java进程所占用的内存，这个数值不能准确的反应堆内存的真实占用情况，因为GC过后这个值是不会变化的，因此内存调优的时候要更多地使用JDK提供的内存查看工具，比如JConsole和Java VisualVM。 　　2、对JVM内存的系统级的调优主要的目的是减少
阅读更多...
JVM 常见异常及内存诊断

JVM 常见异常及内存诊断

栈内存溢出 栈内存大小设置：-Xss size 默认除了window以外的所有操作系统默认情况大小为 1MB，window 的默认大小依赖于虚拟机内存。 栈帧过多导致栈内存溢出 下述示例代码，由于递归深度没有限制且没有设置出口，每次方法的调用都会产生一个栈帧导致了创建的栈帧过多，而导致内存溢出(StackOverflowError)。 示例代码： 运行结果： 栈帧过大导致栈内存
阅读更多...
Spark MLlib模型训练—聚类算法 PIC(Power Iteration Clustering)

Spark MLlib模型训练—聚类算法 PIC(Power Iteration Clustering)

Spark MLlib模型训练—聚类算法 PIC(Power Iteration Clustering) Power Iteration Clustering (PIC) 是一种基于图的聚类算法，用于在大规模数据集上进行高效的社区检测。PIC 算法的核心思想是通过迭代图的幂运算来发现数据中的潜在簇。该算法适用于处理大规模图数据，特别是在社交网络分析、推荐系统和生物信息学等领域具有广泛应用。Spa
阅读更多...
理解java虚拟机内存收集

理解java虚拟机内存收集

学习《深入理解Java虚拟机》时个人的理解笔记 1、为什么要去了解垃圾收集和内存回收技术？ 当需要排查各种内存溢出、内存泄漏问题时，当垃圾收集成为系统达到更高并发量的瓶颈时，我们就必须对这些“自动化”的技术实施必要的监控和调节。 2、“哲学三问”内存收集 what？when？how？ 那些内存需要回收？什么时候回收？如何回收？ 这是一个整体的问题，确定了什么状态的内存可以
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2