dmlc--dgl
Zihao Ye
9f32554296
* change the signature of node/edge filter * upd filter * Support multi-dimension node feature in SPMV * push transformer * remove some experimental settings * stable version * hotfix * upd tutorial * upd README * merge * remove redundency * remove tqdm * several changes * Refactor * Refactor * tutorial train * fixed a bug * fixed perf issue * upd * change dir * move un-related to contrib * tutuorial code * remove redundency * upd * upd * upd * upd * improve viz * universal done * halt norm * fixed a bug * add draw graph * fixed several bugs * remove dependency on core * upd format of README * trigger * trigger * upd viz * trigger * add transformer tutorial * fix tutorial * fix readme * small fix on tutorials * url fix in readme * fixed func link * upd
34 行
1.1 KiB
Python
34 行
1.1 KiB
Python
import torch as th
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
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), 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))
|