dmlc--dgl
Lingfan Yu
29dd22e666
* multi-process version of transformer * lots of fix * fix bugs and accum gradients for multiple batches * many fixes * minor * upd * set torch device * fix bugs * fix and minor * comments and clean up * uncomment viz code
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, 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))
|