ARTICLE DETAIL

资讯详情

深耕网站SEO优化与搜索引擎排名提升的一线实战洞察。

torch-dct源码深度解析:Makhoul算法如何用FFT加速DCT-II离散余弦变换?

torch-dct源码深度解析:Makhoul算法如何用FFT加速DCT-II离散余弦变换? torch-dct源码深度解析Makhoul算法如何用FFT加速DCT-II离散余弦变换【免费下载链接】torch-dctDCT (discrete cosine transform) functions for pytorch项目地址: https://gitcode.com/gh_mirrors/to/torch-dcttorch-dct 是一个基于 PyTorch 的离散余弦变换DCT开源库它把 DCT-II 计算转化为一次 FFT支持 1D/2D/3D 正逆变换、GPU 加速与自动反向传播。本文带你读透其核心源码搞懂 Makhoul 算法到底如何加速 DCT-II 离散余弦变换。 torch-dct 是什么PyTorch 为什么缺一个 DCT 库DCTDiscrete Cosine Transform离散余弦变换是图像压缩如 JPEG、音频编码、信号处理中最经典的正交变换之一它能将能量集中到少数低频系数上天然适合压缩 特征提取。但在 PyTorch 深度学习场景中直接用 numpy/scipy 的 DCT 有三个痛点痛点说明❌ 无法反向传播模型里想放一层可学习前的 DCT 变换就做不了❌ 只能跑 CPU大尺寸数据如图像张量性能瓶颈明显❌ 直接矩阵乘法太慢朴素 DCT-II 是 O(N²) 的稠密矩阵运算torch-dct 的解法利用 Makhoul 算法把 DCT 重写成一次 FFT 复数相乘复杂度降到 O(N log N)并全程走 PyTorch 原生张量运算因此 CPU/GPU 通吃、梯度畅通。目前它实现的能力清单✅ 1-D DCT-I 及其逆变换✅ 1-D DCT-II 及其逆变换即 scaled DCT-III✅ 2-D / 3-D DCT-II 及其逆变换✅LinearDCTGPU 上快约 50 倍的线性层加速方案 torch-dct 安装与快速上手3 行代码完成 DCT 变换pip install torch-dct依赖要求torch0.4.1见 setup.py 中的install_requires。import torch import torch_dct as dct x torch.randn(200) X dct.dct(x) # DCT-II作用于最后一个维度 y dct.idct(X) # 逆变换idct(dct(x)) x注意所有变换都默认作用于张量的最后一个维度前面的维度自动视为 batch 维——所以上图、视频等任意形状数据都能直接进。 torch_dct 源码目录核心文件一览整个库非常小巧核心实现全部在一个文件里文件作用torch_dct/_dct.py全部变换实现dct/idct、dct1/idct1、2D/3D 变换、LinearDCTtorch_dct/__init__.py公开 API一行导出所有对外函数torch_dct/test/test_dct.py与 scipy 对拍的精度测试含 CUDAtorch_dct/test/test_lineardct.pyLinearDCT线性层测试setup.py打包配置声明torch0.4.1依赖⚡ Makhoul 算法揭秘DCT-II 如何变成一次 FFT先回忆 DCT-II 的定义N 点X[k] 2 · Σ x[n] · cos( π/(2N) · (2n1) · (k1/2) )直接按公式算每个 k 都要累加 N 项总复杂度 O(N²)。1984 年 J. Makhoul 的论文给出了关键洞察把 DCT-II 重排、对称补全后它就藏在一个实序列的 FFT 结果里。torch-dct 把这个洞察浓缩成了一行源码位于torch_dct/_dct.py的dct函数v torch.cat([x[:, ::2], x[:, 1::2].flip([1])], dim1)拆解一下这行奇偶魔术偶数下采样点x[0], x[2], …保持顺序奇数下采样点x[1], x[3], …翻转后拼在尾部拼接得到序列v对其做一次 FFTdct_fft_impl(v)FFT 结果乘上一组复数旋转因子e^{iπk/(2N)}源码中用cos(k) ± i·sin(k)构造取实部再乘 2 缩放——DCT-II 完成原始信号 x ──► 奇偶交错重排 v ──► FFT ──► × 旋转因子 ──► 取实部 ──► DCT-II 结果 O(1) 重排 O(N log N) O(N)对比朴素矩阵乘法当 N 4096 时O(N²) ≈ 1677 万次乘加而 FFT 路线只需约 5 万次量级——这就是用 FFT 加速 DCT的速度来源。 dct 函数逐段解读5 个关键步骤dct函数torch_dct/_dct.py第 57-88 行完整流程如下步骤代码要点说明① 展平x.view(-1, N)任意 batch 维压成 2D批处理免费获得② 重排奇偶交错构造vMakhoul 算法核心③ FFTdct_fft_impl(v)一次完整复数 FFT④ 旋转Vc[:, :, 0]*W_r - Vc[:, :, 1]*W_i复数乘旋转因子后取实部⑤ 归一化normortho时按 √N 缩放对齐 scipy 的 ortho 约定一个很贴心的工程细节文件开头的try/except做了版本兼容层——PyTorch ≥ 1.7走新的torch.fft.fft / rfft / irfftPyTorch ≤ 1.6回退到旧的torch.rfft / irfft所以这个库才能声明torch0.4.1的极低版本门槛。 idct 逆变换scaled DCT-III 是怎么做的idcttorch_dct/_dct.py第 91-131 行是 DCT-II 的逆数学上等于一个缩放的 DCT-III库的契约是idct(dct(x)) x。实现思路与正变换镜像对称频谱X除以 2必要时做 ortho 反归一化利用共轭对称性把一维频谱扩展成复序列虚部由翻转的-X构造见V_t_i一行乘旋转因子后做IFFTidct_irfft_impl从交错结果里把偶数、奇数采样点拆回原位x[:, ::2]与x[:, 1::2]。 2D / 3D DCT分离式变换如何免费扩展DCT-II 是可分离的2D 变换 逐行 1D DCT 逐列 1D DCT。dct_2d的实现torch_dct/_dct.py第 134-147 行因此只有三行逻辑X1 dct(x) # 沿最后一维 X2 dct(X1.transpose(-1, -2), normnorm) # 转置后再变一次dct_3d/idct_3d同理多一次transpose(-1, -3)。这意味着1D 核心的速度优势自动传递给高维——处理图像2D和体数据3D都无需额外实现。️ LinearDCTGPU 上快约 50 倍的另一种思路FFT 路线虽然快但 FFT 在 GPU 上的小核性能并不总是最优。LinearDCTtorch_dct/_dct.py第 204-227 行换了一条路继承nn.Linear把整个 DCT 矩阵预先算好对单位阵做dct得到存进权重权重requires_grad False——它是固定变换不参与训练前向就是一次大矩阵乘GEMMGPU 上实测比 FFT 路线快约 50 倍代价要存 N×N 矩阵内存占用更高。搭配apply_linear_2d/apply_linear_3d可把这套线性层方案直接推广到 2D/3D。简单说追求吞吐用LinearDCT追求内存效率用 FFT 版dct。 精度保障与 scipy 对拍的单元测试torch_dct/test/test_dct.py的做法非常硬核以scipy.fftpack.dct为参考实现逐点比对容差仅1e-10float64 级别覆盖 N 2、5、32、111 等大小长度以及 2~11 维的任意张量专门的test_cuda用例有 GPU 时把同一套比对在 CUDA 上再跑一遍torch_dct/test/test_lineardct.py则验证LinearDCT与 FFT 版结果一致float32 容差 1e-3。✅ 总结torch-dct 适用场景速查你的需求推荐用法模型里加一层可微 DCT 变换dct.dct/dct.idctGPU 直接可用JPEG 式图像压缩实验dct_2d/idct_2d分离式自动处理行列体数据 / 视频张量变换dct_3d/idct_3d追求极致 GPU 吞吐LinearDCTapply_linear_2d正交归一化约定传入normortho对齐 scipy 行为一句话总结torch-dct 用 Makhoul 的奇偶重排 FFT 旋转因子三板斧把 O(N²) 的 DCT-II 离散余弦变换降到了 O(N log N)并且全程保留梯度——核心就藏在torch_dct/_dct.py里不到 300 行代码中非常值得一读。【免费下载链接】torch-dctDCT (discrete cosine transform) functions for pytorch项目地址: https://gitcode.com/gh_mirrors/to/torch-dct创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表