动手实现卷积

手写一个卷积操作

主要目的在于熟悉torch中的indexing操作。

首先,要介绍torch中的unfold操作。首先这个函数的输入必须是一个4-dimensional的tensor。还要设置

一些常用的参数,例如kernel_size, stride等等。然后就会被resize成为一个(B, Channel * kernel_size[0] * kernel_size[1], num_of_sifts)的一个tensor。fold操作恰恰相反,有了这些前置知识就可以写出一个卷积操作了。

import torch
import torch.nn.functional as F

input_tensor = torch.randn(2, 3, 4, 4)
conv_kernel = torch.rand(3, 3, 3)

unfold_tensor = F.unfold(input_tensor, (3, 3)) # fold的参数是(input_tensor, (kernel_size[1], kernel_size[2])
conv_kernel_after = conv_kernel.view(conv_kernel.size(0), -1).repeat(1, 3).transpose(0, 1)

multi_tensor = torch.matmul(unfold_tensor.transpose(1, 2), conv_kernel_after).transpose(1, 2)

result = F.fold(multi_tensor, (2, 2), 1) # fold的参数(input, (output_dim[0], output_dim[1]), stride)

conv_kernel_after = conv_kernel.unsqueeze(dim=1).repeat(1,3,1,1)

print((F.conv2d(input_tensor, conv_kernel_after) - result).abs().max())
posted @ 2021-03-17 10:41  John_Ran  阅读(232)  评论(0编辑  收藏  举报