Search code examples
pythonpytorchcoding-styleconventions

Argument convention in PyTorch


I am new to PyTorch and while going through the examples, I noticed that sometimes functions have a different convention when accepting arguments. For example transforms.Compose receives a list as its argument:

transform=transforms.Compose([  # Here we pass a list of elements
   transforms.ToTensor(), 
   transforms.Normalize(
      (0.4915, 0.4823, 0.4468), 
      (0.2470, 0.2435, 0.2616)
   )
]))

At the same time, other functions receive the arguments individually (i.e. not in a list). For example torch.nn.Sequential:

torch.nn.Sequential(  # Here we pass individual elements
    torch.nn.Linear(1, 4),
    torch.nn.Tanh(),
    torch.nn.Linear(4, 1)
)

This has been a common typing mistake for me while learning.

I wonder if we are implying something when:

  • the arguments are passed as a list
  • the arguments are passed as individual items

Or is it simply the preference of the contributing author and should be memorized as is?

Update 1: Note that I do not claim that either format is better. I am merely complaining about lack of consistency. Of course (as Ivan stated in his answer) it makes perfect sense to follow one format if there is a good reason for it (e.g. transforms.Normalize). But if there is not, then I would vote for consistency.


Solution

  • This is not a convention, it is a design decision.

    Yes, torch.nn.Sequential (source) receives individual items, whereas torchvision.transforms.Compose (source) receives a single list of items. Those are arbitrary design choices. I believe PyTorch and Torchvision are maintained by different groups of people, which might explain the difference. One could argue it is more coherent to have the inputs passed as a list since it is as a varied length, this is the approach used in more conventional programming languages such as C++ and Java. On the other hand you could argue it is more readable to pass them as a sequence of separate arguments instead, which what languages such as Python.

    In this particular case we would have

    >>> fn1([element_a, element_b, element_c]) # single list
    

    vs

    >>> fn2(element_a, element_b, element_c) # separate args
    

    Which would have an implementation that resembles:

    def fn1(elements):
        pass
    

    vs using the star argument:

    def fn2(*elements):
        pass
    

    However it is not always up to design decision, sometimes the implementation is clear to take. For instance, it would be much preferred to go the list approach when the function has other arguments (whether they are positional or keyword arguments). In this case it makes more sense to implement it as fn1 instead of fn2. Here I'm giving second example with keyword arguments. Look a the difference in interface for the first set of arguments in both scenarios:

    >>> fn1([elemen_a, element_b], option_1=True, option_2=True) # list
    

    vs

    >>> fn2(element_a, element_b, option_1=True, option_2=True) # separate
    

    Which would have a function header which looks something like:

    def fn1(elements, option_1=False, option_2=False)
       pass
    

    While the other would be using a star argument under the hood:

    def fn2(*elements, option_1=False, option_2=False)
       pass
    

    If an argument is positioned after the star argument it essentially forces the user to use it as a keyword argument...

    Mentioning this you can check out the source code for both Compose and Sequential and you will notice how both only expect a list of elements and no additional arguments afterwards. So in this scenario, it might have been preferred to go with Sequential's approach using the star argument... but this is just personal preference!