모델을 훈련시키기 전에 손실 함수와 옵티마이저에 대해 정의해야 합니다. 이번 예제는 데이터를 분류해야 하는 것으로 크로스 엔트로피(cross entropy) 손실 함수를 사용합니다. 또한, 옵티마이저로는 아담(Adam)을 사용합니다.

코드 2-12 모델의 파라미터 정의

loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

파이토치는 GPU에 최적화된 딥러닝 프레임워크입니다. 하지만 GPU가 없다면 CPU를 사용할 수 있도록 지정해 주어야 합니다. 다음은 GPU가 있다면 GPU를 사용하고, 없다면 CPU를 사용하도록 하는 코드입니다.

코드 2-13 CPU/GPU 사용 지정

if torch.cuda.is_available():
    device = torch.device('cuda') ------ GPU가 있다면 GPU를 사용
else:
    device = torch.device('cpu') ------ GPU가 없다면 CPU를 사용