[feature proposal] Composing transformations with __add__ magic method

TLDR: transform = Resize((299, 299)) + ToTensor()

Motivation

Consider the common use case of building an overall-common-but-differing-at-one-point transformation pipeline for train and test sets, for instance applying augmentation transforms only to train set:

train_transform = transforms.Compose([
    Resize((299, 299), interpolation=1),
    RandomCrop(299, padding=4),
    RandomHorizontalFlip(),
    RandomGrayscale(),
    ToTensor(), 
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    Resize((299, 299), interpolation=1), 
    ToTensor(), 
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

A way of expressing of train_transform builds upon test_transform would be handy. What you could do now is to recursively use transforms.Compose (i.e. Compose([Compose(...resizing...), Compose(...augmentation...], Normalize(...)])), but I don’t find this particularly clean (and intuitively, final transforms objects should be sequential objects with only one level of depth).

Proposed behavior

common_transforms = ToTensor() + Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
augmentation = RandomCrop(299, padding=4) + RandomHorizontalFlip() + RandomGrayscale()
train_transforms =  Resize((299, 299)) + augmentation + common_transforms
test_transforms = Resize((299, 299)) + common_transforms

Note how this mimics the behavior of Python lists (and other iterables):

>>> [1, 2] + [3, 4]
[1, 2, 3, 4]

Implementation sketch

class Transform(object):
    # Right now all transforms directly inherit from object. 
    # Let Transform be a superclass for all transformations

    def __call__(self, pic):
        raise NotImplementedError('Each subclass should implement this method')

    def __repr__(self):
        raise NotImplementedError('Each subclass should implement this method')

    def __add__(self, other):
        if not isinstance(other, Transform):
            raise TypeError('Only transformations can be added')
        if isinstance(self, Compose) and isinstance(other, Compose):
            return Compose(self.transforms + other.transforms)
        if not isinstance(self, Compose) and isinstance(other, Compose):
            other.transforms = self + other.transforms
            return other
        if isinstance(self, Compose) and not isinstance(other, Compose):
            self.transforms = self.transforms + other
            return self
        if not isinstance(self, Compose) and not isinstance(other, Compose):
            return Compose([self, other])

Comments are most welcome!

[0phoff]

I wanted similar behaviour and ended up creating my own Compose class for this:

class Compose(list):
  def __call__(self, data):
    for t in self:
      data = t(data)
    return data

My Compose is basically a callable list… Not sure if it has any drawbacks (it probably has), but this has the advantage of being familiar to most people, as it has all the methods lists have!
A few things you can do with it:

preprocess = Compose([...])

# Concatenate composes
preprocess += Compose([...])

# Concatenate Compose with a list
preprocess += [...]
preprocess.extend([...])
preprocess.append(...)

# Access/change individual elements
step1 = preprocess[0]
preprocess[1] = ...

As you can see, it doesn’t allow you to add a function to a compose like @tomekkorbak wanted and also will not magically turn a sum of 2 transforms in a compose, but I do believe this is quite a clean implementation. All that is needed is to wrap your functions in a list and your done!

I really prefer my Compose, but would like to use the ‘standard PyTorch way’ as much as possible, so it would be cool to implement this into torchvision.

EDIT
You might want to also add a __coerce__(self, other) method to my Compose, if you want to be able to perform operations like [...] + preprocess. I don’t really need that, so I didn’t implement it as it would make everything slightly more complicated…

Example implementation (untested):

def __coerce__(self, other):
  if isinstance(other, Iterable):
    return self, Compose(other)
  else:
    return None

[fmassa]

While I think such a functionality might be somewhat useful (but a bit less efficient than an iterator on a list), I didn’t quite get the drawbacks that you pointed in your message. You could still have something like

common = T.Compose([T.ToTensor(),
                    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

augmentation = T.Compose([T.RandomCrop(299, padding=4),
                          T.RandomHorizontalFlip(),
                          T.RandomGrayscale()])

train_transforms =  T.Compose([T.Resize((299, 299)), augmentation, common_transforms])
test_transforms = T.Compose([T.Resize((299, 299)), common_transforms])

which is still fairly ok, and a bit more efficient than having a long concatenation of composes (that will happen due to the + always creating a new T.Compose).

But I’m not against such proposal. @alykhantejani what do you think?