Hi, I have encountered a problem that relates to the time consumption of the indexing operation.
I have a tensor which is ‘output:[B,C,H,W]’ and I use unfold to transform it like ‘out_patch:[B,C,H//N,W//N,N,N]’. N is the patch size.
And there is a binary mask tensor ‘mask:[B,C,H//N,W//N]’ so that I can select patches from out_patch by: ‘patch=out_patch[mask]’.
There is a similar operation like ‘patch=out_patch[mask]’ in the latter part of the code but the time it consumed is nearly 100 times of the former one. I have tested many times but the problem still occurs.
Later, I found the problem was that there are some nn.conv2d operations running before the indexing operation and the whole code is running on GPU. If I put the indexing operation before the convolutions or run the code on CPU, the time would be the same as the former one.
I wonder what causes this phenomenon and is there any way to avoid it? As in my method, the latter indexing operation has to be put behind the convolutions and the code has to be run on GPU. Running Time is very important to me. Thanks.