Python中4维张量转嵌套列表 无循环合并张量高效方案咨询

完全可以不通过手写显式循环实现该需求，主流数值计算/深度学习框架（NumPy、PyTorch、TensorFlow等）都提供了高度优化的内置张量操作，底层为C/CUDA实现，性能远高于Python级手动循环。

核心实现逻辑 #

仅需两步张量操作即可得到目标结构：

• 第一步：将6个形状为(batch_size, S, S, 1)的输入张量在最后一个维度拼接，得到形状为(batch_size, S, S, 6)的中间张量
• 第二步：将中间张量的两个S维度展平为单个S*S维度，直接得到形状为(batch_size, S*S, 6)的结果张量
• 最后如果需要转为Python原生列表，调用框架内置的转列表方法即可，无需手动遍历元素

各框架代码示例 #

PyTorch 实现 #

import torch

# 示例参数，可替换为你的实际数值
batch_size = 2
S = 3
# t1~t6为你的6个输入张量
t1 = torch.randn(batch_size, S, S, 1)
t2 = torch.randn(batch_size, S, S, 1)
t3 = torch.randn(batch_size, S, S, 1)
t4 = torch.randn(batch_size, S, S, 1)
t5 = torch.randn(batch_size, S, S, 1)
t6 = torch.randn(batch_size, S, S, 1)

# 核心操作
concat_tensor = torch.cat([t1, t2, t3, t4, t5, t6], dim=-1)
result_tensor = concat_tensor.flatten(start_dim=1, end_dim=2)
# 转Python原生列表
result_list = result_tensor.tolist()

# 验证形状：print(result_tensor.shape) 输出为 torch.Size([2, 9, 6])，符合要求

TensorFlow 实现 #

import tensorflow as tf

batch_size = 2
S = 3
t1 = tf.random.normal((batch_size, S, S, 1))
t2 = tf.random.normal((batch_size, S, S, 1))
t3 = tf.random.normal((batch_size, S, S, 1))
t4 = tf.random.normal((batch_size, S, S, 1))
t5 = tf.random.normal((batch_size, S, S, 1))
t6 = tf.random.normal((batch_size, S, S, 1))

concat_tensor = tf.concat([t1, t2, t3, t4, t5, t6], axis=-1)
result_tensor = tf.reshape(concat_tensor, (batch_size, S*S, 6))
result_list = result_tensor.numpy().tolist()

NumPy 实现 #

import numpy as np

batch_size = 2
S = 3
t1 = np.random.randn(batch_size, S, S, 1)
t2 = np.random.randn(batch_size, S, S, 1)
t3 = np.random.randn(batch_size, S, S, 1)
t4 = np.random.randn(batch_size, S, S, 1)
t5 = np.random.randn(batch_size, S, S, 1)
t6 = np.random.randn(batch_size, S, S, 1)

concat_tensor = np.concatenate([t1, t2, t3, t4, t5, t6], axis=-1)
result_tensor = concat_tensor.reshape(batch_size, S*S, 6)
result_list = result_tensor.tolist()

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