MLPs(GMLP) වෙත අවධානය යොමු කරන්න

21from typing import Optional
22
23import torch
24from torch import nn

27class GMLPBlock(nn.Module):

48    def __init__(self, d_model: int, d_ffn: int, seq_len: int):

54        super().__init__()

56        self.norm = nn.LayerNorm([d_model])

58        self.activation = nn.GELU()

60        self.proj1 = nn.Linear(d_model, d_ffn)

62        self.sgu = SpacialGatingUnit(d_ffn, seq_len)

64        self.proj2 = nn.Linear(d_ffn // 2, d_model)

68        self.size = d_model

70    def forward(self, *, x: torch.Tensor, mask: Optional[torch.Tensor] = None):

77        shortcut = x

79        x = self.norm(x)

81        z = self.activation(self.proj1(x))

83        z = self.sgu(z, mask)

85        z = self.proj2(z)

88        return z + shortcut

91class SpacialGatingUnit(nn.Module):

101    def __init__(self, d_z: int, seq_len: int):

106        super().__init__()

108        self.norm = nn.LayerNorm([d_z // 2])

113        self.weight = nn.Parameter(torch.zeros(seq_len, seq_len).uniform_(-0.01, 0.01), requires_grad=True)

117        self.bias = nn.Parameter(torch.ones(seq_len), requires_grad=True)

119    def forward(self, z: torch.Tensor, mask: Optional[torch.Tensor] = None):

128        seq_len = z.shape[0]

130        z1, z2 = torch.chunk(z, 2, dim=-1)

133        if mask is not None:

137            assert mask.shape[0] == 1 or mask.shape[0] == seq_len
138            assert mask.shape[1] == seq_len

140            assert mask.shape[2] == 1

142            mask = mask[:, :, 0]

145        z2 = self.norm(z2)

147        weight = self.weight[:seq_len, :seq_len]

152        if mask is not None:
153            weight = weight * mask

156        z2 = torch.einsum('ij,jbd->ibd', weight, z2) + self.bias[:seq_len, None, None]

159        return z1 * z2