| | |
|---|
| | | torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]]) |
|---|
| | | if torch_version >= (1, 8): |
|---|
| | | import torch.fx |
|---|
| | | torch.fx.wrap('preprocess_for_attn') |
|---|
| | | |
|---|
| | | torch.fx.wrap("preprocess_for_attn") |
|---|
| | | |
|---|
| | | |
|---|
| | | class MultiHeadedAttentionSANMDecoderExport(nn.Module): |
|---|
| | |
|---|
| | | matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) |
|---|
| | | matrix_bd = self.rel_shift(matrix_bd) |
|---|
| | | |
|---|
| | | scores = (matrix_ac + matrix_bd) / math.sqrt( |
|---|
| | | self.d_k |
|---|
| | | ) # (batch, head, time1, time2) |
|---|
| | | scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) |
|---|
| | | |
|---|
| | | return self.forward_attention(v, scores, mask) |
|---|
| | | |
|---|
| | |
|---|
| | | new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) |
|---|
| | | context_layer = context_layer.view(new_context_layer_shape) |
|---|
| | | return self.linear_out(context_layer) # (batch, time1, d_model) |
|---|
| | | |
|---|
