The tensorflow implementation from the paper Attention Augmented Convolutional Networks. Will add a Pytorch implementation of the same soon but currently the torch.einsum is very slow, even on GPU it is slower than Numpy.
To use the layer:
from layer import augmented_conv2d
The paper has shows promising results, especially below:
