[pytorch] 텐서의 새로운 축 만드는 방법 None indexing

pytorch tensor를 다루다 보면 새로운 축을 만들어야 할 때가 있다. 새로운 축을 만드는 방법에는 unsqueeze 함수를 사용하는 방법이 있다. ([pytorch] unsqueeze 함수, 차원 늘리기) 축을 없애는 방법에는 squeeze 함수가 있다. ([pytorch] squeeze사용법: 텐서의 차원 제거하기)축을 만들기 위해 unsqueeze 함수 대신 none indexing을 사용할 수 있다.

 

None indexing 사용방법

tensor = torch.randn(3,2,5)
print(tensor.size())#torch.Size([3, 2, 5])

print(tensor[:,None,:,:].size())#torch.Size([3, 1, 2, 5])

print(tensor[:,:,:,None,None].size())#torch.Size([3, 2, 5, 1, 1])

print(tensor[:,None,None,:,:].size())#torch.Size([3, 1, 1, 2, 5])

새로운 축을 만들고 싶은 인덱스에 None을 집어넣으면 된다.