企业定位是网站建设的,网站负责人信息,西安金融网站建设,集约化网站群建设情况Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境#xff0c;以高效编写自定义 DNN 计算内核#xff0c;并能够在现代 GPU 硬件上以最大吞吐量运行。
更多 Triton 中文文档可访问 →https://triton.hyper.ai/
在本教程中#xff0c;您将编…Triton 是一种用于并行编程的语言和编译器。它旨在提供一个基于 Python 的编程环境以高效编写自定义 DNN 计算内核并能够在现代 GPU 硬件上以最大吞吐量运行。
更多 Triton 中文文档可访问 →https://triton.hyper.ai/
在本教程中您将编写一个内存高效的 Dropout 实现其状态将由单个 int32 seed 组成。这与传统 Dropout 实现不同传统实现通常由与输入 shape 相同的位掩码张量组成。
在这过程中您将学习到以下内容 PyTorch 中 原生实现 Dropout 的局限性。 Triton 中的并行伪随机数生成。
简介
Dropout 是在 [SRIVASTAVA2014] 中引入的一种技术用于改善低数据条件下深度神经网络的性能通常用于正则化。它接受一个向量作为输入并生成相同 shape 的输出向量。输出中的每个标量都有概率 p 被设为零否则直接从输入复制。这使得网络在仅有输入的 1−p 标量时也能表现良好。
在评估阶段为了充分利用网络的能力将 p 设为 0。但是简单地将 p 设为 0 会增加输出的范数可能会人为地降低输出的 softmax temperature。为了防止这种情况发生输出被缩放为 1/(1-p)这使得无论 dropout 概率如何都能保持一致的范数。
Baseline
首先看一下 baseline 的实现。
import tabulate
import torchimport triton
import triton.language as tltriton.jit
def _dropout(x_ptr, # 输入指针x_keep_ptr, # pointer to a mask of 0s and 1s 由 0 和 1 组成的掩码的指针output_ptr, # pointer to the output 输出指针n_elements, # number of elements in the x tensor x 张量的元素数量p, # probability that an element of x is changed to zero 元素 x 被设置为 0 的概率BLOCK_SIZE: tl.constexpr,
):pid tl.program_id(axis0)block_start pid * BLOCK_SIZEoffsets block_start tl.arange(0, BLOCK_SIZE)mask offsets n_elements# Load data# 加载数据x tl.load(x_ptr offsets, maskmask)x_keep tl.load(x_keep_ptr offsets, maskmask)# The line below is the crucial part, described in the paragraph above!# 下一行是上段描述的关键部分output tl.where(x_keep, x / (1 - p), 0.0)# Write-back output# 写回输出tl.store(output_ptr offsets, output, maskmask)def dropout(x, x_keep, p):output torch.empty_like(x)assert x.is_contiguous()n_elements x.numel()grid lambda meta: (triton.cdiv(n_elements, meta[BLOCK_SIZE]), )_dropout[grid](x, x_keep, output, n_elements, p, BLOCK_SIZE1024)return output# Input tensor
# 输入张量
x torch.randn(size(10, )).cuda()
# Dropout mask
# Dropout 掩码
p 0.5
x_keep (torch.rand(size(10, )) p).to(torch.int32).cuda()
#
output dropout(x, x_keepx_keep, pp)
print(tabulate.tabulate([[input] x.tolist(),[keep mask] x_keep.tolist(),[output] output.tolist(),
]))Out:
种子化 Dropout
上述 Dropout 实现效果良好但管理 Dropout 状态可能会变得复杂特别是在考虑反向传播和重新计算/检查点场景时。在这里我们描述一种替代实现它具有以下优点
更小的内存占用。较少的数据移动。简化了在多次调用内核函数时持久化随机性的管理。
生成 Triton 中的伪随机数很简单在本教程中我们将使用 triton.language.rand 函数该函数基于给定的种子和一组 int32 偏移量生成一个块的均匀分布的 float32 值范围在 (0, 1) 内。但如果你需要Triton 也提供其他随机数生成策略。 注意 Triton 的 PRNG 实现基于 Philox 算法详见 [SALMON2011]。 现在将所有内容整合起来。
triton.jit
def _seeded_dropout(x_ptr,output_ptr,n_elements,p,seed,BLOCK_SIZE: tl.constexpr,
):# compute memory offsets of elements handled by this instance# 计算由此实例处理的元素的内存偏移量pid tl.program_id(axis0)block_start pid * BLOCK_SIZEoffsets block_start tl.arange(0, BLOCK_SIZE)# load data from x# 从 x 读取数据mask offsets n_elementsx tl.load(x_ptr offsets, maskmask)# randomly prune it# 随机修剪random tl.rand(seed, offsets)x_keep random p# write-back# 写回output tl.where(x_keep, x / (1 - p), 0.0)tl.store(output_ptr offsets, output, maskmask)def seeded_dropout(x, p, seed):output torch.empty_like(x)assert x.is_contiguous()n_elements x.numel()grid lambda meta: (triton.cdiv(n_elements, meta[BLOCK_SIZE]), )_seeded_dropout[grid](x, output, n_elements, p, seed, BLOCK_SIZE1024)return outputx torch.randn(size(10, )).cuda()
# Compare this to the baseline - dropout mask is never instantiated!
# 与基线相比 - dropout 掩码从未被实例化
output seeded_dropout(x, p0.5, seed123)
output2 seeded_dropout(x, p0.5, seed123)
output3 seeded_dropout(x, p0.5, seed512)print(tabulate.tabulate([[input] x.tolist(),[output (seed 123)] output.tolist(),[output (seed 123)] output2.tolist(),[output (seed 512)] output3.tolist(),
]))Out: 大功告成我们现在有了一个 Triton 内核可以在给定相同种子的情况下应用一致的 dropout 掩码。与传统的 dropout 实现相比这种方法减少了内存开销并简化了状态管理。
练习
扩展内核以处理矩阵并使用一个种子向量 — 每行一个种子。添加对 striding 的支持。挑战实现稀疏 Johnson-Lindenstrauss 变换的内核每次使用种子动态生成投影矩阵。
参考文献 [SALMON2011] John K. Salmon, Mark A. Moraes, Ron O. Dror, and David E. Shaw, “Parallel Random Numbers: As Easy as 1, 2, 3”, 2011 [SRIVASTAVA2014] Nitish Srivastava et al., “Dropout: A Simple Way to Prevent Neural Networks from Overfitting”, JMLR 2014
Download Jupyter notebook: 04-low-memory-dropout.ipynb
Download Python source code: 04-low-memory-dropout.py
Download zipped: 04-low-memory-dropout.zip
