[pytorch] 학습을 위해 train 함수에 꼭 필요한 내용들

def train(model, train_loader, optimizer):

    model.train()

    for batch_idx, (data, target) in enumerate(train_loader):

        data, target = data.to(DEVICE), target.to(DEVICE)

        optimizer.zero_grad()

        output = model(data)

        loss = F.cross_entropy(output, target)

        loss.backward()

        optimizer.step()

---

model.train() 

≫ 모델을 학습 모드로 변환 / 평가 모드는 model.eval() 로 할 수 있다.

 

data, target = data.to(DEVICE), target.to(DEVICE)

≫  각 data 와 target 을 앞서 설정한 DEVICE(GPU 혹은 CPU) 에 보내는 것

 

optimizer.zero_grad()

≫ 반복 때마다 기울기를 새로 계산하므로, 이 함수로 초기화

 

output = model(data)

≫ data를 모델에 넣어서 가설(hypothesis)를 획득

 

loss = F.cross_entropy(output, target)

≫ 가설과 groud truth를 비교하여 loss 계산

 

loss.backward()

≫ loss 를 역전파 알고리즘으로 계산

 

optimizer.step()

≫  계산한 기울기를 앞서 정의한 알고리즘에 맞추어 가중치를 수정