PyTorch Tensor C++源码深度解析:TensorImpl、Storage与view/reshape底层原理
发布时间:2026/10/20 6:17:20 锦皓数字建站

1. 为什么要啃Tensor的C源码很多人用PyTorch写模型写了好几年torch.Tensor这行代码敲了无数遍但真要问一句“tensor.shape到底存在哪里”“tensor.to(device)底层做了什么”“为什么view和reshape行为不一样”能答上来的人其实不多。我自己也是在被一个内存布局相关的bug折磨了整整两天之后才下定决心去翻Tensor的C实现。那个bug的表现是同样的数据、同样的操作在某些情况下view会报错换reshape就没事但reshape之后的数据偶尔又和预期对不上。查了一圈Python层的文档全是“返回一个具有新形状的张量”这种废话根本解释不了我遇到的问题。这篇文章就是那次源码阅读的完整整理。核心目标很明确把torch.Tensor在C层面到底是怎么组织的、关键操作是怎么实现的、有哪些设计上的取舍用能看懂的话讲清楚。适合两类人一是写过PyTorch但想搞清楚底层机制的开发者二是正在做自定义算子或者框架二次开发、需要直接和Tensor C接口打交道的人。读完之后你至少能明白TensorImpl、Storage、TensorOptions这几个核心类之间的关系以及view、reshape、contiguous这些高频操作背后的真实逻辑。需要提前说明的是PyTorch的源码体量非常大一篇文章不可能覆盖所有细节。我会聚焦在Tensor最核心的数据结构和几个最常用的操作上把主干逻辑理清楚。具体的代码路径以PyTorch某个稳定版本为准不同版本之间会有差异但核心设计思想是一致的。2. TensorImpl、Storage与TensorOptions的三层结构2.1 Tensor在C里到底是个什么东西打开c10/core/TensorImpl.h你会看到TensorImpl这个类。它是Tensor在C层面的真正实体Python里那个torch.Tensor对象本质上就是通过pybind11绑定到这个类上的。TensorImpl继承自intrusive_ptr_target这意味着它的生命周期由引用计数管理而不是靠Python的垃圾回收。TensorImpl里面存了什么最关键的几个成员storage_一个Storage对象负责管理实际的内存块。sizes_和strides_描述张量的形状和步长决定了如何从内存块中索引数据。storage_offset_偏移量表示数据从Storage的哪个位置开始。dtype_数据类型比如float32、int64。device_设备信息CPU还是某个加速器。layout_内存布局通常是strided也有sparse等。这里有一个非常重要的设计TensorImpl本身不直接持有数据指针数据在Storage里。Storage才是真正管理内存的类它持有一个DataPtrDataPtr里面包含了实际的内存地址和一个deleter。这种分离设计的好处是多个Tensor可以共享同一个Storage这就是view操作能零拷贝实现的根本原因。我一开始看的时候有个疑惑为什么不把数据指针直接放在TensorImpl里后来想明白了如果数据指针在TensorImpl里那每次做view或者slice都要复制一份数据指针和相关的元信息而且引用计数的管理会变得很复杂。分离之后Storage的引用计数独立管理多个TensorImpl共享同一个Storage时只要增加Storage的引用计数就行非常干净。2.2 Storage和DataPtr的内存管理逻辑Storage的实现主要在c10/core/Storage.h和c10/core/StorageImpl.h。StorageImpl是实际干活的类它包含data_ptr_一个DataPtr对象封装了原始内存指针和释放逻辑。size_bytes_分配的字节数。allocator_指向分配器的指针CPU上通常是CPUAllocator。DataPtr的设计很有意思。它是一个包含void* data和std::functionvoid(void*) deleter的结构。deleter是一个可调用对象当DataPtr被销毁时它会调用deleter来释放内存。这种设计允许不同的分配器使用不同的释放策略比如CPU内存用free加速器内存用对应的释放接口而Storage本身不需要知道具体细节。注意Storage的size_bytes_是字节数不是元素个数。计算元素个数需要除以dtype的大小。这个细节在看源码时很容易搞混我自己就曾经把size_bytes当成元素个数来用结果算出来的shape完全不对。还有一个容易忽略的点Storage可以是未分配的。比如你创建一个空Tensor它的Storage的data_ptr可能是nullptrsize_bytes为0。这种情况下任何试图访问数据的操作都会出问题所以PyTorch在很多操作入口都会检查Tensor是否已经分配了内存。2.3 TensorOptions如何控制创建行为TensorOptions是一个用来描述“我想创建一个什么样的Tensor”的类它包含dtype、device、layout、requires_grad这几个可选字段。它的设计用了Builder模式你可以链式调用auto options torch::TensorOptions() .dtype(torch::kFloat32) .device(torch::kCPU) .requires_grad(false);TensorOptions的一个关键特性是“部分指定”。也就是说你可以只指定dtype不指定device这时候PyTorch会使用默认值。默认值从哪里来从DefaultTensorOptions这个全局配置里来。这就解释了为什么torch.tensor([1,2,3])默认创建的是float32的CPU张量——因为全局默认就是这样的。在实际使用中TensorOptions最常见的坑是类型推断。比如torch.tensor([1, 2, 3])创建的是int64而torch.tensor([1.0, 2.0, 3.0])创建的是float32。这个行为就是由TensorOptions的默认推断逻辑决定的。如果你在做自定义算子开发创建输出Tensor时一定要显式指定dtype不要依赖推断否则很容易出现类型不匹配的问题。3. view、reshape与contiguous的底层差异3.1 view为什么有时候会失败view的语义是“在不复制数据的前提下改变张量的形状”。它能不能成功取决于一个条件新的形状是否和原来的内存布局兼容。具体来说就是新的strides能否通过原来的sizes和strides推导出来。举个例子假设有一个形状为(2, 3)的连续张量strides是(3, 1)。你想把它view成(3, 2)新的strides应该是(2, 1)。这个推导是可行的因为原来的内存是连续的元素排列顺序是[a0,a1,a2,a3,a4,a5]view成(3,2)之后第一行取前两个元素第二行取中间两个第三行取最后两个完全符合内存顺序。但如果原来的张量不是连续的比如是某个大张量的转置strides可能是(1, 3)而不是(3, 1)。这时候你想view成(3, 2)PyTorch就没办法在不复制数据的情况下完成因为内存中的元素排列顺序和新形状要求的顺序不一致。这时候view就会抛出一个错误提示你“view size is not compatible with input tensors size and stride”。我在实际项目中遇到过这种情况对一个transpose之后的张量做view代码在测试集上跑得好好的换了一批数据就报错了。原因就是测试集的数据恰好让某个维度为1而维度为1的张量在stride计算上有特殊处理导致view碰巧能成功。这种bug非常隐蔽因为它的触发条件依赖于具体的数据形状。3.2 reshape做了什么额外工作reshape的语义是“尽量不复制但如果必须复制就复制”。它的实现逻辑是先尝试view如果view失败就调用contiguous得到一份连续的内存副本然后再view。所以reshape永远不会失败除非内存不够但它可能产生一次数据拷贝。这个“可能拷贝”的特性是很多性能问题的根源。如果你在一个训练循环里频繁调用reshape而输入张量恰好是非连续的那每次都会触发一次隐式的内存拷贝和分配。这个开销在大部分情况下不明显但在大规模训练或者高频调用的场景下累积起来就很可观了。提示如果你确定输入是连续的用view比reshape更安全因为它会强制你面对“不连续”这个事实而不是默默帮你拷贝。如果你不确定用reshape但要在性能分析时留意它是否触发了拷贝。3.3 contiguous的触发条件与性能影响contiguous的作用是返回一个内存连续的张量。如果原张量已经是连续的它直接返回自身增加引用计数不做任何拷贝。如果不是连续的它会分配一块新内存按照逻辑顺序把数据搬过去。判断一个张量是否连续看的是它的strides是否和“按行优先展开”的strides一致。对于一个形状为sizes的张量连续条件下的strides应该是从最后一维开始依次为1, sizes[-1], sizes[-1]*sizes[-2], ...。如果实际strides和这个不一致就是不连续的。这里有一个容易误解的地方维度为1的张量其stride可以是任意值不影响连续性判断。比如形状为(1, 3)的张量strides是(100, 1)它仍然是连续的因为第一维大小为1stride是多少都无所谓。PyTorch在判断连续性时会跳过大小为1的维度这个细节在看源码时要注意。4. 从源码看Tensor的内存分配与释放4.1 CPUAllocator的分配策略CPU上的内存分配由CPUAllocator负责实现在c10/core/CPUAllocator.cpp。它底层调用的是posix_memalignLinux或_aligned_mallocWindows保证分配的内存按64字节对齐。为什么要对齐因为现代CPU的SIMD指令比如AVX系列要求内存地址按特定边界对齐不对齐会导致性能下降甚至程序崩溃。CPUAllocator的allocate方法返回一个DataPtrdeleter是free。这里有一个细节PyTorch在分配内存时如果请求的大小为0会返回一个特殊的空指针而不是调用malloc(0)。这是因为malloc(0)的行为在不同平台上不一致有的返回nullptr有的返回一个可以free的指针处理起来很麻烦。我在调试一个内存泄漏问题时曾经追踪到CPUAllocator的分配记录。PyTorch提供了一个环境变量PYTORCH_NO_CUDA_MEMORY_CACHING虽然名字里有CUDA但对CPU分配也有影响可以用来关闭内存缓存方便定位问题。不过这个变量在生产环境不要开会严重影响性能。4.2 引用计数与生命周期管理TensorImpl和StorageImpl都继承自intrusive_ptr_target使用侵入式引用计数。这意味着引用计数的计数器直接存在对象内部而不是像shared_ptr那样存在控制块里。侵入式引用计数的好处是少一次内存间接访问性能更好坏处是对象必须从intrusive_ptr_target继承灵活性差一些。引用的增加和减少通过intrusive_ptr这个智能指针来管理。当你把一个Tensor赋值给另一个变量时引用计数加1当变量离开作用域时引用计数减1减到0时对象被销毁。Storage的引用计数独立于TensorImpl所以即使所有TensorImpl都销毁了只要还有Storage的引用内存就不会被释放。这个机制解释了为什么有时候你删除了一个Tensor内存却没有立即释放。很可能是因为还有另一个Tensor共享着同一个Storage或者某个地方还持有Storage的引用。在Python层面这种情况通常和循环引用或者全局缓存有关。4.3 内存复用与缓存机制PyTorch在CPU侧有一个简单的内存缓存机制主要是为了减少频繁的小块内存分配开销。具体来说CPUAllocator会维护一个空闲块列表当释放内存时不是立即还给操作系统而是放到空闲列表里下次分配时如果大小匹配就直接复用。这个机制在训练循环里非常有用因为每一轮迭代都会创建和销毁大量临时张量如果每次都向操作系统申请和释放内存开销会很大。但这也带来一个问题内存占用看起来比实际需要的高。你可能会发现程序占用的内存一直不降即使已经删除了所有Tensor。这通常不是内存泄漏而是缓存没有释放。如果你确实需要强制释放缓存可以调用torch.cuda.empty_cache()对加速器内存或者手动触发CPU侧的缓存清理。不过CPU侧的缓存清理接口不太稳定不同版本行为不一样生产环境不建议依赖。5. 自定义算子开发中的Tensor C实践5.1 如何正确创建输出Tensor在自定义算子中创建输出Tensor最推荐的方式是用torch::empty配合TensorOptionsauto output torch::empty({batch, channels, height, width}, input.options().dtype(torch::kFloat32));这里用input.options()继承输入张量的device和layout然后显式指定dtype。为什么不直接torch::empty({...}, input.options())因为如果输入是int64而你需要float32输出直接继承dtype会出问题。显式指定dtype是一个好习惯能避免很多类型相关的bug。另一个常见错误是用torch::zeros创建输出。zeros会多做一次内存清零操作如果你的算子会覆盖所有输出元素这次清零就是纯浪费。用empty分配未初始化内存然后直接写入性能更好。5.2 处理非连续输入的策略自定义算子经常会遇到非连续的输入。比如上游做了一个transpose传给你的张量strides是乱的。这时候你有两个选择在算子入口调用input.contiguous()强制转成连续然后再处理。在kernel里直接处理非连续的情况通过strides来索引。第一种方案简单但有一次拷贝开销。第二种方案性能好但kernel逻辑会复杂很多而且容易出错。我的经验是如果算子的计算密度高比如矩阵乘法拷贝的开销相对可以忽略用第一种方案如果算子是memory-bound的比如element-wise操作拷贝开销占比大值得花时间实现第二种方案。注意调用contiguous()之后一定要用返回的新张量而不是原来的。我见过有人在代码里写了input.contiguous()但没接收返回值结果后面用的还是非连续的input性能没提升逻辑还可能出错。5.3 在算子中正确使用view和reshape在C层面view和reshape的接口和Python层一致。但有一个区别C的view返回的是Tensor如果失败会抛异常而Python层可以通过try-except捕获。在自定义算子中如果你不确定输入是否连续用reshape更安全。但要注意reshape在C里的开销比Python层更值得关注因为C算子通常被设计成高性能的一次隐式的内存拷贝可能让整个算子的性能不达标。我的做法是在算子入口先检查input.is_contiguous()如果是连续的用view如果不是根据算子特性决定是报错还是走contiguous路径。torch::Tensor my_op(torch::Tensor input) { TORCH_CHECK(input.is_contiguous(), input must be contiguous); auto flat input.view({-1}); // ... 后续处理 }这种显式检查的好处是问题会在算子入口就暴露出来而不是在某个深层调用里报一个难以理解的错误。6. 源码阅读中容易踩的坑与调试技巧6.1 宏定义与模板带来的阅读障碍PyTorch的C代码大量使用了宏和模板这是阅读时最大的障碍。比如AT_DISPATCH_FLOATING_TYPES这个宏展开后会生成针对float和double两种类型的代码分支。你在源码里看到的是一个宏调用但实际编译出来的代码是两份。阅读这类代码的技巧是先不要试图展开宏而是理解宏的语义。AT_DISPATCH_FLOATING_TYPES的语义就是“对float和double分别执行下面的代码块”。知道这个就够了不需要去看宏的具体展开。等你需要修改或扩展的时候再去研究宏的参数和展开规则。模板的情况类似。TensorImpl有很多模板方法但大部分时候你只需要关心具体类型实例化后的行为。比如TensorImpl::dataT()你只需要知道T是float时返回float指针是int64时返回int64指针不需要理解模板元编程的细节。6.2 用调试器跟踪Tensor操作的实际路径光看代码有时候不够因为PyTorch的调用链很深一个简单的tensor.view()可能经过好几层转发。这时候用调试器gdb或lldb设断点跟踪是最有效的。具体做法在TensorImpl::view或者Tensor::view上设断点然后跑一个最小的Python脚本调用view看调用栈。你会看到从Python层到C层的完整路径包括pybind11的转换层、dispatcher的分发层、最终到具体实现层。这个调用栈能帮你理解PyTorch的架构设计。我建议在调试时编译一个Debug版本的PyTorch虽然编译很慢可能要好几个小时但Debug版本有完整的符号信息调试体验好很多。Release版本经过优化很多变量被优化掉了断点也可能不准。6.3 常见编译错误与解决方案自己编译PyTorch源码时最常见的错误是依赖缺失和版本不匹配。比如CMake版本太低、编译器版本不满足要求、Python头文件找不到等。这些错误的解决方案通常是查官方文档的“从源码构建”章节按照推荐的版本组合来配置环境。另一个常见问题是内存不足。编译PyTorch需要大量内存特别是链接阶段16GB内存的机器可能会卡住。解决办法是减少并行编译的任务数MAX_JOBS4或者增加交换空间。我自己是在32GB内存的机器上编译的全程大约用了2小时。如果你只是想阅读源码而不需要修改其实不需要完整编译。用VS Code或者CLion打开源码目录配置好include路径代码跳转和补全就能正常工作。只有需要实际运行和调试时才需要编译。6.4 版本差异带来的阅读困惑PyTorch的源码在不同版本之间变化很大。比如TensorImpl的成员变量在1.x版本和2.x版本就有不少差异c10命名空间也是后来才引入的。如果你看的教程是基于旧版本的而你的源码是新版本的很多类名和文件路径会对不上。我的建议是先确定你使用的PyTorch版本然后从GitHub上checkout对应版本的源码。不要用main分支的源码去对照旧版本的教程那样只会越看越糊涂。另外PyTorch的代码注释质量参差不齐有些地方注释很详细有些地方完全没有注释这时候只能靠命名和上下文来推断意图。7. 几个高频Tensor操作的源码级解读7.1 squeeze和unsqueeze的零拷贝本质squeeze和unsqueeze都是零拷贝操作它们只修改sizes_和strides_不碰Storage。unsqueeze在指定位置插入一个大小为1的维度对应的stride设为1或者任意值因为大小为1的维度stride不影响索引。squeeze删除所有大小为1的维度或者删除指定的大小为1的维度。这两个操作之所以是零拷贝的是因为它们不改变元素在内存中的排列顺序只是改变了“如何解释”这个内存块。这也是为什么squeeze之后调用view通常能成功——因为内存布局没有变只是维度信息变了。但有一个边界情况如果squeeze删除了一个维度而这个维度的stride不是1那squeeze之后的张量可能变成非连续的。比如形状为(1, 3, 4)、strides为(100, 4, 1)的张量squeeze掉第一维后变成(3, 4)、strides为(4, 1)仍然是连续的。但如果strides是(100, 1, 4)squeeze后变成(3, 4)、strides为(1, 4)就不连续了。这个细节在写高性能代码时要注意。7.2 transpose和permute的stride变换transpose交换两个维度permute重新排列所有维度。它们的实现都是修改sizes_和strides_的对应位置不碰数据。比如形状为(2, 3)、strides为(3, 1)的张量transpose(0, 1)之后sizes变成(3, 2)strides变成(1, 3)。这个变换的数学原理是原来第i维的stride在变换后成为第j维的stride其中j是i在置换中的新位置。对于transpose来说就是交换两个维度的size和stride。transpose之后张量通常是非连续的这是后续view失败的常见原因。如果你需要对transpose后的张量做view要么先contiguous要么用reshape。但要注意contiguous之后的数据排列顺序和原来不同如果你依赖特定的内存顺序需要仔细验证。7.3 narrow和select的偏移量处理narrow在指定维度上取一个子区间select在指定维度上取一个索引。这两个操作会修改storage_offset_因为子区间的起始位置在原Storage中有一个偏移。比如一个形状为(10, 5)的连续张量strides为(5, 1)storage_offset为0。narrow(0, 2, 3)取第2到第4行结果的storage_offset变成2*510sizes变成(3, 5)strides不变。这意味着结果张量的第一个元素对应原Storage的第10个元素。这个偏移量机制是PyTorch内存共享的基础。多个Tensor可以指向同一个Storage的不同区域各自有各自的offset和shape。这种设计非常高效但也意味着修改一个Tensor的数据可能会影响另一个Tensor因为它们共享底层内存。在写代码时如果你不希望这种共享发生需要显式调用clone()。7.4 clone和detach的语义区别clone是深拷贝分配新的Storage复制数据返回一个完全独立的张量。detach是浅拷贝共享同一个Storage但断开计算图返回的Tensor不参与梯度计算。这两个操作经常被混淆。clone之后修改数据不会影响原张量detach之后修改数据会影响原张量因为它们共享内存。clone的开销是O(n)的内存分配和数据拷贝detach的开销几乎为零。在训练循环中如果你需要保存中间结果用于后续分析应该用clone().detach()既复制数据又断开计算图。如果只用detach()保存的张量会随着后续计算被修改因为内存是共享的。这个坑我在实际项目中踩过当时保存了一堆detach()后的张量结果发现它们的内容在训练过程中一直在变完全没法用来做分析。8. 从源码理解PyTorch的设计哲学8.1 为什么PyTorch选择动态图而非静态图从Tensor的C实现可以看出PyTorch的设计是“即时执行”的。每个操作在调用时立即计算没有预先构建计算图再执行的过程。这种设计的好处是调试方便你可以像写普通Python代码一样写模型用print和断点来调试。坏处是性能优化空间受限因为没法在全局视角做算子融合和内存规划。TensorImpl的设计也体现了这个哲学它不包含任何关于计算图的信息计算图的信息在AutogradMeta里通过TensorImpl的autograd_meta_成员关联。这种分离设计让Tensor的核心功能数据存储和索引和自动求导功能解耦各自独立演化。8.2 dispatcher机制如何支撑多后端PyTorch支持CPU、加速器等多种后端同一个torch.add调用在不同后端上有不同的实现。这个分发逻辑由dispatcher负责。当你调用torch.add(a, b)时dispatcher会根据a和b的device、dtype等信息找到对应的kernel来执行。dispatcher的核心是一个注册表每个算子在不同后端上注册自己的实现。注册的key包括算子名、dispatch key比如CPU、加速器、Autograd等。调用时dispatcher根据输入的Tensor计算出dispatch key集合然后按优先级选择合适的kernel。这个机制的好处是扩展性强。新增一个后端只需要注册对应的kernel不需要修改上层调用逻辑。坏处是调用链变长每次调用都要经过dispatcher的分发有一定的开销。PyTorch通过缓存dispatch结果来减少这个开销但第一次调用仍然比直接调用慢。8.3 TensorImpl的侵入式引用计数为什么比shared_ptr更适合前面提到TensorImpl用侵入式引用计数而不是std::shared_ptr。这个选择的原因主要有两个第一是性能。shared_ptr需要维护一个独立的控制块里面存引用计数和deleter。每次拷贝shared_ptr都要访问这个控制块多一次内存间接访问。侵入式引用计数把计数器放在对象内部访问更快。在PyTorch这种高频创建销毁Tensor的场景下这个性能差异很显著。第二是循环引用的处理。shared_ptr的循环引用会导致内存泄漏需要用weak_ptr来打破。侵入式引用计数虽然也有循环引用问题但PyTorch通过精心设计对象关系来避免循环。比如TensorImpl持有Storage的引用但Storage不持有TensorImpl的引用方向是单向的不会形成环。这个设计选择反映了PyTorch对性能的极致追求。在框架层面每一个微小的开销都会被放大到所有用户的所有操作上所以值得花大力气优化。8.4 内存布局的灵活性如何支撑各种操作PyTorch的Tensor支持任意stride这给了极大的灵活性。transpose、permute、narrow、expand等操作都是通过修改stride实现的不需要移动数据。这种设计让PyTorch能高效地表达各种张量操作而不需要为每种操作定义专门的数据结构。但灵活性也有代价。非连续张量的操作通常比连续张量慢因为内存访问模式不友好缓存命中率低。PyTorch在很多算子内部会检查连续性如果不连续就先转成连续再计算。这个转换是隐式的用户通常感知不到但性能影响是实实在在的。理解这一点之后我在写代码时会尽量保持张量连续避免不必要的transpose和permute。如果确实需要转置会在转置后尽快做contiguous而不是让非连续状态一直传递下去。这个习惯让我的训练代码性能提升了大约15%虽然不算巨大但在大规模训练中节省的时间很可观。9. 一些实用的调试与验证手段9.1 用Python层接口验证C层理解看源码时很多行为可以通过Python层的接口来验证。比如你想确认view在什么情况下会失败可以写一段Python代码构造各种形状和stride的张量调用view并捕获异常。通过观察哪些情况成功、哪些失败可以反推C层的判断逻辑。import torch # 连续张量view应该成功 a torch.randn(2, 3) b a.view(3, 2) print(连续张量view成功) # 转置后非连续view应该失败 c a.t() try: d c.view(3, 2) print(非连续张量view成功) except RuntimeError as e: print(f非连续张量view失败: {e}) # reshape应该成功 e c.reshape(3, 2) print(reshape成功)这种验证方式比单纯看代码更直观而且能帮你确认自己对源码的理解是否正确。我经常用这种方法来验证一些边界情况的行为。9.2 用torch._C接口深入底层Python层的torch._C模块暴露了很多C层的接口。比如torch._C._TensorBase就是Tensor的C基类在Python层的绑定。你可以通过它访问一些底层属性import torch a torch.randn(2, 3) print(a.storage().size()) # Storage的字节数 print(a.storage_offset()) # 偏移量 print(a.stride()) # strides print(a.is_contiguous()) # 是否连续这些接口在调试内存布局问题时非常有用。比如你可以通过storage().size()来确认两个Tensor是否共享同一个Storage通过storage_offset()来确认数据的起始位置。9.3 性能分析中如何定位Tensor操作开销PyTorch提供了torch.profiler来做性能分析它可以显示每个算子的耗时包括Tensor操作。如果你怀疑某个reshape或contiguous导致了性能问题可以用profiler来确认。import torch from torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CPU]) as prof: for _ in range(100): a torch.randn(1000, 1000) b a.t().reshape(1000000) c b.sum() print(prof.key_averages().table(sort_bycpu_time_total))在输出中你可以看到reshape和contiguous的调用次数和耗时。如果contiguous的耗时占比很高说明你的代码里有大量非连续张量需要转换值得优化数据布局。9.4 内存泄漏的排查思路Tensor相关的内存泄漏通常有两种原因一是循环引用导致引用计数无法归零二是全局缓存持有Tensor引用不释放。排查循环引用可以用Python的gc模块import gc import torch # 创建一些Tensor a torch.randn(1000, 1000) b torch.randn(1000, 1000) # 手动触发垃圾回收 gc.collect() # 查看无法回收的对象 for obj in gc.garbage: print(type(obj), obj)如果gc.garbage里有Tensor对象说明存在循环引用。常见的循环引用场景是把Tensor存在了某个对象的属性里而这个对象又被Tensor的hook引用了。全局缓存的问题通常出在自定义的缓存字典或者PyTorch内部的缓存上。PyTorch的torch.cuda模块有内存缓存CPU侧也有分配器缓存。这些缓存的设计目的是提高性能但在内存紧张的场景下可能造成困扰。可以通过环境变量或者API来调整缓存行为但要注意对性能的影响。10. 我个人的源码阅读方法10.1 从调用栈入手而非从头读文件我一开始看源码时犯过一个错误从TensorImpl.h的第一行开始读试图理解每一个成员和每一个方法。结果读了半天脑子里全是零散的信息串不起来。后来我改变了策略先写一个简单的Python脚本调用我想理解的操作然后用调试器看调用栈从栈顶往下追。这样每次只关注一个操作的实际路径信息量可控而且能立即看到代码的执行效果。比如我想理解view的实现就在TensorImpl::view上设断点跑一个a.view(3, 2)看调用栈。栈里会显示从Python到C的完整路径包括THPVariable_view、VariableType::view、TensorImpl::view等。沿着这个栈读代码比漫无目的地翻文件高效得多。10.2 用最小复现脚本验证理解看源码时我习惯随手写最小复现脚本来验证理解。比如我读到view的连续性检查逻辑就写一个脚本构造各种stride组合看view是否成功。如果结果和我的理解一致说明理解正确如果不一致说明我漏掉了某些逻辑需要回去重新看。这种“假设-验证”的循环是理解复杂系统最有效的方法。光看代码容易产生“我以为我懂了”的错觉只有实际跑一遍才能确认。10.3 建立自己的源码笔记体系PyTorch源码太庞大不可能全部记住。我建议建立自己的笔记体系记录关键类的职责、关键方法的调用链、容易混淆的概念对比。笔记不用很正式用Markdown文件就行重点是记录自己的理解而不是抄代码。我的笔记里有一张表记录了常见操作是否拷贝数据、是否修改stride、是否修改offset。这张表在我写高性能代码时非常有用能快速判断一个操作的开销。操作是否拷贝数据是否修改stride是否修改offsetview否是否reshape可能是可能transpose否是否narrow否否是clone是否否detach否否否contiguous可能是否这张表是我自己整理的不同版本可能有细微差异但大体逻辑是一致的。每次遇到不确定的操作查一下这张表就能快速判断。10.4 参与社区讨论与源码贡献最后一点看源码不要闭门造车。PyTorch的GitHub仓库有大量的issue和PR讨论很多设计决策的背景和取舍都在这些讨论里。遇到看不懂的地方搜一下相关的issue往往能找到解释。如果发现了文档错误或者代码里的小问题可以尝试提交PR。贡献源码的过程会强迫你更深入地理解代码而且能得到核心开发者的反馈学习效率比自己闷头看高很多。我自己的第一个PR只是修正了一个注释里的拼写错误但通过review过程我学到了很多关于代码规范和设计原则的东西。源码阅读是一个长期积累的过程不要指望一次就能全部理解。每次带着一个具体问题去读解决一个问题积累一个知识点时间长了自然就能形成体系。Tensor的C实现是PyTorch的核心理解了它再看其他部分比如autograd、dispatcher、算子实现就会容易很多。
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。