본문 바로가기

연구/pytorch

[tensor] tensor.size() 에 관하여

반응형

tensor.size()했을 때 나오는 각 숫자들의 의미에 대해서 알아보자.

 

이미지 데이터를 batch단위로 묶어서 학습한다고 생각해보자.

net에 들어가기 전에 tensor의 형태를 출력하면 다음과 같이 나온다.

torch.Size([128, 1, 28, 28])

 

이 때 각각의 숫자가 의미하는게 무엇인지 살펴보자.

128 : mini-batch size(데이터 한 묶음당 들어있는 데이터쌍(x-y) 개수)

1 : channel size

28 : img size

28 : img size

 

이 데이터묶음을 net에 넣고 연산하여 Convolution layer를 모두 거치고 나오면 다음과 같은 형태로 바뀐다.

torch.Size([128, 64, 4, 4])

당연히 batch size 는 변하지 않는다. encoding과정을 거친만큼 img size는 축소했으며 그 만큼 channel size는 증가했다.

이를 FC에 넣기 위해 flatten 하면 다음과 같이 될 것이다.

 

torch.Size([128, 1024])

마찬가지로 batch size 는 변하면 안되고 channel * img size * img size 길이의 tensor가 생긴다.

 


참고로 flatten layer의 코드는 다음과 같이 짜면 된다.

 

 

flatten layer의 사용법은 아래와 같이 쓰면 된다.

 

반응형