原文來自知乎，現摘錄與此

https://zhuanlan.zhihu.com/p/67964081

首先這是一段mnist資料集的基本程式碼。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Net()
optimizer = optim.SGD(model.parameters(), lr=1e-6, momentum=0.5)
train_loader = []

model.train()
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
 

1.初始化(nn.Module.__init__())
init中主要初始化了很多引數，比如buffers，hook等等。根據Net類的程式碼，它會依次初始化各個層。nn.Module.__init__()

 def __init__(self):
        self._backend = thnn_backend
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._backward_hooks = OrderedDict()
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()
        self.training = True
 

關於hook: PyTorch 設計了兩種 hook：register_forward_hook 和 register_backward_hook，
分別用來獲取正/反向傳播時，中間層模組輸入和輸出的 feature/gradient，大大降低了獲取模型內部資訊流的難度。

register_forward_hook的作用是獲取前向傳播過程中，各個網路模組的輸入和輸出。對於模組 module，其使用方式為：module.register_forward_hook(hook_fn) 。
register_backward_hook 的作用是獲取神經網路反向傳播過程中，各個模組輸入端和輸出端的梯度值。

2.前向傳播真正的計算入口點(nn.Module.__call__())

 def __call__(self, *input, **kwargs):
        for hook in self._forward_pre_hooks.values():
            hook(self, input)
        if torch._C._get_tracing_state():
            result = self._slow_forward(*input, **kwargs)
        else:
            result = self.forward(*input, **kwargs)
        for hook in self._forward_hooks.values():  #在執行到這一條語句之前，計算實際上是沒有發生的
            hook_result = hook(self, input, result)
            if hook_result is not None:
                raise RuntimeError(
                    "forward hooks should never return any values, but '{}'"
                    "didn't return None".format(hook))
        if len(self._backward_hooks) > 0:
            var = result
            while not isinstance(var, torch.Tensor):
                if isinstance(var, dict):
                    var = next((v for v in var.values() if isinstance(v, torch.Tensor)))
                else:
                    var = var[0]
            grad_fn = var.grad_fn
            if grad_fn is not None:
                for hook in self._backward_hooks.values():
                    wrapper = functools.partial(hook, self)
                    functools.update_wrapper(wrapper, hook)
                    grad_fn.register_hook(wrapper)
        return result

for hook in self._forward_hooks.values(): 在執行到這一條語句之前，計算實際上是沒有發生的。這一行會在執行forward之前進行，處理預設的hook。

if torch._C._get_tracing_state():
    result = self._slow_forward(*input, **kwargs)
else:
    result = self.forward(*input, **kwargs)

這個地方實現了在不寫C的情況下直接執行forward,有一些自定義操作沒有C，會直接呼叫python的版本。
這一步開始，呼叫了forward方法，首先會呼叫Net類的forward方法，然後會以此呼叫Conv2d的__call__()方法等。
當呼叫Conv2d()的forward方法，其forward方法寫在了torch._C下：

ensor Conv2dImpl::forward(const Tensor& input) {
  if (options.transposed_) {
    return torch::conv_transpose2d(
        input,
        weight,
        bias,
        options.stride_,
        options.padding_,
        options.output_padding_,
        options.groups_,
        options.dilation_);
  }
  return torch::conv2d(
      input,
      weight,
      bias,
      options.stride_,
      options.padding_,
      options.dilation_,
      options.groups_);

然而這依然是一個wrapper，這部分邏輯程式碼最終由aten/c10定義 https://zhuanlan.zhihu.com/p/55966063
最終計算在：

CPU: legacy::cpu::_thnn_conv2d_forward
CUDA: legacy::cuda::_thnn_conv2d_forward

到這裡，一個卷積層的forward操作就結束了，其他層的forward同理。
Conv2d的forward方法執行完成之後接著進行forward_hook和backward_hook的步驟，與之前的forward_pre_hook相似。
到這裡，Conv2d的__call__()方法執行完畢，接下來執行relu之類的邏輯，直到return。

呼叫棧返回Net的forward的返回值，得到loss。

到這裡，前向傳播完成。

3.反向傳播(loss.backward())
loss.backward() 只執行一次，計算完成所有的梯度。

首先，所有的requires_grad為True的張量都會被記錄並被新增進Engine::ready_queue_by_index中，這些tensor都會被以FunctionTask的結構體記錄在ReadyQueue中。

首先在前向傳播的時候，所有requiresgrad==True的物件都會被新增進一個容器中，然後在backward執行之前，首先啟動一個處理引擎，在做了初始化和讀取相關的記錄(包括之前的哪個容器)後呼叫了run_backward方法，然後統一計算出梯度，並返回loss的梯度。