Gene Xu
Gene Xu

Reputation: 649

tensor index manipulation with "..."

Hi I'm new to Pytorch and torch tensors. I'm reading yolo_v3 code and encounter this question. I think it relates to tensor indexing with ..., but it's difficult to search ... by google, so I figure to ask it here. The code is:

        prediction = (
            x.view(num_samples, self.num_anchors, self.num_classes + 5, grid_size, grid_size)
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )
        print (prediction.shape)
        # Get outputs
        x = torch.sigmoid(prediction[..., 0])  # Center x
        y = torch.sigmoid(prediction[..., 1])  # Center y
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height
        pred_conf = torch.sigmoid(prediction[..., 4])  # Conf
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

My understanding is that the prediction will be a tensor with shape of [batch, anchor, x_grid, y_grid, class]. But what does prediction[..., x] do (x=0,1,2,3,4,5)? Is it similar as numpy indexing of [:, x]? If so the calculation of x, y, w, h, pred_conf and pred_cls don't make sense.

Upvotes: 2

Views: 1781

Answers (1)

Dishin H Goyani
Dishin H Goyani

Reputation: 7693

It's call Ellipsis. It indicate unspecified dimensions of ndarray or tensor.

Here, if prediction shape is [batch, anchor, x_grid, y_grid, class] then

prediction[..., 0] # is equivalent to prediction[:,:,:,:,0]
prediction[..., 1] # is equivalent to prediction[:,:,:,:,1]

More

prediction[0, ..., 0] # equivalent to prediction[0,:,:,:,0]

You can also write ... as Ellipsis

prediction[Ellipsis, 0]

Upvotes: 2

Related Questions