Python 卷积矩阵的Propper展开

我目前正在从头开始实施一个litle深度学习框架，因为它很有趣。 显然，一个不到1000年的框架有一个2D卷积的实现。因为我希望至少能够用它来训练一些东西，所以用嵌套循环进行卷积不是一个好主意。 这就是为什么我要展开我的输入矩阵并将其与内核相乘。我已经成功地实现了形状为

（批量大小、dims、img高度、宽度）

的矩阵的前向传递。 现在的问题是，我只能通过大量的换位和重塑来实现，当它进入后传球时，基本上一切都会中断。 右键知道我正在使用内核展开输入

x=（bs，dim，h，w）

K=（K，K）

进入形状

unrolled=（bs，dim，h*w，k**2）

，它看起来很奇怪，但在前传球中仍然有效，但正如我所说的，在后传球中不起作用（可能用它实现工作后传球，但我希望它干净），所以 现在的问题是：使用

shape=（bs，dim，h，w）

to展开输入的支撑形状是什么？如何快速有效地实现这一点

感谢您抽出时间，祝您度过愉快的一天：）