PyTorch复制子张量：将[100,1024,14,14]张量扩展为[1000,1024,14,14]

嘿，这个需求用PyTorch的内置张量操作就能完美解决，我给你分两种场景来讲解——优先满足你想要的不创建副本、共享内存的最优方案，再补充一个允许创建副本的简洁方法：

方法一：不创建副本（共享原张量内存） #

这个方案利用PyTorch的unsqueeze和expand操作，二者都是视图变换，不会额外分配内存，完全共享原张量的数据：

import torch

# 假设你的原始张量已定义
image_features = torch.randn(100, 1024, 14, 14)

# 1. 在第1个维度插入长度为1的新维度，形状变为 (100, 1, 1024, 14, 14)
expanded_tensor = image_features.unsqueeze(1)
# 2. 将新插入的维度扩展为10，形状变为 (100, 10, 1024, 14, 14)，此步骤不复制数据
expanded_tensor = expanded_tensor.expand(-1, 10, -1, -1, -1)  # -1表示保留原维度大小
# 3. 合并前两个维度，得到最终目标形状 (1000, 1024, 14, 14)
result = expanded_tensor.flatten(0, 1)

# 验证效果
print(result.shape)  # 输出: torch.Size([1000, 1024, 14, 14])
# 检查内存是否共享：前10个元素都和原张量第一个子张量共享内存
print(result[0].data_ptr() == image_features[0].data_ptr())  # 输出: True
print(result[9].data_ptr() == image_features[0].data_ptr())  # 输出: True

原理说明 #

• unsqueeze(1)：给每个子张量增加一个“虚拟”维度，让PyTorch可以对这个维度进行扩展操作。
• expand：只修改张量的视图信息，不会复制数据——相当于告诉PyTorch：“这个维度的元素和前一个位置的元素是同一个”，因此内存完全共享。
• flatten(0,1)：把前两个维度（100和10）合并成一个1000的维度，最终形状符合需求。

方法二：允许创建副本的简洁方案 #

如果后续需要对重复后的子张量单独修改（共享内存时修改会影响原张量），或者追求代码简洁，可以使用repeat_interleave，一行代码完成：

result = image_features.repeat_interleave(10, dim=0)

print(result.shape)  # 输出: torch.Size([1000, 1024, 14, 14])

说明 #

repeat_interleave会将原张量的每个元素在指定维度（这里是dim=0）重复指定次数，生成的新张量是原数据的副本，内存占用为原张量的10倍。这种方式代码更短，但会消耗更多内存。

另外，也可以用repeat实现类似效果（同样会创建副本）：

result = image_features.unsqueeze(1).repeat(1, 10, 1, 1, 1).flatten(0, 1)

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