pytorch查询字典、列表维度

时间:2025-04-14 07:11:33

输出tensor变量维度

print(a.shape)

输出字典维度

for key, value in output_dict.items():
    if isinstance(value, torch.Tensor):
        print(f"{key} shape:", value.shape)

输出列表维度

def get_list_dimensions(lst):
    # 基线条件:如果lst不是列表,返回空列表
    if not isinstance(lst, list):
        return []
    # 基线条件:如果lst是空列表,返回[0]
    if not lst:
        return [0]
    # 递归步骤:检查子列表的维度
    sub_dimensions = [get_list_dimensions(item) for item in lst if isinstance(item, list)]
    # 如果没有子列表,则当前列表是一维的
    if not sub_dimensions:
        return [len(lst)]
    # 确定所有子列表的最大维度
    max_dim = max(sub_dimensions, key=len)
    # 返回当前列表长度和子列表的最大维度
    return [len(lst)] + max_dim

# 测试函数
my_list = [[1, 2], [3, 4]]
print(get_list_dimensions(my_list))  # 应该输出 [2, 2]