Pytorch 计算卷积输出的大小

2021年2月10日07:54:31

评论1 1827字阅读6分5秒

摘要这一篇文章会给出两个计算函数，第一个可以方便根据 input size，kernel size，stride 和 padding 计算出 output size 的大小。第二个可以根据 input size 和 output size 计算出 padding 的大小。

文章目录(Table of Contents)

简介

在 Pytorch 中使用卷积会有一个比较麻烦的点，就是当我们设置了 kernel size、stride、pad 之后，不是很好计算 output 的大小。虽然 Pytorch 中提供了计算的式子，但是每次手动计算确实不是很方便。

在这一篇文章中，会将上面的式子写成两个两个函数方便使用：

首先比较常用的是，给定了 input size，stride，padding，kernel size 之后计算出 output size 的大小。例如有下面两个函数。

import math
def num2tuple(num):
return num if isinstance(num, tuple) else (num, num)
def conv2d_output_shape(h_w, kernel_size=1, stride=1, pad=0, dilation=1):
h_w, kernel_size, stride, pad, dilation = num2tuple(h_w), \
num2tuple(kernel_size), num2tuple(stride), num2tuple(pad), num2tuple(dilation)
pad = num2tuple(pad[0]), num2tuple(pad[1])
h = math.floor((h_w[0] + sum(pad[0]) - dilation[0]*(kernel_size[0]-1) - 1) / stride[0] + 1)
w = math.floor((h_w[1] + sum(pad[1]) - dilation[1]*(kernel_size[1]-1) - 1) / stride[1] + 1)
return h, w

首先，利用函数 conv2d_output_shape 来 output 的大小。如下所示：

还有一个比较常用的是给定了 output size 的大小，想要计算 padding 的大小。

import math
def num2tuple(num):
return num if isinstance(num, tuple) else (num, num)
def conv2d_get_padding(h_w_in, h_w_out, kernel_size=1, stride=1, dilation=1):
h_w_in, h_w_out, kernel_size, stride, dilation = num2tuple(h_w_in), num2tuple(h_w_out), \
num2tuple(kernel_size), num2tuple(stride), num2tuple(dilation)
p_h = ((h_w_out[0] - 1)*stride[0] - h_w_in[0] + dilation[0]*(kernel_size[0]-1) + 1)
p_w = ((h_w_out[1] - 1)*stride[1] - h_w_in[1] + dilation[1]*(kernel_size[1]-1) + 1)
return (math.floor(p_h/2), math.ceil(p_h/2)), (math.floor(p_w/2), math.ceil(p_w/2))