Some stupid question about tensor operations in Pytorch

1

Hi there, I ran into a problem during the recent project. In the “forward” part, the “for” I use slowed the program a lot which I cannot accept.

截屏2022-08-01 02.24.29
截屏2022-08-01 02.24.292062×900 129 KB

As the sample shows above, how could I convert the “for” loop into matrix operation in pytorch or einsum making the program run faster ? Or how to think about problem like this ?

Sorry for my poor English.

Thanks a lot!

2

Could you post an executable code snippet with the missing definitions of e.g. couplings_lookup by wrapping it into three backticks ```, please?

3 
couplings_lookup = torch.tensor(range(1,901)).reshape(6,6,5,5)


fccl = torch.tensor([1,2,3,4,5,0])
xx = torch.tensor([3,2,2,2,1,0])
coupling = []
for idx, pos in enumerate(fccl):
    ccoupling_mask = torch.zeros(couplings_lookup.shape)
    ccoupling_mask[fccl[:idx],pos,xx[:idx],:] = 1
    # ccoupling_mask[fccl[:idx],pos,:,:] = 1
    # ccoupling_mask[:,pos,xx[:idx],:] = 1
    ccoupling = torch.sum(couplings_lookup * ccoupling_mask, axis=(0,1,2))
    print(ccoupling)
    coupling.append(ccoupling)

Thanks for the reply!

4

Maybe give me some hits about how to use “scatter” to accomplish this? Many thanks.

Sorry for the bother .