bias_emb: [batch_size, 1, emb_size]
(self, bias_emb)
| 161 | self.softmax = nn.Softmax(dim=-1) |
| 162 | |
| 163 | def forward(self, bias_emb): |
| 164 | """ |
| 165 | bias_emb: [batch_size, 1, emb_size] |
| 166 | """ |
| 167 | att_scores = torch.matmul(bias_emb, self.index.transpose(-1, -2)) # [batch_size, clusters_k] |
| 168 | att_scores = self.softmax(att_scores) |
| 169 | |
| 170 | # [batch_size, input_size, output_size] |
| 171 | para_new = torch.matmul(att_scores, self.array) # [batch_size, input_size*output_size] |
| 172 | para_new = para_new.view(-1, self.output_size, self.input_size) |
| 173 | |
| 174 | return para_new |
| 175 | |
| 176 | def reg_loss(self, reg_weights=1e-2): |
| 177 | loss_1 = reg_weights * self.array.norm(2) |
nothing calls this directly
no outgoing calls
no test coverage detected