Pytorch MNIST에 CNN모델 만들기
모두를 위한 딥러닝 - 파이토치 강의 참고
Convolution layer를 활용해서 MNIST이미지를 학습시킬 CNN 모델을 만들어보자.
아래 CNN 구조와 비슷한 모델을 설계해본다.
torch.nn.Sequential을 통해 위의 그림에서 나타내는 각 layer와 Fully Connected layer를 표현할 수 있다.
1 | class CNN(torch.nn.Module): |
Conv2d와MaxPool2d을 지닌 layer를 통과하면서 output의 shape은 게속 변하게 된다.이러한 과정에서 Fully Connected layer를 지날때는 이전 layer에서 나온 output의 shape를 알아야한다.
이때 이전 포스트에서 알아본 공식을 통해 output shape을 계산할 수 있다.
또한, forward를 진행하면서 각각의 layer를 통과한 output의 shape를 출력하면서 확인할 수도 있다.
CNN 모델을 배우기 전에는 선형함수만을 통해 MNIST이미지를 학습했다면 이번 과정에서는 지금까지 배운 Dropout, weight initialization, Convolution layer 등을 사용하는 모델을 만들 수 있었다.
학습과 평가는 기존 MNIST 학습 코드와 유사하고, Full Code 링크를 통해 확인할 수 있다.