오늘은 tensor를 다루는 여러 기능을 가진 패키지인 einops에 대해 알아보겠습니다. Einops는 Pytorch보다 직관적이고 간단하게 tensor 연산이 가능하다는 장점이 있어 ViT를 비롯한 최근 논문에서 자주 쓰이고 있습니다. 오늘은 Einops의 여러 기능 중 rearrange, reduce, repeat 을 이용한 간단한 tensor 연산을 수행해보겠습니다.
0. 예시로 사용할 MNIST Dataset 불러오기
Einops를 이용해 Tensor연산을 수행할 예시 이미지로 MNIST DATASET을 사용하기 위해 torchvision의 MNIST dataset을 불러옵니다.
from torchvision.datasets import MNIST
train_mnist = MNIST(download_root,transform=mnist_transform,train=True,download=True)MNIST dataset은 손글씨 숫자 데이터로써 각 image 당 (1,28,28)의 shape을 갖습니다.
image, label = train_mnist[10]
plt.imshow(image.permute(1,2,0), cmap = 'gray')
본 포스팅에서는 6개의 손글씨 숫자 이미지를 활용하겠습니다.
stack_image = torch.stack([train_mnist.__getitem__(10)[0],
train_mnist.__getitem__(1)[0],
train_mnist.__getitem__(3)[0],
train_mnist.__getitem__(4)[0],
train_mnist.__getitem__(5)[0],
train_mnist.__getitem__(8)[0]],dim=0)
stack_image.shape #torch.Size([6, 1, 28, 28])6개의 image를 0차원 기준으로 쌓아 (6,1,28,28) size의 tensor가 생성되었습니다. tensor의 각 차원은 batch,channel, height, width를 의미하며 앞으로 (b, c, h, w)와 같이 표현할 예정입니다.
1. rearrange 함수 설명
rearrange 함수는 tensor의 차원을 결합하거나 재구성하여 tensor의 모양을 변경할 때 활용됩니다.
첫번째 예시로 (1,28,28) 이었던 image tensor를 (28,28,1)의 (h, w, c) tensor로 바꿔보겠습니다.
from einops import rearrange
#pytorch
image.permute(1,2,0)
#einops
rearrange(image,'c h w -> h w c')코드는 pytorch가 더 짧을지라도 tensor의 모양이 어떻게 바뀐건지에 대한 직관적인 인식은 einops가 더 쉽고 간편합니다. rearrange 함수를 통해서 여러 이미지를 가로, 세로 기준으로 결합시키는 것도 가능합니다.
h_stack_image = rearrange(imgs,'b h w c -> (b h) w c') #(b h): b,h 간의 flatten
plt.imshow(h_stack_image, cmap = 'gray')
print(imgs.shape,h_stack_image.shape) #전: torch.Size([6, 28, 28, 1]) 후: torch.Size([168, 28, 1])
세로기준으로 결합
w_stack_image = rearrange(imgs,'b h w c -> h (b w) c')
plt.imshow(w_stack_image,cmap = 'gray')
print(imgs.shape,w_stack_image.shape) #전: torch.Size([6, 28, 28, 1]) 후: torch.Size([28, 168, 1])
가로기준으로 결합
또한 한줄로만 나열하는 것이 아닌 23의 이미지 행렬, 32의 이미지 행렬 등 6개의 batch image를 이용해 다양한 형태의 image tensor를 만들 수 있습니다.
stack_image_32 = rearrange(imgs, '(b1 b2) h w c -> (b2 h) (b1 w) c',b1=2)
plt.imshow(stack_image_32,cmap = 'gray')
print(imgs.shape,stack_image_32.shape) #전: torch.Size([6, 28, 28, 1]) 후: torch.Size([84, 56, 1])
위의 코드를 보면 batch(6)을 b1과 b2로 나누고, b1=2로 지정해주면서 자연스럽게 b2=3으로 지정되어 (2*h, 3*w) size의 결합된 image tensor가 생성되었습니다.
rearrange 함수는 batch만 재배열할 수 있는 것이 아닙니다. 너비나 높이를 2배로 늘리거나 줄이는 것도 가능합니다.
stack_image_61 = rearrange(imgs,'b h (w1 w2) c -> (h w2) (b w1) c', w2=2)
plt.imshow(stack_image_61,cmap = 'gray')
print(imgs.shape,stack_image_61.shape) #전: torch.Size([6, 28, 28, 1]) 후: torch.Size([56, 84, 1])
width를 기준으로 너비는 2(w1)배 줄이고, 높이는 2(w2)배 늘림으로써 tensor의 성분수는 그대로 유지하면서 tensor의 모양을 위의 그림과 같이 바꿨습니다. 이때 괄호() 안의 순서나 w1,w2의 순서가 바뀌면 의도와 다른 이미지가 생성되니 주의해야합니다.
rearrange가 어떤 방식으로 tensor의 성분을 움직이는지는 아래의 코드를 통해 이미지를 생성해서 확인할 수 있습니다.
stack_image_6 = rearrange(imgs,'(b1 b2) h w c -> h (b1 b2 w) c', b1=2)
plt.imshow(stack_image_6,cmap = 'gray')
print(stack_image_6.shape) #torch.Size([28, 168, 1])
stack_image_62 = rearrange(imgs,'(b1 b2) h w c -> h (b2 b1 w) c', b1=2)
plt.imshow(stack_image_23,cmap = 'gray')
print(stack_image_23.shape) #torch.Size([28, 168, 1])
# b1 b2 -> (2,3)
두 이미지 tensor의 shape은 동일합니다. 그러나 image의 배치 순서가 다릅니다. 그 이유는 무엇일까요? 그 이유는 두 image tensor의 output pattern이 다르기 때문입니다. 첫번째 image는 width의 배열이 (b1 b2) 원래 batch의 배열과 동일하고, 두번째 image는 배열이 (b1 b2)가 아닌 transposed된 형태의 (b2 b1)이기 때문입니다. einops를 활용할 때는 이 점을 유의하여 사용해야 하겠습니다.
2. reduce 함수 설명
reduce함수는 차원을 축소하거나, tensor의 크기를 축소하는데 활용됩니다.
먼저 차원을 축소하는 예시입니다. pytorch에서는 차원을 축소할 때 squeeze함수를 통해서 성분이 한개인 차원만 축소가 가능했는데, reduce함수는 그와 관계없이 차원을 축소하는 데, 지정 함수(mean, max, min 등)을 통해 다른 차원들에 aggregation 연산을 수행해 정보가 유지될 수 있도록 합니다.
reduce_image = reduce(imgs,'b h w c -> h w c','mean')
plt.imshow(reduce_image,cmap='gray') #축소된 차원의 성분 6개의 값이 다른 차원에 평균되서 들어감
print(imgs.shape,reduce_image.shape) #전: torch.Size([6, 28, 28, 1]) 후: torch.Size([28, 28, 1])
batch 차원을 축소하고 mean함수로 다른 차원들에 aggregation을 하였더니 위의 그림과 같이 이미지가 흐릿해진 상태로 결합되었습니다. 만일 max함수를 쓰면 어떻게 될까요?
max함수는 위와 같이 모든 batch의 max 값을 유지하다보니 batch 내의 모든 image가 그대로 쌓여 결합된 것 처럼 남아 있습니다.
이렇게 reduce함수는 차원을 축소할 수도 있지만 성분의 수 즉, tensor의 크기를 줄일 수도 있습니다.
reduce_image = reduce(imgs, 'b (h r1) (w r2) c -> h (b w) c','mean',r1=2, r2=2)
plt.imshow(reduce_image,cmap='gray') #image 사이즈를 줄임
print(imgs.shape,reduce_image.shape) #전: torch.Size([6, 28, 28, 1]) 후: torch.Size([14, 84, 1])
위에서 rearrange를 이용해서 batch가 width를 기준으로 결합했을 때는 (28*128*1)의 사이즈를 가졌었는데, reduce함수를 통해 image의 size를 가로, 세로 모두 2배 줄였을 때 (14*84*1)의 사이즈가 된 것을 알 수 있습니다.
3. repeat 함수 설명
repeat함수는 pytorch의 repeat이나 expand와 같이 tensor의 크기 확장에 활용됩니다. 이 역시 사용자가 지정하는 pattern에 따라서 여러 방향으로 확장 될 수 있는데, 아래의 예시들을 보며 설명하겠습니다.
repeat_image = repeat(imgs[0], 'h w c -> h (5 w) c') #width를 기준으로 5배 repeat
plt.imshow(repeat_image,cmap='gray')
print(imgs.shape,repeat_image.shape) #전: torch.Size([28, 28, 1]) 후: torch.Size([28, 140, 1])
(28,28,1)의 이미지를 가로 방향으로 5배 반복하여 (28,140,1) 크기의 반복된 이미지 tensor를 생성하였습니다.
당연히 한 방향만이 아닌 가로, 세로 두 방향 모두 반복도 가능합니다.
repeat_image = repeat(imgs[0], 'h w c -> (2 h) (3 w) c') #세로 방향으로 2번 repeat, 가로 방향으로 3번 repeat
plt.imshow(repeat_image,cmap='gray')
print(imgs[0].shape,repeat_image.shape) #전: torch.Size([28, 28, 1]) 후: torch.Size([56, 84, 1])
이렇게만 보면 pytorch의 repeat 연산과 비슷해 보이는데 einops의 repeat함수는 역시 pattern을 일부 변형하여 새로운 형태의 image tensor로의 변형이 가능합니다.
repeat_image = repeat(imgs[0], 'h w c -> (2 h) (w 3) c')
plt.imshow(repeat_image,cmap='gray')
print(imgs.shape,repeat_image.shape)
이 이미지는 지금까지 봤던 동일 크기 이미지의 단순 반복이 아니라 이미지의 사이즈를 변형하고 이를 반복하였습니다. 위의 코드는 width가 반복될 때 같은 width를 3번 반복하는 게 아니라 3을 width만큼 반복하는 pattern을 넣었습니다. 이는 123 123 123 과 111 222 333이 1,2,3을 각각 3번 반복했지만 형태가 다르듯 사용자가 주문하는 repeat 형태에 따라 tensor가 변경될 수 있는 유연성을 갖고 있다고 할 수 있겠습니다.
마무리
오늘은 Einops 공식 페이지(https://einops.rocks/1-einops-basics/)의 튜토리얼을 참고하여 Einops의 사용법에 대해서 알아보았습니다. Tensor 연산의 직관성과 유연성있게 수행할 수 있는 패키지 이기에 유용하게 활용할 수 있을 것 같습니다. 앞으로 Tensor연산을 최대한 다양하게 해보면서 익숙해지면 좋을 것 같다고 생각합니다.
