학습에서 사용될 옵티마이저와 손실 함수를 정의합니다.
코드 6-31 옵티마이저 및 손실 함수 정의
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()
torchsummary를 이용하여 모델의 네트워크를 다시 살펴보겠습니다. 모델의 네트워크를 파라미터와 함께 보여 주고 있기 때문에 내부적으로 어떤 일들이 이루어지는지 예측하면서 살펴볼 수 있습니다.
코드 6-32 모델 네트워크 구조 확인
from torchsummary import summary
summary(model, input_size=(3, 256, 256))
다음은 torchsummary를 이용한 네트워크 구조를 출력한 결과입니다.
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 63, 63] 23,296
ReLU-2 [-1, 64, 63, 63] 0
MaxPool2d-3 [-1, 64, 31, 31] 0
Conv2d-4 [-1, 192, 31, 31] 307,392
ReLU-5 [-1, 192, 31, 31] 0
MaxPool2d-6 [-1, 192, 15, 15] 0
Conv2d-7 [-1, 384, 15, 15] 663,936
ReLU-8 [-1, 384, 15, 15] 0
Conv2d-9 [-1, 256, 15, 15] 884,992
ReLU-10 [-1, 256, 15, 15] 0
Conv2d-11 [-1, 256, 15, 15] 590,080
ReLU-12 [-1, 256, 15, 15] 0
MaxPool2d-13 [-1, 256, 7, 7] 0
AdaptiveAvgPool2d-14 [-1, 256, 6, 6] 0
Dropout-15 [-1, 9216] 0
Linear-16 [-1, 4096] 37,752,832
ReLU-17 [-1, 4096] 0
Dropout-18 [-1, 4096] 0
Linear-19 [-1, 512] 2,097,664
ReLU-20 [-1, 512] 0
Linear-21 [-1, 2] 1,026
================================================================
Total params: 42,321,218
Trainable params: 42,321,218
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.75
Forward/backward pass size (MB): 10.90
Params size (MB): 161.44
Estimated Total Size (MB): 173.10
----------------------------------------------------------------