CODEWITHSUNDEEP
pythonpytorch
Yan W
Yan W

Reputation: 21

how the code '-input[range(target.shape[0]),target]' works?

I'm learing pytorch.Reading the official tutorial,I met the preplexing code. input is a tensor, so is target.

def nll(input,target):
    return -input[range(target.shape[0]),target].mean()

And the pred is:
pred

target is:
target

the '-input[range(target.shape[0]),target]' is:
'input[range(target.shape[0],target]'

Output shows this is not substracting target from input or merging two tensors

Upvotes: 1

Views: 208

Answers (2)

Toye Kunle
Toye Kunle

Reputation: 1

I encountered similar code somewhere and I was curious as to how it works:

My thought: I think it follows general rule of slicing.

-input for example = 64,10 # 64 different outputs with 10 individual probabilities

-target = 64 # 64 targets for 64 samples, bunch of indexes

-input[row,col] ~ input[range(64), target]

-the range(64) ensures all the rows are considered while the target picks the column index of each row.

-That said, it is a smart way of writing cross-enrtropy code since all you care is the probab

Upvotes: 0

Shai
Shai

Reputation: 114786

The code input[range(target.shape[0]), target] simply picks, from each row i of input the element at column indicated by the corresponding element of target, that is target[i].
In other words, if out = input[range(target.shape[0]), target] then out[i] = input[i, target[i]].

This is very similar to torch.gather.

Upvotes: 2

Related Questions