einops 张量处理

最新推荐文章于 2024-10-17 13:50:17 发布

CV温故知新

最新推荐文章于 2024-10-17 13:50:17 发布

阅读量773

点赞数

文章标签：深度学习 pytorch cnn

本文链接：https://blog.csdn.net/lansebingxuan/article/details/122495763

版权

安装

einops支持 python 3.5 及更高版本，安装命令 pip install einops

代码链接：arogozhnikov/einops: Deep learning operations reinvented (for pytorch, tensorflow, jax and others) (github.com)

用法

可以用来调整tensor的维度

from einops import rearrange

img = torch.randn(1,3,256,256)
Rearrange(img,'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = 32, p2 = 32)

Rearrange用于维度调整，( )表示里面的两个数相乘，中间维度看作h*p1，然后给出p1的数值，这样系统会自动把中间那个维度拆解。

以上面代码为例，输入img大小为(1,3,256,256)，p1 = 32，256自动拆成8*32，因此h=w=8，维度由(b,c,h*p1,w*p2)这个四维矩阵调整为(b,h*w,p1*p2*c)的三维矩阵，即(1,64,3072)

layer构建

Rearrange可以直接用于构建模型的一个层

from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Rearrange

model = Sequential(
    Conv2d(3, 6, kernel_size=5),
    MaxPool2d(kernel_size=2),
    Conv2d(6, 16, kernel_size=5),
    MaxPool2d(kernel_size=2),
    # flattening
    Rearrange('b c h w -> b (c h w)'),  
    Linear(16*5*5, 120), 
    ReLU(),
    Linear(120, 10), 
)