pytorch学习笔记--手写conv2d-flatten版本

参考b站up主 deep_thoughts

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
#不考虑batchsize channel  Conv2d  flatten 版本
def ManualConv2d_MatrixCaculation(input , kernel , bias = 0 , stride = 1 , padding = 0 ):
     if(padding > 0):
          input = F.pad(input , (padding ,padding ,padding ,padding))
     input_h , input_w = input.shape
     kernel_h , kernel_w = kernel.shape
    
     output_w  = (math.floor((input_h - kernel_h)/stride) + 1) # conv output wide 
     output_h  = (math.floor((input_w - kernel_w)/stride) + 1) # conv output height
     output = torch.zeros(output_h , output_w) # initial output matrix 
     
     region_matrix = torch.zeros(output.numel() , kernel.numel()) # 存储所由的拉平后的特征区域
     kernel_matrix = kernel.reshape((kernel.numel() , 1))  # kernel 的列向量的形式
     row_index = 0
     for i in range(0  , input_h - kernel_h + 1 , stride): # 对高度维进行遍历
          for j in range(0 , input_w - kernel_w + 1 , stride): # 对宽度维进行遍历
               region = input[i : i + kernel_w , j : j + kernel_w] # kernel 正在覆盖的区域
               region_vector = torch.flatten(region)
               region_matrix [row_index] = region_vector
               row_index += 1
     output_matix = region_matrix @ kernel_matrix
     output = output_matix.reshape((output_h , output_w)) + bias
     return output

input = torch.randn(5 ,5)
kernnel = torch.randn(3 ,3)
bias = torch.randn(1)

mat_mul_conv_output = ManualConv2d_MatrixCaculation(input , kernnel , bias = bias , padding= 1)

pytorch_api_conv_output = F.conv2d(input.reshape((1 , 1, input.shape[0] , input.shape[1])),
                                                 kernnel.reshape((1 , 1, kernnel.shape[0] , kernnel.shape[1])),
                                                 padding = 1 ,
                                                 bias = bias ).squeeze(0).squeeze(0)

flag = torch.allclose(mat_mul_conv_output , pytorch_api_conv_output)
print(flag)
print(mat_mul_conv_output)
print(pytorch_api_conv_output)