efficient.py

#

1import math
2
3import torch
4from torch import nn
5
6from labml_helpers.module import Module
7from labml_nn.transformers import MultiHeadAttention

#

Spatial Depth Wise Convolution

This is actually slower

10class SpatialDepthWiseConvolution(Module):

#

d_k is the number of channels in each head

17    def __init__(self, d_k: int, kernel_size: int = 3):

#

21        super().__init__()
22        self.kernel_size = kernel_size

#

We use PyTorch's Conv1d module. We set the number of groups to be equal to the number of channels so that it does a separate convolution (with different kernels) for each channel. We add padding to both sides and later crop the right most kernel_size - 1 results

27        rng = 1 / math.sqrt(kernel_size)
28        self.kernels = nn.Parameter(torch.zeros((kernel_size, d_k)).uniform_(-rng, rng))

#

x has shape [seq_len, batch_size, heads, d_k]

30    def forward(self, x: torch.Tensor):

#

35        res = x * self.kernels[0].view(1, 1, 1, -1)
36
37        for i in range(1, len(self.kernels)):
38            res[i:] += x[:-i] * self.kernels[i].view(1, 1, 1, -1)
39
40        return res

#

Multi-DConv-Head Attention (MDHA)

We extend our original implementation of Multi-Head Attention and add the spatial depth-wise convolution to query, key and value projections.

43class MultiDConvHeadAttention(MultiHeadAttention):

#

51    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
52        super().__init__(heads, d_model, dropout_prob)

#

Multi-Head Attention will create query, key and value projection modules self.query , self.key , and self.value .

We combine a spatial depth-wise convolution layer to each of them and replace self.query , self.key , and self.value .

59        self.query = nn.Sequential(self.query, SpatialDepthWiseConvolution(self.d_k))
60        self.key = nn.Sequential(self.key, SpatialDepthWiseConvolution(self.d_k))
61        self.value = nn.Sequential(self.value, SpatialDepthWiseConvolution(self.d_k))