pytorch(3d、4d张量转换)

维度转换

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])

相关推荐
爱偷懒的。。1 天前
基于 WebSocket 协议的实时弹幕通信机制分析-抖音
网络·python·websocket·网络协议·学习·js
cllsse1 天前
pytest学习
软件测试·python·pytest
海琴烟Sunshine1 天前
leetcode 88.合并两个有序数组
python·算法·leetcode
B站_计算机毕业设计之家1 天前
数据分析:Python懂车帝汽车数据分析可视化系统 爬虫(Django+Vue+销量分析 源码+文档)✅
大数据·爬虫·python·数据分析·汽车·可视化·懂车帝
一条星星鱼1 天前
深度学习中的归一化:从BN到LN到底是怎么工作的?
人工智能·深度学习·算法·归一化
Rock_yzh1 天前
AI学习日记——深度学习
人工智能·python·深度学习·神经网络·学习
DrugOne1 天前
GAMESS 在 Ubuntu 24.04 平台上的编译与配置
linux·python·ubuntu·drugone
路长冬1 天前
记录一次深度学习+SSH的配置
深度学习
C嘎嘎嵌入式开发1 天前
(5)100天python从入门到拿捏------流程控制语句
服务器·网络·python
Siren_dream1 天前
python进阶_Day5
python