x_fft = torch.fft.rfft(x, dim=1, norm=ortho)

下面给出一个例子，假设我们有如下数据，其形状为
[
(\text{batch} \times \text{ts_d},, \text{seg_num},, d_{\text{model}}) = (4,,4,,3)
]

这里我们设定：

• batch = 2
• ts_d = 2
• seg_num = 4（也就是每个信号长度为 4）
• d_model = 3（每个时刻有 3 个特征）

因此，共有 4 个信号（可以认为是 4 个独立样本），每个信号是一个 4×3 的矩阵。比如我们定义数据如下（为了方便说明，这里各个样本的数据取不同值）：

x = tensor([[ [ 1,  2,  3],    // 第1个样本，第1个通道的信号： [1, 4, 7, 10] 后面会说明[ 4,  5,  6],[ 7,  8,  9],[10, 11, 12] ],[ [ 2,  3,  4],[ 5,  6,  7],[ 8,  9, 10],[11, 12, 13] ],[ [ 3,  4,  5],[ 6,  7,  8],[ 9, 10, 11],[12, 13, 14] ],[ [ 4,  5,  6],[ 7,  8,  9],[10, 11, 12],[13, 14, 15] ]
])

这里每个最外层的元素对应一个样本（总数为4，即 batch×ts_d），每个样本内部是一个 4×3 的矩阵，其中：

• 第一维（索引0到3）为 seg_num，代表一个信号的时间步数；
• 第二维（索引0到2）为 d_model，对应每个时刻的 3 个特征。

计算过程说明

我们调用

x_fft = torch.fft.rfft(x, dim=1, norm='ortho')

这里说明：

• dim=1 表示对每个样本中的 seg_num 维度进行快速傅里叶变换（FFT），即对每个通道（d_model）独立地对长度为4的序列进行变换。
• norm='ortho' 表示采用正交归一化，即每个变换乘以 ( \frac{1}{\sqrt{N}} )（这里 ( N=4 )）。

1. FFT 后的形状

由于输入沿 dim=1 的长度为 4，为了利用实数输入的共轭对称性质，rfft 只计算非负频率部分，其返回的频率个数为
[
\text{freq_num} = \frac{\text{seg_num}}{2} + 1 = 3.
]
因此，计算后的输出形状为
[
(4,, 3,, 3)
]
即 4 个样本，每个样本有 3 个频率成分，每个频率成分对应 3 个特征（d_model）。

2. 对单个样本、单个通道的计算示例

为了详细说明计算过程，我们以第 0 个样本的第 0 个通道为例。
在该样本中，第 0 个通道的信号为：
[
x[0, :, 0] = [1,, 4,, 7,, 10]
]
我们对这个长度为 4 的实数序列进行 rfft，正交归一化版本的离散傅里叶变换公式为
[
X[k] = \frac{1}{\sqrt{4}} \sum_{n=0}^{3} x[n], e^{-i,2\pi,k,n/4},\quad k=0,1,2.
]

• 频率 ( k=0 ):
  指数部分 ( e^{-i,2\pi,0,n/4} = 1 ) 对所有 ( n ) 都为 1。
  计算：
  [
  X[0] = \frac{1}{2} (1 + 4 + 7 + 10) = \frac{22}{2} = 11.
  ]
  结果为实数 11。

• 频率 ( k=1 ):
  指数部分为：

  • ( n=0 )：( e^{-i,0} = 1 )
  • ( n=1 )：( e^{-i,2\pi/4} = e^{-i,\pi/2} = -i )
  • ( n=2 )：( e^{-i,\pi} = -1 )
  • ( n=3 )：( e^{-i,3\pi/2} = i )

  对应计算：
  [
  \begin{aligned}
  X[1] &= \frac{1}{2} \Bigl[ 1\cdot1 + 4\cdot(-i) + 7\cdot(-1) + 10\cdot i \Bigr] \
  &= \frac{1}{2} \Bigl[ 1 - 4i - 7 + 10i \Bigr] \
  &= \frac{1}{2} \Bigl[ (1-7) + (-4i+10i) \Bigr] \
  &= \frac{1}{2} \Bigl[ -6 + 6i \Bigr] = -3 + 3i.
  \end{aligned}
  ]

• 频率 ( k=2 ):
  此时 ( e^{-i,2\pi,2,n/4} = e^{-i,\pi,n} = (-1)^n )。
  计算：
  [
  \begin{aligned}
  X[2] &= \frac{1}{2} \Bigl[ 1\cdot1 + 4\cdot(-1) + 7\cdot1 + 10\cdot(-1) \Bigr] \
  &= \frac{1}{2} \Bigl[ 1 - 4 + 7 - 10 \Bigr] \
  &= \frac{1}{2} \Bigl[ -6 \Bigr] = -3.
  \end{aligned}
  ]
  注意 ( k=2 )（等于 seg_num/2，当 seg_num 为偶数时）对应的频谱值为实数。

3. 对其他样本和通道

对于形状中每个样本（共 4 个）以及每个通道（3 个），rfft 都是沿着 seg_num 维度分别计算。也就是说，对于第 i 个样本和第 j 个通道，其信号为
[
x[i, :, j] = [x[i,0,j],, x[i,1,j],, x[i,2,j],, x[i,3,j]]
]
计算过程与上面完全类似。如果某个信号的取值不同，则各个频率分量 ( X[k] ) 也会不同；但计算步骤始终按照下式进行：
[
X[k] = \frac{1}{\sqrt{4}} \sum_{n=0}^{3} x[i, n, j], e^{-i,2\pi,k,n/4},\quad k=0,1,2.
]

4. 结果的解释

最终，调用

x_fft = torch.fft.rfft(x, dim=1, norm='ortho')

后得到的 x_fft 张量形状为 (4, 3, 3)。其中：

• 第一维（4）对应 batch*ts_d，每个样本独立处理；
• 第二维（3）对应非负频率的个数，即 ( k=0,1,2 )；
• 第三维（3）对应每个样本的 d_model 特征。

对每个样本、每个特征通道，你都得到了其频谱表示。正交归一化保证了变换前后能量保持一致（即 Parseval 定理成立）。

总结

• 我们构造了一个符合形状 ((4,4,3)) 的数据，其中 4 表示 batch×ts_d，4 表示 seg_num，3 表示 d_model。
• 对每个样本中的每个通道，我们沿 seg_num（长度为 4）的维度进行 rfft，得到 3 个频率成分。
• 以样本 0，第 0 通道为例，我们详细计算了频率 ( k=0,1,2 ) 的傅里叶变换，分别得到 11、(-3+3i) 和 (-3)（乘以归一化因子 1/2 后的结果）。

这种方法能够推广到更高维度的输入，只要指定正确的 dim 参数，即可沿着特定维度计算 FFT。