如何在PyTorch中遍历未知形状张量及任意模型的可学习权重

1. 遍历形状未知的整个张量 #

不管张量维度是多少，你可以通过两种核心方式实现全元素遍历：

方式1：展平为一维后遍历（最简洁） #

利用flatten()或view(-1)将任意形状的张量转为一维，直接遍历每个元素：

import torch

# 示例：任意形状的张量
random_tensor = torch.randn(2, 5, 3)

# 使用flatten()展平
for elem in random_tensor.flatten():
    print(elem.item())  # 转换为Python数值输出

# 使用view(-1)展平（适合需要原地修改的场景）
for elem in random_tensor.view(-1):
    print(elem.item())

方式2：保留原始索引的遍历 #

如果需要获取每个元素在原张量中的位置索引，可以用torch.ndindex()生成所有维度的索引组合：

for idx in torch.ndindex(random_tensor.shape):
    print(f"索引{idx}对应的元素值：{random_tensor[idx].item()}")

2. 遍历并修改模型所有可学习权重 #

针对任意nn.Module的可学习参数，无需提前知道参数维度，以下是实用方案：

场景1：统一修改所有权重值 #

如果要把所有权重设为同一个值（比如你的示例中的0.5），直接用张量的原地填充方法，效率最高：

import torch
import torchvision.models as models

model = models.resnet50()

with torch.no_grad():
    for param in model.parameters():
        # 填充所有元素为0.5，自动适配任意维度
        param.data.fill_(0.5)
        # 验证修改结果（取展平后的第一个元素）
        print(param.flatten()[0].item())  # 输出0.5

场景2：逐个自定义修改权重 #

如果需要对每个元素单独处理（比如按自定义逻辑修改），展平参数后遍历即可：

with torch.no_grad():
    for param in model.parameters():
        flat_param = param.data.flatten()
        for i in range(len(flat_param)):
            # 替换为你的自定义修改逻辑，示例：将每个权重乘以0.8
            flat_param[i] = flat_param[i] * 0.8

关于递归的疑问 #

完全不需要使用递归，PyTorch的张量操作天生支持任意维度的批量处理，展平操作已经覆盖了所有维度的遍历需求，代码更简洁且执行效率更高。

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