如何在PyTorch中并行计算不同权重参数与输入组合下的神经网络输出？
2026-4-14
如何在PyTorch中并行计算不同权重参数与输入组合下的神经网络输出？
我正在用PyTorch做机器学习相关的工作，目前有一个继承自nn.Module的神经网络类，网络结构已经实现好了。
现在我有两个列表：一个是存储不同参数（权重、偏置）的parameters = [[...], ..., [...]]，另一个是存储不同输入的inputs = [[...], ..., [...]]，所有输入的尺寸和维度都是一致的。
我需要计算每一组输入和参数组合对应的网络输出，最终得到一个如下的结果矩阵：
| parameters[0] | parameters[1] | ... | parameters[n] | |
|---|---|---|---|---|
| input[0] | result[0][0] | result[0][1] | ... | result[0][n] |
| input[1] | result[1][0] | result[1][1] | ... | result[1][n] |
| ... | ... | ... | ... | ... |
| input[m] | result[m][0] | result[m][1] | ... | result[m][n] |
最直观的做法是用两层循环来计算，这在CPU上运行没问题，但我想在GPU上实现并行计算——毕竟每个result[i][j]都是可以独立计算的，应该能利用批量计算来提速。有没有合适的解决方案可以推荐给我？
备注：内容来源于stack exchange，提问作者Oleksandr N.
