Custom Convolution

$BatchSize = N$
$Input Feature Spatial Size = H_{in} \times W_{in}$
$Output Feature Spatial Size = H_{out} \times W_{out}$
$Kernel Spatial Size = H_{kernel} \times W_{kernel}$
$Input Channels = C_{in}$
$Output Channels = C_{out}$

in pytorch,
$Input Feature Size = (N,C_{in},H_{in},W_{in})$
$Output Feature Size = (N,C_{out},H_{out},W_{out})$
$Kernel Size = (C_{out},C{in},H_{kernel},W_{kernel})$

unfold input feature, $size = (N,C_{in} \times H_{kernel} \times W_{kernel}, H_{out} \times W_{out})$
after transpose, $size = (N, H_{out} \times W_{out},C_{in} \times H_{kernel} \times W_{kernel})$

kernel flatten, $size = (C_{out},C{in} \times H_{kernel} \times W_{kernel})$
after transpose, $size = (C{in} \times H_{kernel} \times W_{kernel},C_{out})$

after matmul, $size = (N,H_{out} \times W_{out},C_{out})$
after transpose, $size = (N,C_{out},H_{out} \times W_{out})$

after fold, $size = (N,C_{out},H_{out},W_{out})$