# 维度转换

import torch

from einops import rearrange

print(torch.cuda.is_available())

 

# to_3d

# 把四维的张量转换为三维的张量，输入形状（b,c,h,w），输出形状（b,hw,c）

def to_3d(x):

    return rearrange(x, 'b c h w -> b (h w) c')

# to_4d

# 把三维的张量转换为四维的张量，输入形状（b,hw,c），输出形状（b,c,h,w）

def to_4d(x,h,w):

    return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)

# 测试

if __name__ == '__main__':

    # 创建一个四维张量

    tensor_4d = torch.randn(2, 3, 4, 5) # 形状（批大小2, 通道数3, 高度4, 宽度5）

    # 转换为三维张量

    tensor_3d = to_3d(tensor_4d)     # 形状（批大小2, 20, 3）

    # 转换为四维张量

    height, width = 4, 5

    tensor_4d_back = to_4d(tensor_3d, height, width)  # 形状（批大小2, 通道数3, 高度4, 宽度5）

    print(tensor_4d.shape)  

    print(tensor_3d.shape)  # 输出：torch.Size([2, 20, 3])

    print(tensor_4d_back.shape)  # 输出：torch.Size([2, 3, 4, 5])