This is a problem I've now encountered three times in the past year. I appreciate that in certain scenarios, a vectorized solution will be better and considerably faster. However IMHO, there is a trade off between discovering a vectorizing solution and using what is essentially a for loop (or double for loop). Discovering the vectorized solution (if there exists indeed one) may take a lot more effort, trial and error, research and so on. The code that is the simplest form (the double for loop in this case) almost always ends up being a bottle neck for me, but takes very little time to implement and test. Here's an example: @torch.jit.script def seq_prob(t_samples: torch.Tensor): i = 0 probs = [0] * len(t_samples) for t_i in torch.unbind(t_samples): for t_k in torch.unbind(t_samples): is_same = torch.all(torch.isclose(t_i, t_k, rtol=1e-05, atol=1e-08, equal_nan=False)) if is_same is True: probs += 1 i += 1 return probs The torch.unbind simply treats the outer dimension as the iterable. In some cases, I have spent a considerable amount of time to derive a vectorized form of the loop, which usually results in masking, cumsum, index select and a variety of built-in pytorch methods, complicating the logic compared to a for-loop but making it faster. Similarly, using CUDA and @torch.jit sometimes helps (but not always). Therefore my questions are: When using some form of for loop in pytorch (e.g., torch.unbind or torch.chunk or anything similar) whose aim is to iterate over a dimension and perform an operation Is there a way, a golden standard, some option, which can speed things up (vectorizing excluded)? If vectorizing is the only option, then what would be a good first plan of attack? Take as an example the code shown above which computes observations of a value within a sample set given some tolerance. Continue reading...
