pytorch 여러 loss function의 결합

YOLO와 같이 logit에 대해 분할하여 여러 loss를 적용시키는 경우 pytorch에서 autograd가 정상 작동하는지 확인해보았다.

결론만 말하자면 logit을 인덱스 슬라이싱 할 경우 computational graph에서 인덱스 슬라이싱한 부분이 추적이 되어 logit의 일부만 떼어서 loss를 계산하더라도 적상적으로 역전파가 되는 것을 확인할 수 있다.

logit = model(x)

loss = criterion(logit[:5], y[:5])

와 같은 코드라면 model이 logit[:5]까지 영향을 미치는 부분들만 autograd된다.