CoderXL's Blog

Back

PyTorch 库 API 总结Blur image

1. torch#

生成张量#

.tensor(data) 创建一个张量, 可以传入一个数(将创建标量)或者一个列表,甚至一个 numpy 数组。

.[dist](size, ...) 按照 size 指定的形状创建一个符合分布 [dist] 的张量,其中 [dist] 可以取:

  • zeros 全零
  • ones 全一
  • full(, fill_value) 全部为 fill_value
  • empty 内存全部为未初始化的数据
  • rand[0,1)[0,1) 上均匀分布
  • randint(, low, high, size, ...)[low,high)[\mathrm{low}, \mathrm{high}) 中均匀分布的整数
  • randn 均值为 00,标准差为 11 的正态分布

.[dist]_like(input: torch.Tensor) 创建一个新张量,形状与 input 相同,符合分布 [dist],见上。

.arange(n: int) 创建一个 0n10\sim n-1 的一维等差数列张量。

.linspace(start, end, steps, *, ...) 传入上下限,创建一个长度为 steps 的一维张量,值为从 startend 的等差数列。

.normal(mean: Tensor, std: Tensor, ...)

  • 要求 meanstd 元素个数相同。创建一个与 mean 形状相同的张量,其中的各个元素都是符合 meanstd 对应位置元素所指定的正态分布的随机数,其中 mean 是均值,std 是标准差。

.normal(mean: Tensor, std: float, ...)

  • 和上方的类似,只是 std 对所有元素都相同。

.normal(mean: float, std: Tensor, ...)

  • 和上方的类似,只是返回的张量形状由 std 确定,且 mean 对所有元素都相同。

.normal(mean: float, std: float, size: Tuple[int], ...)

  • 创建一个形状为 size 的张量,每个元素都是服从 meanstd 所指定的正态分布的随机数,其中 mean 是均值,std 是标准差。

多个张量的操作#

.cat(tensors: Sequence[torch.Tensor], dim=0, ...) 从各张量当前的第 dim 维拼接成一个新张量。要求各张量除了第 dim 维外,形状一致。生成的新张量维度不增加。

.stack(tensors: Sequence[torch.Tensor], dim=0, ...) 新增一个维度作为新张量的第 dim 维,要求 dim 介于 00 和各张量维度数(含)之间。要求各张量形状完全一致。生成的新张量维度加一。

非原地张量变换#

可以从 torch 调用,但一般从 torch.Tensor 实例上调用,详见 #2. torch.Tensor

.[manip](input: torch.Tensor, ...) 返回对张量进行 [manip] 操作后的新张量,[manip] 可以为:

  • 各种数学函数,比如 abs | absolutecosceilsum 等等。
  • 张量操作,比如 reshapetpermute 等等。

原地张量变换#

不能从 torch 调用,必须要从 torch.Tensor 实例上调用,见 #2. torch.Tensor

其它#

.no_grad() 可以配合 with 创建暂时停止记录计算图的上下文。在这个上下文中,原本 requires_grad=True 的叶子张量可以进行原位操作,并且在上下文结束后恢复记录计算图。

  • with torch.no_grad():

2. torch.Tensor#

元数据获取#

.size() 获取该张量的维度信息,numpy 风格写法。

.shape 结果和 .size() 相同,属性风格写法,更推荐。

.numel() 返回该张量的总元素个数(number of elements)。

.requires_grad 属性,获取该张量是否在被计算图跟踪。

非原地张量变换#

统一形式是 .[manip](...),返回一个张量,原张量不受改变。以下介绍常用的 [manip] 取值。

数学函数#

由于十分常见且数量众多,只作为备忘录,介绍重要但易混淆的几个

.arg[max|min]() 将该张量展平为一维之后,返回最大(小)值所在位置的下标。

.arg[max|min](dim, keepdim=False) 生成一个新张量,每个元素为原张量对应的在 dim 维度方向延伸的排上的最大(小)值所在位置的索引。该张量相比原张量缺少了第 dim 维(keepdim=False),或者第 dim 维长度为 11keepdim=True)。

张量操作#

PyTorch 的非原地张量操作一般来说只构建新的元数据,不复制张量内容,因此内存和时间开销都接近甚至优于原地操作。

.clone() 返回该张量的一个副本,元数据与数据内存均分离。新张量允许梯度信息在反向传播时流回原张量。

.data 属性,不建议使用。会返回该张量的裸的(不带计算图/梯度跟踪的)内部数据张量,可以进行原位修改以改变原张量,在较早的版本中被用来规避 requires_grad=True 叶子张量不能原位修改的问题。这样做被认为是不安全的,因此被 .no_grad() 的方式取代。

.type(dtype=None, ...) 如果未提供 dtype,则返回该张量的类型;否则如果该张量已经是 dtype 类型,直接返回该张量;否则产生一个新张量,是将原张量转换为 dtype 后的结果。

.detach() 返回一个从当前计算图中分离的新张量,具有独立的元数据,但与原张量共享数据内存,对新张量的原地修改将会影响原张量。可以和 .clone() 配合使用。

.view(shape) 可以传入一个 Iterable[int] 或者多个 int,作为新的维度信息。会返回一个当前张量的新 view,共享内存,但是有不同的元数据。.view 要求当前张量内存连续(.contiguous==True),否则会报错。.view 不改变当前张量。

.reshape(shape) 强化版的 .view,在当前张量不连续的时候,复制张量到新的内存中去,恢复 .contiguous==True;其他时候与 .view 一致。一般都使用 .reshape 而不是 .view.reshape 不改变当前张量。

  • 可以在至多一个维度上使用 -1 来允许 torch 自行推断维度大小。
    比如:shape: [3,4,5] .reshape(-1, 4) 会变成 shape: [15, 4]

