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

10class SpatialDepthWiseConvolution(Module):

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

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

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

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

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)

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))