Linear层（矩阵乘法）的优化

综合讨论区
#1

我在实现nerf（神经辐射场）with taichi, 有一个步骤是Linear层，目前的实现是参考官方的:
https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/autodiff/diff_sph/diff_sph.py#L57

但感觉还可以优化，能否给点建议？

矩阵大小： [4096*128, 128] @ [128, 128]

2 个赞
#2

你好，这里是希望有一个很大的batch size吗？

#3

对，batch是1e6量级的， 向量大小大约是128左右

#4

这是我现在的实现，由于batchsize过大，ndrange会爆int32范围，我给除以了4，放到循环里面了。
https://github.com/Feynman1999/taichi-nerf/blob/main/train.py#L139

1 个赞
#5

好的，我来结合代码看一下

1 个赞
#6

我现在只实现了前向的过程，还没有测试梯度部分。在cpu上我用profile看，90%以上的时间都在linear层上。gpu上的情况我还没分析。