[CUDA] Pytorch c++/cuda extenstion tutorial 6

CUDA Programming Tutorial 6

• Tutorial 5에서 trilinear interpolation을 forward pass로 구현했지만, 아직까진 backward pass는 못합니다.
• So if we want to calculate some loss based on this variable, it is not able to compute the correct gradient to update our features,

backward pass를 implement하여 gradient가 back으로 passed 되도록 cuda code를 작성해보겠습니다.

cuda backward implementation에서 가장 중요한 2가지 스텝은 다음과 같습니다.

• First, we need to compute all outputs’ partial derivatives w.r.t. all trainable inputs.

사용할 Loss를 $L$이라고 합시다.

trainable inputs인 $f_1$ ~ $f_4$의 미소 변화량과 Loss의 관계에서 미분인 $\frac{dL}{df_1}$, $\frac{dL}{df_2}$, $\frac{dL}{df_3}$, $\frac{dL}{df_4}$을 구합니다.

Since we already computed all partial derivaties of “all outputs” w.r.t. “all trainable inputs”, we certainly can represent all these $dL/d_{something}$.

This is why we did the $df/d_{inputs}$ in the first place.

Loss에 대해서 모든 inputs에 대한 $dL$, 즉, $\frac{dL}{df_1}$, $\frac{dL}{df_2}$, $\frac{dL}{df_3}$, $\frac{dL}{df_4}$을 정의하였으니까, backward pass를 코딩할 수 있습니다.

For $\frac{dL}{d_{feats}}$, the original feats has shape (N,8,F), $\frac{dL}{d_{feats}}$ has the same shape (a tensor always has the same shape as its gradient).

tensor와 tensor의 gradient는 항상 같은 shape을 가짐을 기억합시다.

backward pass 미분계산 요약

• To calculate partial derivates, we imagine a virtual Loss $L$

• then there will be the partial derivates(s) of this $L$ w.r.t the output(s), in our case $\frac{dL}{f}$

  

• We need to establish, from $\frac{dL}{df}$, the formulae of the partial derivates of $L$ w.r.t the inputs, so $\frac{dL}{df_1}$ ~ $\frac{dL}{df_4}$.

  • This can be achieved using chain rule and the partial derivates of the output(s) w.r.t the input(s).

    

• Once we have the forumula, last thing we need to do is to put the value one by one into our tensor.

  

Backward pass code

• forward code를 copy & paste하고, fill in the partial derviates, change the input output tensors and their shape.

  

• 그리고 fowrad function에서 해준 것처럼 .cpp, .h 파일들에 backward function을 선언해줍니다.

• PYBIND11로 forward pass, backward pass가 python과 연동될 수 있도록 이름을 정의해줍니다.

  

  

• 마지막으로 setup.py의 current working dir에서 pip install .로 library를 build 해주는 것을 잊지 맙시다.

Pytorch에서 torch.autograd.Function을 call하여 fw, bw functions을 wrap 해주어야 합니다.

• ctx: the abbreviation of context, it is used to store some variables that we need for the backward pass.
• ctx is an argument that is always required and should be the first argument.

forward pass

• intuitively we only need to return the output
• but since in the backward partial derivatives, we need both the points and feats variables (c.f. for formulae).

• in the formulae there are u,v which can only be calculated from points, so we need to save them.

  

  다음 코드로 fw에서 partial derivatives에 관련된 inputs (feats, points)를 저장합니다.

  ctx.save_for_backward(feats, points)
  

backward pass

• the following argument(s) is what we imagined earlier, $\frac{dL}{d_{output(s)}}$.

• 즉, backprop에 사용되는 Loss term $L$에 output feat_interp이 주는 미소 변화량 $\frac{dL}{dfeat \ interp}$를 backward pass에 넣습니다.

  

• If you have multiple outputs, your backward function will have as many arguments as your outputs, all with the form $\frac{dL}{d_{output}}.
• Loss term $L$과 관련된 outputs이 여러 개면, backward pass에 $L$에 대한 outputs의 미소 변화량들을 모두 넣어줍니다.

backward pass에서는 먼저 forward pass에서 저장된 tensors들을 가져옵니다.

• backward pass에 inputs으로 작성한 3개의 inputs을 넣어주어 $\frac{dL}{d_{feats}}$인 the partial derivatives of the loss w.r.t the inputs을 얻습니다.

• backward pass에서 return하는 value는 forward pass의 inputs의 수와 같습니다.
• feats, points가 input인데, backward pass에서 return 되는 value는 partial derivative of the loss w.r.t that input 입니다.
• Because the first input ins feats, the first return value is $\frac{dL}{d_{feats}}$.

• Why is the second None? That’s because we said that in the problem definition, we assume that the points are fixed, so never changed. So the None here means we won’t have any gradient for the second input, points.

  

이제 backward pass 까지 required_grad = True가 되는지 확인해봅시다.

• Trilinear_interpolation_cuda.apply(feats2, points)는 forward pass와 backward pass를 torch.autograd.Function로 상속받은 class로 구현한 함수 입니다.
  • 이 class는 pytorch에서 .apply()로 사용합니다.

• trilinear_interpolation_py(feats, points)는 pytorch로 구현한 함수입니다.

• 둘의 값 차이는 없음을 torch.allclose(out_py, out_cuda)로 확인했습니다.
  • out_py는 pytorch로 구현한 함수의 output
  • out_cuda는 cuda로 forward pass, backward pass를 구현한 함수의 output
• 속도에서 forward에는 pytoch나 cuda나 거의 속도 차이가 없지만, backward pass에서 cuda가 pytorch보다 10배 빠름을 확인했습니다.

감사합니다.

Reference

Pytorch+cpp/cuda extension 教學 tutorial 6 - English CC -