TransPose: Keypoint Localization via Transformer

이 논문은 pose estimation 분야에 트랜스포머를 적용해 2021년 iccv 에 게재되었는데 cnn 대신 트랜스포머를 적용해 파라미터 수를 줄이고 빠른 속도로 cnn 기반 모델과 비슷한 성능을 달성했습니다. 기존 cnn 기반 모델은 네트워크 구조가 깊고 비선형적이어서 feature 해석이 불가능하고, 특히 cnn 모델들은 localization 보다는 이미지 classification에 초점을 맞추고 있습니다.

따라서 트랜스포머를 사용하면 얻을 수 있는 이점은 두 가지가 있는데, 이미지의 공간 의존성을 명시적으로 포착할 수 있고 컨볼루션의 receptive field 로는 여러 layer를 쌓아 글로벌 의존성을 포착해야 하는 것에 반해 한 개의 어텐션 layer로 모든 pairwise location에 대해 글로벌 의존성을 포착할 수 있습니다. 오른쪽 이미지의 빨간 박스를 보시면 가려진 오른쪽 발목 키포인트 위치를 예측하기 위해 의존하고 있는 영역들을 알 수 있는데, 왼쪽 발목, 허벅지, 오른쪽 다리에 의존하고 있는 모습을 볼 수 있습니다. 이 어텐션 맵은 3번째 어텐션 layer에서 뽑은 것입니다. 

Overall Architecture

이 논문에서 사용한 방법으로는 먼저 이미지 feature를 뽑아 트랜스포머에 넣어주게 되는데 이미지 feature를 뽑는 모델로는 hrnet(TransPose-H-S)과 resnet(TransPose-R)을 사용했고, 인코딩 시 나오는 히트맵을 사용해 activation maximize position을 키포인트로 예측하기 때문에 트랜스포머 구조 중 인코더만 사용합니다.

앞서 설명드린 대로 이미지에서 백본 네트워크를 통해 피처맵을 뽑은 후 패치로 나누고(여기서는 1 pixel = 1 patch) flatten 하여 트랜스포머에 인코딩 레이어에 넣어주게 됩니다. 뽑아준 image feature에 포지션 인코딩을 더해주고, weight 행렬 3개를 사용해 쿼리와 키, value로 projection 합니다. 쿼리와 키는 이미지 피쳐에 포지션 임베딩을 더해준것, value는 이미지 피처를 그대로 사용하고 있습니다. 이후 multi head attention을 사용해 feature를 추출하고 (head는 1개 사용) 쿼리와 키를 행렬곱해 어텐션 스코어, 즉 모든 키와의 유사성을 계산해 weight vecter 를 구합니다.

#Position-wise FFN 

src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))

특히 마지막 어텐션 layer에서 모든 이미지 위치에서의 컨트리뷰션을 모아 예측된 키포인트 히트맵에서 max activation을 만드는 역할을 합니다. 이후의 Fead forward network 와 head는 위치별로 컨트리뷰션의 상대적 비율을 변경하지 않고 선형변환을 해주는 것을 의미합니다.

# head

self.final_layer = nn.Conv2d(
	in_channels=d_model,
	out_channels=cfg['MODEL']['NUM_JOINTS'],
	kernel_size=extra['FINAL_CONV_KERNEL'],
	stride=1,
	padding=1 if extra['FINAL_CONV_KERNEL'] == 3 else 0
)

# forward

x = self.final_layer(x)

head는 위 코드로 나타낼 수 있는데, 트랜스포머의 output이 채널을 키포인트 수로 줄여주는 1x1 conv layer입니다. 결과로 나온 키포인트 히트맵의 activation의 최대 위치는 키포인트 위치를 의미합니다.

Dependencies

$$J = {j|A_{i,j},j (θ^∗,I) ≥ δ}, $$

이미지 I의 특정 키포인트의 위치가 i일 때, element j가 δ 보다 더 높은 attention 스코어를 갖는데 이는 element j가 i 위치 예측에 컨트리뷰션을 많이 준다는 의미입니다.

Dependency area

$A_{i,:}$ Ai: 는 i위치가 의존하고 있는 영역들을 나타내줍니다. 이 영역은 첫 그림 오른쪽의 map들로 나타납니다.

Affected area

$𝐀_{:,𝒋}$ A:j는 j 가 영향을 주는 영역들을 나타내줍니다. 

결과는 다음과 같습니다. 4배 정도 적게 파라미터를 사용하고도 cnn기반 모델들과 비슷한 성능을 내었다고 합니다.

논문: TransPose

코드: https://github.com/yangsenius/TransPose

GitHub - yangsenius/TransPose: PyTorch Implementation for "TransPose: Keypoint localization via Transformer", ICCV 2021.