[이거 어떡하지] Expected target size (, ), got torch.Size([, ]) (nn.CrossEntropyLoss 에서)

증상:

ValueError: Expected target size(256, 84), got torch.Size([256, 64])

image classification이나 language model을 학습시키는 과정에서 nn.CrossEntropyLoss 를 사용하실 때 위와 같은 에러가 발생할 때가 있습니다. (특히 language model처럼 batch 한 샘플마다 sequence별로 classification을 해야 하는 경우)

 

sequence sample별로 output으로 num_classes만큼의 dimension이 뽑히게 냈는데 dimension이 안맞으면 에러가 발생할 수 있습니다.

저렇게 output이 나오도록 model을 구성하고 loss = loss_function(y_hat, y) 하면 loss 계산이 되지 않을까? - 어림도 없음.

language model의 경우에는 입력으로

[batch_size, len_input]

의 shape를 갖는 tensor가 들어갑니다.

 

model output으로는

[batch_size, len_output, num_classes], 혹은 [batch_size, num_classes, len_output]

의 shape를 갖는 tensor가 나오도록 구성하셨을 겁니다.

 

위 사진의 경우는

batch_size == 256, len_output == 64, num_classes == 84

로 설정되어 있습니다.

 

해결 방법:

상단 포스트 사진의 경우에는 y_hat.permute(0, 2, 1)로 [batch_size, num_classes, len_output]로 맞춰서 해결할 수 있습니다.

심화: class 수를 2번째 dimension( y_hat[1] )에다가 넣고 나머지 dimension이 model output과 true y가 같도록 한다.

(model output) y_hat.shape = [B, num_classes, dim1, dim2, dim3, ...]
(true y) y.shape = [B, dim1, dim2, dim3, ...]

예를 들어 y.shape가 [64, 8, 32, 16, 128]이면 model output은 [64, num_classes, 8, 32, 16, 128]로 맞춰야합니다.