다음은 모델의 학습 가능한 파라미터 수를 model.parameters()를 이용하여 확인해 보겠습니다. 전반적인 코드의 흐름과는 관련이 없는 부분이지만 정보 확인을 위한 코드입니다. 참고로 해당 정보는 앞에서 실행했던 summary로도 확인이 가능합니다.

코드 6-13 학습 가능한 파라미터 수 확인

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f'The model has {count_parameters(model):,} trainable parameters')

다음 실행 결과를 보니 46,038,242개의 파라미터를 학습해야 함을 확인할 수 있습니다. 꽤 많은 파라미터를 학습해야 하기 때문에 추후 데이터를 모델에 적용했을 때 긴 시간이 걸릴 것을 예상할 수 있습니다. 그런 이유로 캐글에서 내려받은 이미지 전체를 사용하는 것이 아닌 일부 이미지만 예제에서 사용합니다.

The model has 46,038,242 trainable parameters

모델의 네트워크 구성이 완료되었기 때문에 이제 옵티마이저와 손실 함수를 정의합니다.

코드 6-14 옵티마이저와 손실 함수 정의

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) ------ ①
criterion = nn.CrossEntropyLoss()