维度转换
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])