TensorFlow中complex64型3D张量转float32型(B,N,2K)张量的实现

我有一个维度为(B, N, K)的complex64类型3D张量，需要将其转换为维度为(B, N, 2*K)的float32类型张量，要求每个复数的实部和虚部相邻排列。示例如下：

原复数张量（维度(2, 3, 2)）：

[[[ 2.51-0.49j  0.80-0.74j]
  [-0.01+0.34j -2.04+0.70j]
  [ 0.02-1.85j -0.38+1.66j]]

 [[ 0.54+0.49j  0.28+1.75j]
  [-1.52-1.72j  0.68+0.17j]
  [-0.89+0.32j -1.88+0.15j]]]

转换后的浮点张量（维度(2, 3, 4)）：

[[[ 2.51  -0.49  0.80  -0.74]
  [-0.01  0.34   -2.04  0.70]
  [ 0.02 -1.85   -0.38  1.66]]

 [[ 0.54  0.49   0.28  1.75]
  [-1.52  1.72   0.68  0.17]
  [-0.89  0.32  -1.88  0.15]]]

注：示例中已减少小数位数以提升可读性。

PyTorch 实现 #

方法1：提取实部虚部后拼接 #

通过分别提取张量的实部和虚部，在最后一个维度完成拼接，是最直观的实现方式：

import torch

# 构造示例复数张量
complex_tensor = torch.tensor([
    [[2.51-0.49j, 0.80-0.74j],
     [-0.01+0.34j, -2.04+0.70j],
     [0.02-1.85j, -0.38+1.66j]],
    [[0.54+0.49j, 0.28+1.75j],
     [-1.52-1.72j, 0.68+0.17j],
     [-0.89+0.32j, -1.88+0.15j]]
], dtype=torch.complex64)

# 提取实部、虚部
real_part = complex_tensor.real  # 形状: (B, N, K)
imag_part = complex_tensor.imag  # 形状: (B, N, K)

# 在最后一维拼接，得到目标形状(B, N, 2*K)
float_tensor = torch.cat([real_part, imag_part], dim=-1)
# 或通过堆叠后展平最后两个维度实现
float_tensor = torch.stack([real_part, imag_part], dim=-1).flatten(start_dim=-2)

方法2：直接内存视图转换 #

由于complex64在内存中是以[实部, 虚部]的连续形式存储的，可直接转换数据类型后调整形状，效率更高：

float_tensor = complex_tensor.view(torch.float32).view(complex_tensor.shape[0], complex_tensor.shape[1], -1)

TensorFlow 实现 #

方法1：提取实部虚部后拼接 #

和PyTorch逻辑一致，通过TensorFlow的实部虚部提取函数完成拼接：

import tensorflow as tf

# 构造示例复数张量
complex_tensor = tf.constant([
    [[2.51-0.49j, 0.80-0.74j],
     [-0.01+0.34j, -2.04+0.70j],
     [0.02-1.85j, -0.38+1.66j]],
    [[0.54+0.49j, 0.28+1.75j],
     [-1.52-1.72j, 0.68+0.17j],
     [-0.89+0.32j, -1.88+0.15j]]
], dtype=tf.complex64)

# 提取实部、虚部
real_part = tf.math.real(complex_tensor)
imag_part = tf.math.imag(complex_tensor)

# 拼接得到目标张量
float_tensor = tf.concat([real_part, imag_part], axis=-1)
# 或堆叠后展平
float_tensor = tf.reshape(tf.stack([real_part, imag_part], axis=-1), shape=[tf.shape(complex_tensor)[0], tf.shape(complex_tensor)[1], -1])

方法2：类型转换后调整形状 #

利用TensorFlow的类型转换和形状调整函数，直接完成格式转换：

float_tensor = tf.reshape(tf.cast(complex_tensor, tf.float32), shape=[tf.shape(complex_tensor)[0], tf.shape(complex_tensor)[1], -1])

内容的提问来源于stack exchange，提问作者Alveena