PyTorch torch 命名空间完全指南张量运算、加速器、随机采样与序列化全景解析【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文以 PyTorch 官方 API 参考文档 docs/source/torch.md 为骨架系统梳理顶层torch命名空间的核心能力涵盖张量创建与索引、加速器CUDA/XPU/MPS 等体系、随机采样、序列化、并行控制、梯度禁用、各类数学运算、符号数字与编译优化等模块。读完本文你将能快速定位torch各 API 的用途、底层实现位置与正确使用姿势并能在实际项目中熟练运用张量操作、随机数控制、模型序列化与梯度上下文管理等高频能力。torch 模块概述torch是 PyTorch 的核心顶层命名空间官方文档通过 docs/source/torch.md 以automoduleautosummary的方式自动生成 API 索引。该文档覆盖了从张量基础操作到编译优化的完整技术栈张量的类型判断、默认配置与创建/索引操作加速器设备Accelerator抽象与异步执行模型随机数生成器Generator与各类采样函数模型与张量的序列化save/load线程并行控制梯度计算的局部禁用与恢复从逐元素运算到 BLAS/LAPACK 的完整数学运算体系符号数字SymInt/SymFloat/SymBool、导出export、控制流cond与编译compile等新特性。Tensors张量类型判断与全局默认配置张量类型谓词文档首先列出了一组用于判断张量属性的谓词函数它们全部接收一个输入并返回布尔值函数作用torch.is_tensor判断对象是否为torch.Tensor实例torch.is_storage判断对象是否为 Storage张量的底层存储torch.is_complex判断张量 dtype 是否为复数类型torch.is_conj判断张量是否为共轭位conjugate bit置位的视图torch.is_floating_point判断张量 dtype 是否为浮点类型torch.is_inference判断张量是否处于推理模式inference modetorch.is_neg判断张量是否为负数位negative bit置位的视图torch.is_nonzero判断张量是否含有非零元素torch.is_same_size判断两个张量形状是否相同torch.is_signed判断张量 dtype 是否为有符号类型torch.numel返回张量元素总数其中is_complex、is_floating_point、is_signed这类谓词直接与 dtype 体系挂钩可结合 torch/_C/init.pyi.in 中定义的 dtype 枚举理解其分类逻辑。全局默认配置set_default_dtype/get_default_dtype设置与读取默认浮点 dtype默认torch.float32影响torch.tensor、torch.zeros等未显式指定 dtype 的创建操作。set_default_device/get_default_device设置与读取张量创建的默认设备。设置后不带 device 参数的创建操作会落到该设备上CPU 上不可用需加速器环境。set_default_tensor_type旧式 API通过张量类型如torch.FloatTensor设置默认类型新代码推荐使用set_default_dtype。set_printoptions控制张量的打印格式精度、行宽、科学计数法等用于调试输出。set_flush_denormal控制在 x86 架构上将次正规数denormal刷新为零可提升特定场景性能。Creation Ops张量创建操作文档将创建操作分为两大类基础创建函数tensor-creation-ops 一节直接构造torch.tensor、torch.as_tensor尽量共享内存、torch.asarrayNumPy 风格、torch.as_strided、torch.scalar_tensor、torch.complex、torch.polar、torch.heaviside与外部数据互操作torch.from_numpy与 NumPy 数组共享内存、torch.from_dlpack遵循 DLPack 协议零拷贝交换张量、torch.frombuffer从实现了缓冲区协议的对象创建张量、torch.from_file内存映射文件稀疏张量sparse_coo_tensor、sparse_csr_tensor、sparse_csc_tensor、sparse_bsr_tensor、sparse_bsc_tensor分别对应五种稀疏布局全量填充zeros/zeros_like、ones/ones_like、full/full_like、empty/empty_like、empty_permuted、empty_strided、eye、arange、range、linspace、logspace量化相关quantize_per_tensor、quantize_per_tensor_dynamic、quantize_per_channel、dequantize、empty_quantized。随机采样创建函数文档明确说明torch.rand、rand_like、randn、randn_like、randint、randint_like、randperm属于随机采样类在 random-sampling 一节统一列出同时可用torch.empty配合 Tensor 的 in-place 随机方法如uniform_从更广泛的分布中采样。Indexing, Slicing, Joining, Mutating Ops索引、切片、连接与修改这一节是torch中最庞大的一族 API可归纳为几个子类形状变换与视图reshape、permute/permute_copy、transpose/transpose_copy、t/t_copy、squeeze/squeeze_copy、unsqueeze/unsqueeze_copy、swapaxes、swapdims、moveaxis、movedim、flatten位于 Other Operations、ravel、tile、narrow/narrow_copy、unbind/unbind_copy、split/split_copy、split_with_sizes_copy、chunk、tensor_split、vsplit/dsplit/hsplit。拼接与堆叠cat、concat、concatenate三者同一族、stack、column_stack、row_stack、dstack、hstack、vstack。索引与选取index_select、index_add、index_copy、index_put_、index_reduce、gather、take、take_along_dim、select/select_copy、nonzero/nonzero_static、argwhere、unravel_index、masked_select、masked_fill、where。就地修改copy/scatter 系列view_copy、slice_copy、slice_scatter、slice_inverse、select_scatter、diagonal_copy、diagonal_scatter、as_strided_copy、as_strided_scatter、expand_copy、detach/detach_copy、alias_copy、fill、put、scatter、scatter_add、scatter_reduce、segment_reduce、values_copy、view_as_complex_copy、view_as_real_copy、crow_indices_copy、col_indices_copy、ccol_indices_copy、row_indices_copy、indices_copy。值得注意的是文档中出现了大量以_copy结尾的函数如select_copy、slice_copy、permute_copy。这是 PyTorch 为函数化functionalization与编译后端引入的“返回副本而非视图”的语义便于在torch.compile、torch.export等场景下追踪张量别名关系避免视图操作被改写时引入无法表达的别名语义。Accelerators加速器抽象与异步执行什么是 Accelerator文档accelerators给出了 PyTorch 仓库内部对 Accelerator 的统一定义与 CPU 协同工作以加速计算的torch.device。这些设备采用异步执行模型通过torch.Stream与torch.Event完成同步。同时假设同一时刻宿主机上只有一个这样的加速器可用从而可以把当前加速器作为固定内存pinned memory、Stream 的 device_type、FSDP 等概念的默认设备。当前仓库承认的加速器设备无特定顺序包括CUDA对应文档 docs/source/cuda.rstMTIAMeta 自研推理加速器对应 docs/source/mtia.rstXPUIntel 扩展设备对应 docs/source/xpu.rstMPSApple Metal Performance Shaders对应 docs/source/mps.rstHPUHabana 设备PrivateUse1一类通过torch.library私有扩展机制接入的外部设备具体设备不在 PyTorch 仓库内。与 fork 的交互延迟加速器初始化这是加速器一节最具实操价值的忠告。PyTorch 生态中大量工具如 DataLoader 的子进程、算子内并行依赖fork创建子进程因此任何会阻止后续 fork 的操作都应尽量延迟——而绝大多数加速器的初始化恰好具有这种副作用。文档给出了明确的实践准则torch.accelerator.current_accelerator()默认是编译期检查因此始终 fork 安全反之向该函数传入check_availableTrue或调用torch.accelerator.is_available()通常会触发运行时可用性探测从而阻止后续 fork部分后端提供了 opt-in 的实验性开关使用 CUDA 设备时设置环境变量PYTORCH_NVML_BASED_CUDA_CHECK1可以让运行时可用性检查变为 fork 安全基于 NVML 而非初始化 CUDA context。实践中这意味着在启动任何多进程multiprocessing/DataLoader之前不要调用is_available()这类会初始化设备上下文的 API把设备探测尽量推迟到子进程内部或 fork 之后。torch.Stream与torch.Event是该异步模型的基石Stream 提供同一设备上算子执行的串行化队列Event 用于记录与等待执行到某个点二者是 CUDA 流/事件的统一抽象可迁移到其他加速器。Generators 与 Random sampling随机数体系Generatortorch.Generator是管理随机数生成状态的句柄负责维护随机数引擎的状态种子与偏移可跨调用共享。文档在 generators 一节列出并指出torch.default_generator返回默认的 CPUtorch.Generator。所有随机采样函数均可通过generator参数传入自定义 Generator 以获得可复现的独立随机流。随机种子管理在 random-sampling 一节文档列出了五个种子与状态管理函数其实现位于 torch/random.pytorch.manual_seed(seed)设置 CPU 随机种子同时会尽力设置当前加速器设备的种子返回设置后的Generatortorch.seed()用非确定性的系统熵初始化默认 Generator返回生成的种子torch.initial_seed()返回默认 Generator 的初始种子torch.get_rng_state()返回表示 CPU RNG 状态的张量用于保存与恢复随机状态torch.set_rng_state(new_state)恢复 CPU RNG 状态。复现实验的标准做法是在脚本开头调用torch.manual_seed(seed)而get_rng_state/set_rng_state可用于在训练中途精确保存/恢复随机流如断点续训、对比实验。采样函数族bernoulli、multinomial、normal、poisson、rand、rand_like、randint、randint_like、randn、randn_like、randperm构成标准采样 API。文档特别说明rand*系列与 in-place 方法可以组合出更丰富的分布。In-place 随机采样Tensor 上还定义了若干 in-place 随机方法inplace-random-sampling用于直接在已有张量上填充随机值Tensor.bernoulli_()torch.bernoulli的就地版本按伯努利分布以概率 p 生成 0/1Tensor.cauchy_()柯西分布Tensor.exponential_()指数分布Tensor.geometric_()几何分布Tensor.log_normal_()对数正态分布Tensor.normal_()torch.normal的就地版本正态分布Tensor.random_()离散均匀分布Tensor.uniform_()连续均匀分布。典型用法torch.empty(3, 4).uniform_(0, 1)生成[0, 1)上的均匀随机矩阵比torch.rand更灵活可任意指定范围与分布。Quasi-random sampling准随机采样torch.quasirandom.SobolEnginesobolengine 模板实现了 Sobol 序列低差异拟随机数生成器用于需要比伪随机数更均匀覆盖采样空间的任务如数值积分、超参数搜索、金融蒙特卡洛其单例接口与torch.rand类似但生成的序列具有更好的空间均匀性。Serialization张量与模型序列化torch.save与torch.load是 PyTorch 序列化的入口底层实现在 torch/serialization.pysave(obj, f, pickle_modulepickle, pickle_protocol2, ...)torch/serialization.pyobj可以是任意可 pickle 对象张量、模块、状态字典f可以是文件路径或文件类对象PyTorch 惯例使用.pt扩展名保存张量。自 1.6 版本起默认采用基于 zipfile 的新格式_use_new_zipfile_serializationTrue旧格式仍可被load读取。序列化会保留存储共享storage sharing即多个张量共享同一底层存储时反序列化后共享关系依然成立。load(f, map_location..., ...)torch/serialization.py支持map_location参数将张量重定位到目标设备如map_locationcpu、map_location{cuda:1: cuda:0}。import torch x torch.tensor([0, 1, 2, 3, 4]) torch.save(x, tensor.pt) # 保存到文件 buffer io.BytesIO() torch.save(x, buffer) # 保存到内存缓冲 y torch.load(tensor.pt, map_locationcpu) # 加载并重定位到 CPU序列化辅助函数同样位于 torch/serialization.pycheck_module_version_greater_or_equal(module, version)L489校验 torch 模块版本是否满足要求用于兼容性检查default_restore_location(storage, location)L713load的默认位置恢复回调负责根据map_location将存储放置到正确设备。Parallelism线程并行控制torch提供两级线程控制接口声明位于 torch/_C/init.pyi.inset_num_threads(n)/get_num_threads()控制算子内并行intra-op parallelism即单个算子内部使用的线程数对应底层 THPModule_setNumThreads/GetNumThreadsset_num_interop_threads(n)/get_num_interop_threads()控制算子间并行inter-op parallelism即并行执行多个相互独立的算子时使用的线程池大小torch.init_num_threads()显式初始化线程池必须在任何其他调用之前执行torch.fork/torch.wait文档中列出的并行原语用于并行任务编排。需要注意set_num_threads应在算子执行前调用且不同库如与 OpenMP 的交互可能影响实际生效情况set_num_interop_threads一旦线程池启动后便不可再更改。Locally disabling gradient computation局部禁用梯度这是文档中带有完整示例的核心章节torch-rst-local-disable-grad。PyTorch 提供三个上下文管理器用于在局部范围禁用/启用梯度计算其实现位于 torch/autograd/grad_mode.pytorch.no_grad()禁用梯度计算输出张量不追踪梯度推理、评测的标准选择torch.enable_grad()在no_grad外层作用域内局部重新启用梯度torch.set_grad_enabled(mode)按布尔参数动态开关梯度既可作上下文管理器也可作函数调用。关键语义这些上下文管理器是线程局部的thread local即它们只影响当前线程如果通过threading模块把计算发送到其他线程梯度开关不会随之生效。文档给出的官方示例 x torch.zeros(1, requires_gradTrue) with torch.no_grad(): ... y x * 2 y.requires_grad False is_train False with torch.set_grad_enabled(is_train): ... y x * 2 y.requires_grad False torch.set_grad_enabled(True) # 也可以作为函数使用 y x * 2 y.requires_grad True torch.set_grad_enabled(False) y x * 2 y.requires_grad False源码层面grad_mode.py 中set_grad_enabled的__enter__/__exit__直接调用torch._C._set_grad_enabled(mode)并保存前一状态用于恢复no_gradL22与enable_gradL89继承_NoParamDecoratorContextManager除上下文管理器外还可作为装饰器修饰函数。除上述三个外同一小节还列出torch.is_grad_enabled()查询当前线程梯度是否启用torch.autograd.grad_mode.inference_modegrad_mode.py推理模式比no_grad更强——它保证其中执行的算子完全不与 autograd 交互不创建图、不记录版本计数因而更快且允许视图改写等优化适用于数据加载、模型评测等确定不会与 autograd 交互的场景torch.is_inference_mode_enabled()查询当前是否处于推理模式。三者对比no_grad关闭梯度追踪但保留与 autograd 的交互能力inference_mode更进一步禁止任何 autograd 交互性能最优enable_grad用于在禁用的外层重新打开。注意set_grad_enabled的文档说明其不适用于前向模式自动微分forward-mode AD。Math operations完整数学运算体系Constants数学常量torch暴露四个数学常量分别是math模块对应常量的别名torch.e自然对数底数约 2.7183math.e别名torch.inf正无穷浮点值math.inf别名torch.nan非数值math.nan别名torch.pi圆周率约 3.1416math.pi别名。Pointwise Ops逐元素运算这是最庞大的 API 族覆盖算术add、sub/subtract、mul/multiply、div/divide、true_divide、floor_divide、pow、float_power、rsub、addcdiv、addcmul、lerp、remainder、fmod三角/反三角sin、cos、tan、asin/arcsin、acos/arccos、atan/arctan、atan2/arctan2、sinh、cosh、tanh、asinh/arcsinh、acosh/arccosh、atanh/arctanh指数对数exp、exp2、expm1、log、log10、log2、log1p、logaddexp、logaddexp2、logit、xlogy取整与符号ceil、floor、round、trunc/fix、frac、sign、sgn、signbit、abs/absolute、negative/neg、positive、copysign、nextafter特殊函数erf、erfc、erfinv、digamma、polygamma、lgamma、i0、igamma、igammac、mvlgamma、hypot、rsqrt、sqrt、square、reciprocal角度转换deg2rad、rad2deg逻辑与位运算logical_and、logical_or、logical_xor、logical_not、bitwise_and、bitwise_or、bitwise_xor、bitwise_not、bitwise_left_shift、bitwise_right_shift裁剪clamp、clip、clamp_min、clamp_max含 in-place 变体复数angle、imag、real、conj_physical、resolve_conj、resolve_neg量化fake_quantize_per_tensor_affine、fake_quantize_per_channel_affine量化感知训练前向传播中的伪量化激活与归一化算子softmax、nan_to_num、fill_、zero_以及sigmoid/tanh等。大量逐元素算子同时提供 in-place 版本下划线后缀如abs_、sqrt_、zero_、round_in-place 版本在内存复用与性能敏感路径中更受青睐。Reduction Ops归约运算极值max、min、amax、amin、aminmax、argmax、argmin统计量mean、nanmean、median、nanmedian、mode、std、var、std_mean、var_mean、quantile、nanquantile、norm、norm_except_dim、nuclear_norm、dist求和类sum、nansum、prod、logsumexp、count_nonzero、all、any去重unique、unique_consecutive其他hash_tensor张量内容哈希。归约运算大多支持dim指定维度、keepdim保持维度数返回张量或(values, indices)元组如max/min带 dim 时。Comparison Ops比较运算逐元素比较eq、ne/not_equal、ge/greater_equal、gt/greater、le/less_equal、lt/less张量级比较equal形状与值全等、allclose/isclose带容差的近似比较测试断言的首选元素属性isfinite、isinf、isposinf、isneginf、isnan、isreal、isin极值与排序maximum、minimum、fmax、fmin、argsort、sort、topk、kthvalue、msort。Spectral Ops频谱运算stft/istft短时傅里叶变换及其逆变换用于时频分析bartlett_window、blackman_window、hamming_window、hann_window、kaiser_window生成各类窗函数是音频与信号处理任务如torchaudio类应用的基础组件。Other Operations其他运算涵盖卷积与池化conv1d、conv3d、conv_transpose1d/2d/3d、convolution、avg_pool1d、max_pool1d、max_pool3d、adaptive_avg_pool1d、adaptive_max_pool1d、归一化batch_norm*、group_norm、instance_norm、native_layer_norm、rms_norm、RNN/LSTM 单元lstm、gru、rnn_relu、rnn_tanh及*_cell变体、损失函数ctc_loss、kl_div、poisson_nll_loss、margin_ranking_loss、hinge_embedding_loss、triplet_margin_loss、cosine_embedding_loss、自注意力相关scaled_dot_product_attention族、Dropout 系列dropout_、alpha_dropout、feature_alpha_dropout、feature_dropout以及einsum、meshgrid、cartesian_prod、cdist、pdist、corrcoef、cov、histc、histogram、histogramdd、bincount、searchsorted、bucketize、kron、tensordot、flip、fliplr、flipud、rot90、roll、repeat_interleave、block_diag、broadcast_to、broadcast_tensors、broadcast_shapes、atleast_1d/2d/3d、vander、tril_indices、triu_indices、diff、gradient等百余个算子。此外还包含一批设备/厂商专用算子cudnn_*、miopen_*、mkldnn_*、fbgemm_*与量化辅助q_scale、q_zero_point、q_per_channel_scales、q_per_channel_zero_points、q_per_channel_axis、choose_qparams_optimized、fused_moving_avg_obs_fake_quant、int_repr普通用户应优先使用对应的高层 API如torch.nn.functional这些底层算子由框架内部调度。BLAS and LAPACK Operations线性代数运算矩阵乘法mm、bmm、matmul、mv、dot、vdot、inner、outer、ger、addmm、addbmm、baddbmm、addmv、addr、chain_matmul、saddmm、spmm、dsmm、hsmm分解与求解svd、svd_lowrank、pca_lowrank、lobpcg、lu、lu_solve、lu_unpack、cholesky_inverse、cholesky_solve、inverse、pinverse、triangular_solve、geqrf、orgqr、ormqr行列式与迹det、logdet、slogdet、trace、matrix_power、matrix_exp数值积分辅助trapz、trapezoid、cumulative_trapezoid。这些算子通常下探到 BLAS/LAPACK 或厂商加速库cuBLAS、oneDNN 等实现是深度学习矩阵运算性能的关键路径。Foreach Operationstorch.foreach提供批量逐元素运算 API如对张量列表整体执行add、mul等相比 Python 循环逐个调用foreach 算子可显著减少调度开销、提升多张量场景如优化器参数更新的性能。文档说明迁移期间私有torch._foreach_*API 仍保留以兼容旧代码但推荐使用公开的torch.foreach接口。Utilities实用工具文档的 Utilities 一节覆盖五类工具自动混合精度autocastset_autocast_enabled、is_autocast_enabled、set_autocast_dtype、get_autocast_dtype及 CPU/GPU/IPU/XLA 各自的开关与 dtype 查询get_autocast_cpu_dtype、get_autocast_gpu_dtype等、clear_autocast_cache、autocast_increment_nesting/autocast_decrement_nesting。autocast 在保持数值精度的前提下自动为部分算子选择低精度 dtype如 float16/bfloat16以加速训练推理。异常检测set_anomaly_enabled、is_anomaly_enabled、is_anomaly_check_nan_enabled用于定位 NaN 梯度产生的源头会显著降低性能仅调试使用。确定性算法use_deterministic_algorithms、are_deterministic_algorithms_enabled、is_deterministic_algorithms_warn_only_enabled、set_deterministic_debug_mode、get_deterministic_debug_mode、set_float32_matmul_precision、get_float32_matmul_precision用于追求可复现的运算结果或平衡精度与性能。设备与类型工具get_device、get_device_module、result_type、can_cast、promote_types、compiled_with_cxx11_abi查询 libtorch 是否以 C11 ABI 编译、is_distributed、is_vulkan_available、vmap向量化映射、typename返回对象的 Python 类型名。IR 与解析工具import_ir_module、import_ir_module_from_buffer、parse_ir、parse_schema、parse_type_comment、merge_type_from_type_comment面向 TorchScript 与 IR 处理的底层能力。Type Information 与 Symbolic Numbers类型体系与符号数字TensorType是张量类型信息的自动文档类。在符号数字Symbolic Numbers方面文档展示了动态形状系统中的三类核心对象torch.SymInt、torch.SymFloat、torch.SymBool分别表示符号整数、符号浮点与符号布尔是torch.compile动态形状dynamic shapes追踪的基石允许形状表达为符号表达式而非具体数值符号运算辅助函数sym_float、sym_int在符号与具体数值间转换、sym_max、sym_min、sym_sum、sym_not、sym_iteif-then-else、sym_sqrt、sym_fresh_size、sym_constrain_range、sym_constrain_range_for_size为符号值附加取值范围约束供编译器生成守卫。这些 API 主要面向框架开发与高级编译用户普通训练代码通常经由torch.compile间接受益。Export、Control Flow 与 Optimizations新特性区Export Pathtorch.export文档以原型prototype标注用于将模型导出为可序列化、可移植的中间表示并附有generated/exportdb/index的导出用例索引。文档明确警告该特性仍处于原型阶段未来可能发生破坏性变更。Control Flowtorch.cond同样标注为原型用于在符号执行/编译图内表示数据依赖的条件分支是torch.compile与导出流程中表达动态控制流的关键算子。Optimizationstorch.compile文档将torch.compile列为编译优化入口详见仓库内 torch.compiler 用户指南它将 Python 级模型图捕获、算子融合与后端代码生成一体化是当前 PyTorch 性能优化的主路径。Operator Tagstorch.Tag为算子提供元数据标签如torch.Tag.pointwise、torch.Tag.dynamic_output_shape等供编译后端与调度系统按类别处理算子。附录文档中隐藏的模块追踪torch.md 末尾还通过空py:module声明对一批子模块做了仅用于追踪的登记torch.contrib、torch.utils.backcompat、torch.utils.hipifyROCm 构建内部使用、torch.utils.model_dump提供get_model_info、get_inline_skeleton模型信息工具、torch.utils.viz、torch.quasirandom、torch.return_types、torch.serialization含StorageType、torch.signal.windows.windows、torch.sparse.semi_structured、torch.storage、torch.torch_version、torch.types、torch.version与torch.compiler.config。这些声明帮助文档系统正确索引子模块而不强制生成页面也从侧面勾勒出torch命名空间的外延边界。小结torch顶层命名空间是 PyTorch 一切能力的入口从张量的创建、索引与逐元素/归约/线性代数运算到加速器抽象与异步执行再到随机数控制、序列化、线程并行与梯度上下文管理以及面向编译前沿的符号数字、导出与控制流原语。阅读本文后建议按需深入对应源码文件梯度管理见 torch/autograd/grad_mode.py序列化见 torch/serialization.py随机数见 torch/random.py线程与底层接口声明见 torch/_C/init.pyi.in即可形成从 API 到实现的完整知识闭环。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考