Pytorch手动实现卷积

某些时候可能需要用到类似卷积操作滑动窗的思想，实现一些操作。具体实现如下，首先导入必要的包

import torch
import numpy as np
import torch.nn as nn

方便起见（同时为了验证方法正确性），构建简单的卷积输入

input = torch.from_numpy(np.array([1,2,3,4,5,6,7,8,9])).view((1,1,3,3)).float()

查看输入内容

In [5]: input.shape                                                                                                                                                                                                                        
Out[5]: torch.Size([1, 1, 3, 3])

In [6]: input                                                                                                                                                                                                                              
Out[6]: 
tensor([[[[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]]])

其中，输入张量存储方式按照pytorch正常卷积操作，即BxCxKxK

手动构造卷积核

In [7]: kernel = torch.ones((1,2,2,1))                                                                                                                                                                                                     

In [8]: kernel                                                                                                                                                                                                                             
Out[8]: 
tensor([[[[1.],
          [1.]],

         [[1.],
          [1.]]]])

卷积核尺寸为CxKxKxCout。

接下来是手动实现卷积的核心，利用Unfold函数，Unfold可以将输入数据按照卷积时候的滑动窗cube方式展开

unfold = nn.Unfold(kernel_size=2) 

对输入进行操作

unin = unfold(input) 

查看Unfold之后的尺寸

In [14]: unin.shape                                                                                                                                                                                                                        
Out[14]: torch.Size([1, 4, 4])

其各个轴分别为Bx(K*K)xN，其中N是滑动窗的数量。

为了实现卷积，需要对Unfold之后的数据重新调整维度，如下

In [21]: unin = unin.view((1,1,-1,2,2))                                                                                                                                                                                                    

In [22]: unin.shape                                                                                                                                                                                                                        
Out[22]: torch.Size([1, 1, 4, 2, 2])

其中，unin的尺寸为BxCx(K*K)xH'xW'

对卷积核重新调整维度

In [24]: kernel = kernel.view((1,-1,1))                                                                                                                                                                                                    

In [25]: kernel.shape                                                                                                                                                                                                                      
Out[25]: torch.Size([1, 4, 1])

其中kernel的维度为Cin x (K*K) x Cout

维度调整好后，实现卷积的乘积求和操作，采用einsum求和方式

In [26]: out = torch.einsum('ijkmn,jkl->ilmn',unin,kernel)                                                                                                                                                                                 

In [27]: out.shape                                                                                                                                                                                                                         
Out[27]: torch.Size([1, 1, 2, 2])

查看输入数据及卷积后的结果

In [29]: input                                                                                                                                                                                                                             
Out[29]: 
tensor([[[[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]]])

In [28]: out                                                                                                                                                                                                                               
Out[28]: 
tensor([[[[12., 16.],
          [24., 28.]]]])