How to understand PyTorch's source code

cakeeatingpolarbear · September 19, 2017, 9:16pm

Hi all,
I was wondering where (or how) I should approach try to understand PyTorch’s source code? I’m particularly interested in understanding the implementation of the autograd module in PyTorch. I’ve read through all of the docs in regards to autograd and Variables, I also have a good understanding of how the back propagation equations are derived. Though they didn’t help me much in understanding how the modules are implemented. As a result I’ve just been using a debugger to try to step through a toy example, though even then I get lost in the complexity of it all. Am I going about this the wrong way? Or if anyone has attempted this would they care so share how they approached it?

smth · September 20, 2017, 4:27am

These two posts will help get through the code structure and build system:

http://pytorch.org/2017/05/11/Internals.html

http://pytorch.org/2017/06/27/Internals2.html

It doesn’t touch the autograd system though.
The autograd system is moved into C now, and is multi-threaded, so stepping through the python debugger is probably a bit pointless.

Here’s a pointer to very old source code, where all the logic was in Python, in the early days of PyTorch. The logic is largely the same still, but moved into C.

github.com

pytorch/pytorch/blob/v0.1.1/torch/autograd/engine.py

from collections import Counter

class ExecutionEngine(object):
    def __init__(self):
        pass

    def _compute_dependencies(self, function):
        dependencies = {}
        seen = {function}
        queue = [function]
        while len(queue) > 0:
            fn = queue.pop()
            for prev_fn, arg_id in fn.previous_functions:
                if prev_fn not in dependencies:
                    dependencies[prev_fn] = [Counter() for _ in prev_fn.output_ids]
                output_idx = prev_fn.output_ids[arg_id]
                dependencies[prev_fn][output_idx][fn] += 1
                if prev_fn not in seen:
                    queue.append(prev_fn)
                    seen.add(prev_fn)

This file has been truncated. show original

github.com

pytorch/pytorch/blob/v0.1.1/torch/autograd/function.py

from collections import OrderedDict
from .variable import Variable

class Function(object):

    def __init__(self):
        self.previous_functions = None
        self.output_ids = None
        self.needs_input_grad = None
        self.backward_hooks = OrderedDict()

    def __call__(self, *input):
        return self._do_forward(*input)

    def _do_forward(self, *input):
        unpacked_input = tuple(arg.data for arg in input)
        raw_output = self.forward(*unpacked_input)
        if not isinstance(raw_output, tuple):
            raw_output = (raw_output,)
        self.needs_input_grad = tuple(arg.creator.requires_grad for arg in input)

This file has been truncated. show original

cakeeatingpolarbear · September 20, 2017, 9:45am

Thanks I’ll give it a read!

christianperone · May 4, 2018, 1:49pm

I also did an article about it here.