[Pytorch Error] Dimension out of range (expected to be in range of [-1, 0], but got 1)

에러 발생

.py 파일로 모듈화하는 과정에서 다음과 같은 에러가 발생했습니다.

IndexError: Dimension out of range (expected to be in range of [-1,0], but got 1)

결론!

결론부터 말하자면, flatten(x,1)로 에러를 해결했습니다.

과정 설명

현재 ResNet18을 사용하여 이미지 다중분류 task를 하고 있습니다. ResNet과 관련하여 model.py는 완성해뒀지만, 추후에 더욱 다양한 실험을 위해 ResNet 이외에도 VGG, EfficientNet 등 다른 모델들도 model.py에 추가하여, 어떤 feature extractor를 입력하는가에 따라 자동으로 모델이 돌아갈 수 있도록 코드를 수정해나가고 있었습니다. 코드를 조금 더 단순화하고 간결하게 정리하는 도중에 다음 에러가 발생하였고, 이를 해결하기 위해 여러 시도를 해보았습니다.

++ 기존 코드에서 오류가 발생한 부분은 다음과 같습니다.

> def forward(self, x):
	x = self.model(x)
    x = self.fc(x.squeeze()) ## nn.linear()// 마지막 layer에서 출력노드를 class 개수로 수정하여 추가함.
    return x

1. squeeze 제거함 -> 오류 발생

1 epoch가 돌아가고 나서 바로 오류가 발생했기 때문에, 모델을 돌렸을 때 출력값을 뽑아서 수정 전의 출력값과 비교해봤지만 크게 다르진 않았습니다.
따라서, 먼저 forward 모듈에서 squeeze()를 제거하고 코드를 돌려봤습니다. 이때, mat1과 mat2의 shape이 다르다며, RuntimeError가 출력되었습니다. 이 방법이 틀렸음을 알 수 있었습니다...

2. view(-1) 추가 -> 오류 발생

두 번째 방법으로는 view를 추가해주는 방법을 사용했습니다. 검색 결과 중에서, 모델을 돌리는 py 파일에서 loss = criterion(outputs, labels)를 loss = criterion(outputs.view(-1), labels)로 수정하면 해결할 수 있다는 글을 보고 바로 시도해보았습니다. 이때에도 오류가 발생했는데, 캡쳐를 해두지 않았.,,,,습...니다.. 어쨌든 이 방법도 제 오류를 해결할 수 없었습니다..!

3. flatten() 추가 -> 오류 해결

마지막으로 시도해 본 방법은 flatten()을 추가해주는 것이었습니다. 다음 stackoverflow을 참고했습니다.
저도 모델을 2개의 부분으로 나눈 후, 출력 layer를 수정하여 사용한 것이기 때문에 이 방법을 시도해봤습니다. 위의 forward()를 다음과 같이 수정했습니다.

> def forward(self, x):
	x = self.model(x)
    x = torch.flatten(x,1) ## self.flatten으로 수정 예정
    x = self.fc(x)
    return x

flatten을 추가하여 코드를 수행했을 때, 코드가 무사히 잘 돌아가고 있는 것을 확인할 수 있었습니다.

이렇게 오류를 해결하고 난 뒤, train.py 코드를 실행해주면, 다음과 같이 출력이 됩니다.
이때 best acc를 갖는 모델을 저장해주고, inference.py에서 그 모델을 불러온 뒤 test 해보면 다음과 같이 오류 없이 잘 수행됨을 알 수 있습니다.

+ 미래의 나에게..

혹시라도 이런 오류가 또 발생한다면.. 해결 방법이 생각이 나지 않는다면.. 너의 velog를 확인하렴....!!