개요
PyTorch에서 텐서(tensor)의 요소들을 정렬하려면 torch.sort() 메서드를 사용하면 됩니다. 이 메서드는 두 개의 텐서를 반환하는데, 첫 번째 텐서는 정렬된 요소 값들을 담고 있으며, 두 번째 텐서는 원본 텐서에서 각 요소가 위치했던 인덱스 정보를 담고 있습니다.
기본적으로 정렬은 오름차순(ascending order)으로 수행되며, 내림차순으로 정렬하고 싶다면 descending=True 옵션을 추가하면 됩니다. 또한 2차원 텐서의 경우 dim 매개변수를 활용해 행(row) 단위 또는 열(column) 단위로 정렬할 수 있습니다.
정렬 절차
필요한 라이브러리 임포트: 아래 모든 예제에서 필요한 파이썬 라이브러리는 torch입니다. 미리 설치되어 있는지 확인하세요.
텐서 생성: PyTorch 텐서를 생성한 뒤 화면에 출력합니다.
torch.sort(input, dim) 호출: 생성된 텐서를 정렬하기 위해 torch.sort(input, dim)을 계산하고 그 결과를 새로운 변수 v에 할당합니다. 여기서 input은 입력 텐서이며, dim은 정렬이 수행될 차원입니다. 행 단위로 정렬하려면 dim=1, 열 단위로 정렬하려면 dim=0으로 설정합니다.
결과 접근: 정렬된 값들이 담긴 텐서는 v[0]으로, 정렬된 요소들의 원래 인덱스 텐서는 v[1]로 접근할 수 있습니다.
결과 출력: 정렬된 값 텐서와 인덱스 텐서를 화면에 출력합니다.
예제 1: 1차원 텐서 정렬
다음 파이썬 프로그램은 1차원(1D) 텐서의 요소들을 정렬하는 방법을 보여줍니다.
# Python program to sort elements of a tensor
# import necessary library
import torch
# Create a tensor
T = torch.Tensor([2.334,4.433,-4.33,-0.433,5, 4.443])
print("Original Tensor:\n", T)
# sort the tensor T
# it sorts the tensor in ascending order
v = torch.sort(T)
# print(v)
# print tensor of sorted value
print("Tensor with sorted value:\n", v[0])
# print indices of sorted value
print("Indices of sorted value:\n", v[1])출력 결과
Original Tensor:
tensor([ 2.3340, 4.4330, -4.3300, -0.4330, 5.0000, 4.4430])
Tensor with sorted value:
tensor([-4.3300, -0.4330, 2.3340, 4.4330, 4.4430, 5.0000])
Indices of sorted value:
tensor([2, 3, 0, 1, 5, 4])출력 결과에서 볼 수 있듯이, 음수를 포함한 실수 요소들이 오름차순으로 정렬되었으며, 두 번째 텐서에는 각 정렬된 값이 원본 텐서에서 어느 위치에 있었는지를 나타내는 인덱스가 저장됩니다.
예제 2: 2차원 텐서 정렬
다음 파이썬 프로그램은 2차원(2D) 텐서의 요소들을 열 단위(dim=0)와 행 단위(dim=1)로 각각 정렬하는 방법을 보여줍니다.
# Python program to sort elements of a 2-D tensor
# import the library
import torch
# Create a 2-D tensor
T = torch.Tensor([[2,3,-32],
[43,4,-53],
[4,37,-4],
[3,-75,34]])
print("Original Tensor:\n", T)
# sort tensor T
# it sorts the tensor in ascending order
v = torch.sort(T)
# print(v)
# print tensor of sorted value
print("Tensor with sorted value:\n", v[0])
# print indices of sorted value
print("Indices of sorted value:\n", v[1])
print("Sort tensor Column-wise")
v = torch.sort(T, 0)
# print(v)
# print tensor of sorted value
print("Tensor with sorted value:\n", v[0])
# print indices of sorted value
print("Indices of sorted value:\n", v[1])
print("Sort tensor Row-wise")
v = torch.sort(T, 1)
# print(v)
# print tensor of sorted value
print("Tensor with sorted value:\n", v[0])
# print indices of sorted value
print("Indices of sorted value:\n", v[1])출력 결과
Original Tensor:
tensor([[ 2., 3., -32.],
[ 43., 4., -53.],
[ 4., 37., -4.],
[ 3., -75., 34.]])
Tensor with sorted value:
tensor([[-32., 2., 3.],
[-53., 4., 43.],
[ -4., 4., 37.],
[-75., 3., 34.]])
Indices of sorted value:
tensor([[2, 0, 1],
[2, 1, 0],
[2, 0, 1],
[1, 0, 2]])
Sort tensor Column-wise
Tensor with sorted value:
tensor([[ 2., -75., -53.],
[ 3., 3., -32.],
[ 4., 4., -4.],
[ 43., 37., 34.]])
Indices of sorted value:
tensor([[0, 3, 1],
[3, 0, 0],
[2, 1, 2],
[1, 2, 3]])
Sort tensor Row-wise
Tensor with sorted value:
tensor([[-32., 2., 3.],
[-53., 4., 43.],
[ -4., 4., 37.],
[-75., 3., 34.]])
Indices of sorted value:
tensor([[2, 0, 1],
[2, 1, 0],
[2, 0, 1],
[1, 0, 2]])마무리
torch.sort() 메서드를 사용하면 1차원뿐만 아니라 다차원 텐서의 요소도 손쉽게 정렬할 수 있습니다. dim 매개변수로 행 또는 열 단위 정렬을 지정하고, 필요하다면 descending=True 옵션으로 내림차순 정렬도 가능합니다. 정렬된 값과 함께 원본 인덱스까지 반환되므로, 데이터 분석이나 머신러닝 전처리 과정에서 순위 기반 연산을 구현할 때 매우 유용하게 활용할 수 있습니다.