Search code examples
pythonperformancepytorch

How to efficiently process tensors with a batch dimension in PyTorch?

I have a tensor with bounding boxes of the shape [batch, n_boxes, 10]. My goal is to preprocess those boxes in the first and second dimension. I do the following which is horribly slow.

      for b in range(tensor.shape[0]):
            for idx in range(tensor.shape[1]):
                tensor[b ,idx, 0] = 2 * ((tensor[b, idx, 0] - x_min) / (x_max - x_min)) - 1
                tensor[b, idx, 1] = 2 * ((tensor[b, idx, 1] - y_min) / (y_max - y_min)) - 1

For the preprocessing I can use numpy, but for the postproessing not, as it happens right before the loss computation and I need to keep the gradient history intact.

How can I do this more efficiently in PyTorch?

Solution

  • I think this should do the trick:

    tensor[...,0] = 2*(tensor[...,0] - x_min) / (x_max - x_min)) - 1
tensor[...,1] = 2*(tensor[...,1] - y_min) / (y_max - y_min)) - 1