torch.where(), 근데 이제 loss를 만들 때 많이 곁들인

🤔 Problem

PyTorch에서 자주 쓰는 메서드들은 어느 정도 체득하면서 빠르게 핸들링할 수 있도록 하는 것이 좋은 거 같습니다. 이번 글은 그러한 메서드들 중에 자주 쓰이는 torch.where()에 대해서 정리해 보았습니다.

 

사실은 매우 간단합니다만, 완전히 체득을 해두는 게 앞으로 작업에 편할 거 같아서 일부러 글을 씁니다.

out = torch.where(condition, input, other)

condition은 BoolTensor입니다. True와 False만을 담고 있으며, True인 element 자리에 input이 들어가고, False인 element 자리에 other이 들어가게 됩니다.

1. condition은 Boolean Masking으로 나타내기

보통 condition은 원래의 tensor에 Boolean masking을 하여, 하나의 BoolTensor로 나타냅니다.

(Boolean masking이란 Broadcasting을 통해 다차원 배열 혹은 텐서에서 스칼라 값에 대한 조건을 만족하면, 해당 자리에 True, 반대이면 False로 처리하는 기법을 의미합니다.) 

그래서, torch.where()를 쓸 때 자주 쓰는 것은 Boolean Masking으로 condition을 나타내는 것입니다.

a = torch.tensor([1, 2, 3, 4, 5])
condition = a > 2
print(a)
>>> tensor([False, False, True, True, True])

2. input과 other는 꼭 스칼라 값이 아니어도 됩니다.

굳이 글로 남기려는 이유는 이 스킬 때문입니다. 위 글만 보면, input과 other는 특정 스칼라 값이어야만 할 거 같습니다. 

공식 문서를 보면, 꼭 그렇지는 않습니다. input과 output이 스칼라 값이면 스칼라 값을 그대로 넣어주는 것도 맞지만, tensor라면 해당 위치(indices)에 대한 tensor의 값을 넣어줄 수 있다고 명시되어 있습니다.

그래서, Boolean masking 이전 원래 tensor에 대해서 '내 조건을 만족한다면 그대로 놔두고, 아니라면 모두 0으로 처리해라' 같은 작업을 수행하고 싶다면, 아래와 같은 코드로 간단하게 처리가 가능합니다.

x = torch.tensor([~~~~~])
condition = x > 0 # tensor element에서 0보다 높으면 True
out = torch.where(condition, x, 0)

이를 조금 더 응용해서 원하는 위치에 대해서만 특정 연산을 수행하는 복잡한 코드도 한 줄이면, 작성이 가능합니다.

 

이러한 torch.where()는 특정 조건에 상응하는 손실 값을 반영시키기 위해서 Loss function을 구현할 때, 많이 사용합니다.

---

📂 Reference

https://pytorch.org/docs/stable/generated/torch.where.html

torch.where — PyTorch 2.3 documentation