Python 实现 TensorFlow 和 PyTorch 验证卷积 convolution 函数矩阵化计算及反向传播

摘要

本文使用纯 Python 实现 TensorFlow 和 PyTorch 验证卷积 convolution 函数矩阵化计算及反向传播.

相关

原理和详细解释, 请参考文章 :

卷积 convolution 函数的矩阵化计算方法及其梯度的反向传播

系列文章索引 :
https://blog.youkuaiyun.com/oBrightLamp/article/details/85067981

正文

1. Batch2ConvMatrix类和Conv2d类

文件目录 : vanilla_nn/convolution_matrix.py

import numpy as np


class Batch2ConvMatrix:
    def __init__(self, stride, kernel_h, kernel_w):
        self.stride = stride
        self.kernel_h = kernel_h
        self.kernel_w = kernel_w

        self.x = None
        self.conv_size = None

    def __call__(self, x):
        self.x = x
        x_nums, x_channels, x_height, x_width = np.shape(self.x)

        conv_height = int((x_height - self.kernel_h) / self.stride) + 1
        conv_width = int((x_width - self.kernel_w) / self.stride) + 1

        scan = np.zeros((x_nums, conv_height, conv_width,
                         x_channels, self.kernel_h, self.kernel_w))

        for n in range(x_nums):
            for h in range(conv_height):
                for w in range(conv_width):
                    for c in range(x_channels):
                        start_h = h * self.stride
                        start_w = w * self.stride
                        end_h = start_h + self.kernel_h
                        end_w = start_w + self.kernel_w

                        scan[n, h, w, c] = \
                            x[n, c, start_h:end_h, start_w:end_w]

        conv_matrix = scan.reshape(x_nums * conv_height * conv_width, -1)
        self.conv_size = [x_nums, x_channels, conv_height, conv_width]
        return conv_matrix

    def backward(self, dx2m):
        dx = np.zeros_like(self.x)
        kh = self.kernel_h
        kw = self.kernel_w
        xn, xc, ch, cw = self.conv_size

        dx2m = dx2m.reshape((xn, ch,