dmlc--dgl
Hongzhi (Steve), Chen
be8763fa65
Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
35 行
1.1 KiB
Python
35 行
1.1 KiB
Python
import numpy as np
|
|
import torch as th
|
|
import torch.nn as nn
|
|
|
|
from .layers import clones
|
|
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
"Multi-Head Attention"
|
|
|
|
def __init__(self, h, dim_model):
|
|
"h: number of heads; dim_model: hidden dimension"
|
|
super(MultiHeadAttention, self).__init__()
|
|
self.d_k = dim_model // h
|
|
self.h = h
|
|
# W_q, W_k, W_v, W_o
|
|
self.linears = clones(nn.Linear(dim_model, dim_model, bias=False), 4)
|
|
|
|
def get(self, x, fields="qkv"):
|
|
"Return a dict of queries / keys / values."
|
|
batch_size = x.shape[0]
|
|
ret = {}
|
|
if "q" in fields:
|
|
ret["q"] = self.linears[0](x).view(batch_size, self.h, self.d_k)
|
|
if "k" in fields:
|
|
ret["k"] = self.linears[1](x).view(batch_size, self.h, self.d_k)
|
|
if "v" in fields:
|
|
ret["v"] = self.linears[2](x).view(batch_size, self.h, self.d_k)
|
|
return ret
|
|
|
|
def get_o(self, x):
|
|
"get output of the multi-head attention"
|
|
batch_size = x.shape[0]
|
|
return self.linears[3](x.view(batch_size, -1))
|