Pytorch源码搜索与分析

PyTorch的的代码主要由C10、ATen、torch三大部分组成的。其中:

C10

C10,来自于Caffe Tensor Library的缩写。这里存放的都是最基础的Tensor库的代码,可以运行在服务端和移动端。PyTorch目前正在将代码从ATen/core目录下迁移到C10中。C10的代码有一些特殊性,体现在这里的代码除了服务端外还要运行在移动端,因此编译后的二进制文件大小也很关键,因此C10目前存放的都是最核心、精简的、基础的Tensor函数和接口。

C10目前最具代表性的一个class就是TensorImpl了,它实现了Tensor的最基础框架。继承者和使用者有:

Variable的Variable::Impl

SparseTensorImpl

detail::make_tensor(storage_impl, CUDATensorId(), false)

Tensor(c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> tensor_impl)

c10::make_intrusive<at::TensorImpl, at::UndefinedTensorImpl>

值得一提的是,C10中还使用/修改了来自llvm的SmallVector,在vector元素比较少的时候用以代替std::vector,用以提升性能;

ATen

ATen,来自于 A TENsor library for C++11的缩写;PyTorch的C++ tensor library。ATen部分有大量的代码是来声明和定义Tensor运算相关的逻辑的,除此之外,PyTorch还使用了aten/src/ATen/gen.py来动态生成一些ATen相关的代码。ATen基于C10,Gemfield本文讨论的正是这部分;

torch._VF

torch._VF is a subset of those functions are mapped to ATen functions in torch/jit/_builtins.py.

例如如果沿着torch.nn.dropout去找,会发现找到torch._VF这个库,IDE无法继续navigate:

这时候直接去github的pytorch源码repo(https://github.com/pytorch/pytorch/tree/b3bb234e16780ea3f3d749d2c8c156c9245eb797)搜dropout_就能很快看到在cpu上cpp的实现:

加上cuda的限制就是gpu的实现:

另外在pytorch的老旧版本,像dropout这种由python实现的算子在pytorch v1.0.4版本的https://github.com/pytorch/pytorch/blob/v0.1.4/torch/nn/functions/dropout.py 这个目录下还能找到,但是github并没有提供历史tag的搜索功能,不如就去https://github.com/colesbury/pytorch-old/tree/master/torch/nn/functions 这个古早目录下面找吧

Torch

Torch,部分代码仍然在使用以前的快要进入历史博物馆的Torch开源项目,比如具有下面这些文件名格式的文件:

TH* = TorcH

THC* = TorcH Cuda

THCS* = TorcH Cuda Sparse (now defunct)

THCUNN* = TorcH CUda Neural Network (see cunn)

THD* = TorcH Distributed

THNN* = TorcH Neural Network

THS* = TorcH Sparse (now defunct)

THP* = TorcH Python

PyTorch会使用tools/setup_helpers/generate_code.py来动态生成Torch层面相关的一些代码,这部分动态生成的逻辑将不在本文阐述,你可以关注Gemfield专栏的后续文章。

部分转载自

  1. https://zhuanlan.zhihu.com/p/55966063
相关推荐
区块链蓝海几秒前
$TEA将于6月4日主网启动并同步登陆多个主流平台
人工智能·区块链
咖啡星人k1 分钟前
云端开发环境架构演进:容器隔离、文件系统与AI深度集成
人工智能·架构
小小王app小程序开发8 分钟前
AI数字人小程序开发玩法深度解析:功能架构、技术实现与落地场景
人工智能·架构
想你依然心痛11 分钟前
HarmonyOS 6(API 23)实战:打造“光味智厨“——AI烹饪新体验
人工智能·华为·ar·harmonyos·智能体
人工智能培训14 分钟前
数字孪生建模常用方式有哪些?
人工智能·深度学习·机器学习·容器·知识图谱
CCC:CarCrazeCurator16 分钟前
【DriveGen 文件详解】04——evaluate.py
人工智能·自动驾驶·transformer
老高学长17 分钟前
金融机构文档加密软件哪个好|合规与安全兼顾|2026新测评
网络·人工智能·安全
闻道参看17 分钟前
生成式智能搜索下的流量卡位攻略:初创个体如何甄选高兼容性的 GEO 优化 服务商
人工智能
Herlie19 分钟前
6款可编辑AI海报工具深度横测(2026)
大数据·人工智能
轻刀快马20 分钟前
跨越“拟人”的最后一道天堑:大模型强化学习(RLHF/RLAIF)底层原理解析
人工智能·深度学习·机器学习