Spaces:
Sleeping
Sleeping
File size: 3,228 Bytes
2d48951 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
# transformer modules
class AddNorm(nn.Module):
"""残差连接后进行层归一化"""
def __init__(self, normalized, dropout):
super(AddNorm, self).__init__()
self.dropout = nn.Dropout(dropout)
self.ln = nn.LayerNorm(normalized)
def forward(self, x, y):
return self.ln(x + self.dropout(y))
class PositionWiseFFN(nn.Module):
"""基于位置的前馈⽹络"""
def __init__(self, ffn_input, ffn_hiddens,mlp_bias=True):
super(PositionWiseFFN, self).__init__()
self.ffn = nn.Sequential(
nn.Linear(ffn_input, ffn_hiddens, bias=mlp_bias),
nn.ReLU(),
nn.Linear(ffn_hiddens, ffn_input, bias=mlp_bias),
)
def forward(self, x):
return self.ffn(x)
class PositionalEncoding(nn.Module):
"""位置编码"""
def __init__(self, num_hiddens, dropout, max_len=1000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(dropout)
# 创建⼀个⾜够⻓的P
self.P = torch.zeros((1, max_len, num_hiddens))
X = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1) / torch.pow(10000, torch.arange(0, num_hiddens, 2,
dtype=torch.float32) / num_hiddens)
self.P[:, :, 0::2] = torch.sin(X)
self.P[:, :, 1::2] = torch.cos(X)
def forward(self, X):
X = X + self.P[:, :X.shape[1], :].to(X.device)
return self.dropout(X)
class AttentionEncode(nn.Module):
def __init__(self, dropout, embedding_size, num_heads,ffn=False):
super(AttentionEncode, self).__init__()
self.dropout = dropout
self.embedding_size = embedding_size
self.num_heads = num_heads
self.seq_len = 50
self.is_ffn = ffn
self.att = nn.MultiheadAttention(embed_dim=self.embedding_size,
num_heads=num_heads,
dropout=0.6
)
self.addNorm = AddNorm(normalized=[self.seq_len, self.embedding_size], dropout=self.dropout)
self.FFN = PositionWiseFFN(ffn_input=self.embedding_size, ffn_hiddens=self.embedding_size*2)
def forward(self, x):
bs,_,_ = x.size()
MHAtt, _ = self.att(x, x, x)
MHAtt_encode = self.addNorm(x, MHAtt)
if self.is_ffn:
ffn_in = MHAtt_encode # bs,seq_len,feat_dims
ffn_out = self.FFN(ffn_in)
MHAtt_encode = self.addNorm(ffn_in,ffn_out)
return MHAtt_encode
class FAN_encode(nn.Module):
def __init__(self, dropout, shape):
super(FAN_encode, self).__init__()
self.dropout = dropout
self.addNorm = AddNorm(normalized=[1, shape], dropout=self.dropout)
self.FFN = PositionWiseFFN(ffn_input=shape, ffn_hiddens=(2*shape))
self.ln = nn.LayerNorm(shape)
def forward(self, x):
#x = self.ln(x)
ffn_out = self.FFN(x)
encode_output = self.addNorm(x, ffn_out)
return encode_output |