(self, dim, num_heads, dropout=0.1, eps=1e-5)
| 14 | class SelfAttention(nn.Module): |
| 15 | |
| 16 | def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5): |
| 17 | assert dim % num_heads == 0 |
| 18 | super().__init__() |
| 19 | self.dim = dim |
| 20 | self.num_heads = num_heads |
| 21 | self.head_dim = dim // num_heads |
| 22 | self.eps = eps |
| 23 | |
| 24 | # layers |
| 25 | self.q = nn.Linear(dim, dim) |
| 26 | self.k = nn.Linear(dim, dim) |
| 27 | self.v = nn.Linear(dim, dim) |
| 28 | self.o = nn.Linear(dim, dim) |
| 29 | self.dropout = nn.Dropout(dropout) |
| 30 | |
| 31 | def forward(self, x, mask): |
| 32 | """ |