|
|
|
|
|
from __future__ import absolute_import |
|
from __future__ import division |
|
from __future__ import print_function |
|
|
|
import copy |
|
import logging |
|
import math |
|
|
|
from os.path import join as pjoin |
|
|
|
import torch |
|
import torch.nn as nn |
|
import numpy as np |
|
|
|
from torch.nn import CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm |
|
from torch.nn.modules.utils import _pair |
|
from scipy import ndimage |
|
|
|
|
|
ACT2FN = {"gelu": torch.nn.functional.gelu, "relu": torch.nn.functional.relu} |
|
|
|
|
|
class Attention(nn.Module): |
|
def __init__(self, config): |
|
super(Attention, self).__init__() |
|
self.num_attention_heads = config["num_heads"] |
|
self.attention_head_size = int(config['hidden_size'] / self.num_attention_heads) |
|
self.all_head_size = self.num_attention_heads * self.attention_head_size |
|
|
|
self.query = Linear(config['hidden_size'], self.all_head_size) |
|
self.key = Linear(config['hidden_size'], self.all_head_size) |
|
self.value = Linear(config['hidden_size'], self.all_head_size) |
|
|
|
|
|
self.out = Linear(self.all_head_size, config['hidden_size']) |
|
self.attn_dropout = Dropout(config["attention_dropout_rate"]) |
|
self.proj_dropout = Dropout(config["attention_dropout_rate"]) |
|
|
|
self.softmax = Softmax(dim=-1) |
|
|
|
def transpose_for_scores(self, x): |
|
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) |
|
x = x.view(*new_x_shape) |
|
return x.permute(0, 2, 1, 3) |
|
|
|
def forward(self, hidden_states): |
|
|
|
mixed_query_layer = self.query(hidden_states) |
|
mixed_key_layer = self.key(hidden_states) |
|
mixed_value_layer = self.value(hidden_states) |
|
|
|
query_layer = self.transpose_for_scores(mixed_query_layer) |
|
key_layer = self.transpose_for_scores(mixed_key_layer) |
|
value_layer = self.transpose_for_scores(mixed_value_layer) |
|
|
|
attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) |
|
attention_scores = attention_scores / math.sqrt(self.attention_head_size) |
|
attention_probs = self.softmax(attention_scores) |
|
attention_probs = self.attn_dropout(attention_probs) |
|
|
|
context_layer = torch.matmul(attention_probs, value_layer) |
|
context_layer = context_layer.permute(0, 2, 1, 3).contiguous() |
|
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) |
|
context_layer = context_layer.view(*new_context_layer_shape) |
|
attention_output = self.out(context_layer) |
|
attention_output = self.proj_dropout(attention_output) |
|
return attention_output |
|
|
|
|
|
class Mlp(nn.Module): |
|
def __init__(self, config): |
|
super(Mlp, self).__init__() |
|
self.fc1 = Linear(config['hidden_size'], config["mlp_dim"]) |
|
self.fc2 = Linear(config["mlp_dim"], config['hidden_size']) |
|
self.act_fn = ACT2FN["gelu"] |
|
self.dropout = Dropout(config["dropout_rate"]) |
|
self._init_weights() |
|
|
|
def _init_weights(self): |
|
nn.init.xavier_uniform_(self.fc1.weight) |
|
nn.init.xavier_uniform_(self.fc2.weight) |
|
nn.init.normal_(self.fc1.bias, std=1e-6) |
|
nn.init.normal_(self.fc2.bias, std=1e-6) |
|
|
|
def forward(self, x): |
|
x = self.fc1(x) |
|
x = self.act_fn(x) |
|
x = self.dropout(x) |
|
x = self.fc2(x) |
|
x = self.dropout(x) |
|
return x |
|
|
|
|
|
class Block(nn.Module): |
|
def __init__(self, config): |
|
super(Block, self).__init__() |
|
self.flag = config['num_heads'] |
|
self.hidden_size = config['hidden_size'] |
|
self.ffn_norm = LayerNorm(config['hidden_size'], eps=1e-6) |
|
self.ffn = Mlp(config) |
|
self.attn = Attention(config) |
|
self.attention_norm = LayerNorm(config['hidden_size'], eps=1e-6) |
|
|
|
def forward(self, x): |
|
h = x |
|
|
|
x = self.attention_norm(x) |
|
x = self.attn(x) |
|
x = x + h |
|
|
|
h = x |
|
x = self.ffn_norm(x) |
|
x = self.ffn(x) |
|
x = x + h |
|
return x |
|
|
|
|
|
class Encoder(nn.Module): |
|
def __init__(self, config): |
|
super(Encoder, self).__init__() |
|
|
|
self.layer = nn.ModuleList() |
|
self.encoder_norm = LayerNorm(config['hidden_size'], eps=1e-6) |
|
for _ in range(config["num_layers"]): |
|
layer = Block(config) |
|
self.layer.append(copy.deepcopy(layer)) |
|
|
|
def forward(self, hidden_states): |
|
for layer_block in self.layer: |
|
hidden_states = layer_block(hidden_states) |
|
encoded = self.encoder_norm(hidden_states) |
|
|
|
return encoded |
|
|
|
|
|
|
