diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py index 92ac7e4..61a84c1 100644 --- a/fastai_contrib/models.py +++ b/fastai_contrib/models.py @@ -70,6 +70,120 @@ class MultiBatchBiLMModel(BiLMModel): outputs.append(o) return self.concat(raw_outputs), self.concat(outputs) +class BiAttentionPoolingClassifier(nn.Module): + r" [WIP] BiLM Pooling with self attention" + + def __init__(self, layers:Collection[int], drops:Collection[float]): + super().__init__() + mod_layers = [] + activs = [nn.ReLU(inplace=True)] * (len(layers) - 2) + [None] + for n_in,n_out,p,actn in zip(layers[:-1],layers[1:], drops, activs): + mod_layers += bn_drop_lin(n_in, n_out, p=p, actn=actn) + self.self_attn = MultiHeadAttention(n_head=8, d_model=1, d_k=64, d_v=64, dropout=0.1) + self.layers = nn.Sequential(*mod_layers) + + def pool(self, x:Tensor, bs:int, is_max:bool): + "Pool the tensor along the seq_len dimension." + f = F.adaptive_max_pool1d if is_max else F.adaptive_avg_pool1d + return f(x.permute(1,2,0), (1,)).view(bs,-1) + + def forward(self, input:Tuple[Tensor,Tensor])->Tuple[Tensor,Tensor,Tensor]: + raw_outputs, outputs = input + output = outputs[-1] + assert len(output.size()) == 4, 'Expected input dimension 4' + sl, bs, em_sz, passes = output.size() + + f_avgpool = self.pool(output[...,0], bs, False) + f_mxpool = self.pool(output[...,0], bs, True) + b_avgpool = self.pool(output[..., 1], bs, False) + b_mxpool = self.pool(output[..., 1], bs, True) + x = torch.cat([output[-1][..., 0], f_mxpool, f_avgpool, + output[-1][..., 1], b_mxpool, b_avgpool,], 1) + x = x.unsqueeze(-1) + x, _ = self.self_attn(x, x, x) + + x = self.layers(x) + return x, raw_outputs, outputs + +class ScaledDotProductAttention(nn.Module): + r""" + Scaled Dot-Product Attention + based on: https://github.com/jadore801120/attention-is-all-you-need-pytorch + """ + + def __init__(self, temperature, attn_dropout=0.1): + super().__init__() + self.temperature = temperature + self.dropout = nn.Dropout(attn_dropout) + self.softmax = nn.Softmax(dim=2) + + def forward(self, q, k, v): + + attn = torch.bmm(q, k.transpose(1, 2)) + attn = attn / self.temperature + + attn = self.softmax(attn) + attn = self.dropout(attn) + output = torch.bmm(attn, v) + + return output, attn + +class MultiHeadAttention(nn.Module): + r""" + Multi-Head Attention module + based on: https://github.com/jadore801120/attention-is-all-you-need-pytorch + """ + + def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1): + super().__init__() + + self.n_head = n_head + self.d_k = d_k + self.d_v = d_v + + self.w_qs = nn.Linear(d_model, n_head * d_k) + self.w_ks = nn.Linear(d_model, n_head * d_k) + self.w_vs = nn.Linear(d_model, n_head * d_v) + nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k))) + nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k))) + nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v))) + + self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5)) + self.layer_norm = nn.LayerNorm(d_model) + + self.fc = nn.Linear(n_head * d_v, d_model) + nn.init.xavier_normal_(self.fc.weight) + + self.dropout = nn.Dropout(dropout) + + def forward(self, q, k, v): + + d_k, d_v, n_head = self.d_k, self.d_v, self.n_head + + + sz_b, len_q, _ = q.size() + sz_b, len_k, _ = k.size() + sz_b, len_v, _ = v.size() + residual = q + + q = self.w_qs(q).view(sz_b, len_q, n_head, d_k) + k = self.w_ks(k).view(sz_b, len_k, n_head, d_k) + v = self.w_vs(v).view(sz_b, len_v, n_head, d_v) + + q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk + k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk + v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dv + + x, attn = self.attention(q, k, v) + + x = x.view(n_head, sz_b, len_q, d_v) + x = x.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv) + + x = self.dropout(self.fc(x)) + x = self.layer_norm(x + residual) + + return x, attn + class BiPoolingLinearClassifier(PoolingLinearClassifier): "Create a linear classifier with pooling." @@ -163,4 +277,4 @@ def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_s model.reset() return model -#endregion \ No newline at end of file +#endregion