.contiguous() 返回一个内存连续的张量,其中包含与该张量相同的数据。如果该张量已经是指定的内存格式,则此函数返回该张量本身。

.t() 返回二维张量的转置。

.transpose(dim0, dim1) 返回新张量,交换了原张量的两个维度,由 dim0dim1 指定。

.permute(dims: Tuple[int]) 返回新张量,等于原张量按照 dims 指定的顺序重新排列维度。

.squeeze(dim: Optional[Union[[int]]]) 去除长度只有 11 的维度。当指定 dim 的时候,只在这些维度上尝试去除。可以生成标量。

.unsqueeze(dim: int) 新增一个维度作为新张量的第 dim维,要求 dim 介于 00 和原张量维度数(含)之间,该维长度为 1。数据与原张量相同。

.flatten(start_dim=0, end_dim=-1) 将张量展平为一维。如果指定 start_dimend_dim,那么只展平它们之间的维度。

原地张量变换#

统一形式是 .[manip]_(...),对该张量进行原地修改,返回该张量。注意,对 requires_grad=True 的张量不可以使用原地变换!原因很显然,原位修改无法被记录到计算图中,因而会导致梯度计算错误。如果确实需要原位修改(比如初始化的时候),可以使用 with torch.no_grad(): 上下文,或者使用不建议的属性 .data 获取裸的张量内容进行修改。

数学函数#

一般都有原地形式

张量操作#

一般有原地形式,除了几个特例:

.clone_()

.view_()

.reshape_()

.contiguous_()

特别地,虽然可以使用 .detach_(),但存在限制:view 张量不能原地分离。

此外,还有一些只支持原地形式的张量变换:

.fill_(value) 使用 value 填充该张量。

.zero_() 使用 00 填充该张量。

.uniform_(from=0, to=1) 将该张量填充为从连续均匀分布中采样的数字,其中 from 为下限,to 为上限。

.normal_(mean=0, std=1, ...) 使用随机数填充该张量,其中每个元素都是服从 meanstd 所指定的正态分布的随机数,其中 mean 是均值,std 是标准差。

.random_(from=0, to=None, ...) 使用离散均匀分布在 [from, to  1][\mathrm{from}, \mathrm{to} - 1] 上的数字填充该张量。如果未指定,值通常仅受该张量自身数据类型的限制。
然而,对于浮点类型,如果未指定,范围将是 [0, 2mantissa][0, 2^{\mathrm{mantissa}}],以确保每个值都是可表示的。例如,torch.tensor(1, dtype=torch.double).random_() 将在 [0, 253][0, 2^{53}] 上均匀分布。

.bernoulli_(p=0.5, ...) 使用 p 所指定的伯努利分布填充该张量。注意:存在 .bernoulli() 函数,但其含义与此函数相去较远,因此不认为此函数有非原地形式!

.requires_grad_(requires_grad=True) 指定该张量做好计算梯度的准备(会给该张量增加一个 .grad 属性用于存储梯度;效果等同于在创建张量的时候传入 requires_grad=True;比直接修改属性 .requires_grad 更安全)。

其它#

.backward() 计算反向传播。

.numpy() 转换为 numpy 数组(仅限 cpu 张量)。

3. torch.utils.data#

数据集模块

.Dataset#

映射式数据集的抽象类

  • 应当实现 .__getitem__(),还可以实现 .__len__().__getitems__()
.TensorDataset(*tensors)#

.Dataset 的一个子类,接受任意多的张量,只要每个张量在第 0 维(即样本总数)上大小相同。

  • 该类不会自动推断输入的张量的用途,不会给张量按照 featureslabels 等方式分类。 它只是做好容器的本职工作。
    但它又比纯粹的张量进了一步,实现了比较方便的索引和随机化方法。
    .TensorDataset 的功能不能被一个大张量代替,因为它允许打包多个只在一个维度上形状相同的张量,而这些张量本身是无法拼接、堆叠的。
.IterableDataset#

可迭代数据集的抽象类

  • 应当实现 .__iter__()
.DataLoader(dataset, batch_size=1, shuffle=False, ...)#

接受上述两种数据集作为 dataset 输入,作用是提供封装好的随机取样 batch 迭代器(省去了 构建 indices + 手动 shuffle)。它自身提供 .__iter__().__next__(),可以直接作为迭代器使用。

4. torch.nn#

神经网络模块

.Module 神经网络模块的基类#
.Sequential 容纳模块的顺序容器类#

相比 .ModuleList,该类除了容纳模块,还会以级联的方式链接相邻模块

.forward() 方法接受任何输入,并将其传递给它包含的第一个模块。然后,它将输出按顺序“链接”到后续每个模块的输入,最后返回最后一个模块的输出。

.Linear(in_features: int, out_features: int, bias=True, ...) 线性层#

对输入数据 A\boldsymbol{A} 应用仿射线性变换:y=wAT+b\boldsymbol{y}=\boldsymbol{w}\boldsymbol{A}^T+\boldsymbol{b},其中 w\boldsymbol{w}b\boldsymbol{b} 是可以学习的参数。

in_features 指定每个输入样本的大小(具体含义可见线性回归模型#符号化表述out_features 指定每个输出样本的大小 bias 指示是否应学习加性偏置

变量#

weight: torch.Tensor 形状为 (out_features, in_features) bias: torch.Tensor 形状为 (out_features)

PyTorch 库 API 总结
https://blog.leosrealms.top/blog/2025-11-14-pytorch-library-api-summary
Author CoderXL
Published at 2025年11月14日
Comment seems to stuck. Try to refresh?✨