here is my model
# coding: utf-8
import torch
import torch.nn as nn
class ScaleAndShift(nn.Module):
def __init__(self, vocab_size):
super().__init__()
self.scale = nn.Parameter(torch.zeros(vocab_size), requires_grad=True)
self.shift = nn.Parameter(torch.zeros(vocab_size), requires_grad=True)
def forward(self, x):
outputs = torch.exp(self.scale) * x + self.shift
return outputs
class SelfModulatedLayerNormalization(nn.Module):
def __init__(self, embed_dim: int, num_hidden: int):
super().__init__()
self.layer_norm = nn.LayerNorm(embed_dim)
self.beta_dense_1 = nn.Linear(num_hidden * 2, num_hidden // 2) # 乘以2是因为encoder是双向LSTM
self.beta_dense_2 = nn.Linear(num_hidden // 2, embed_dim)
self.gamma_dense_1 = nn.Linear(num_hidden * 2, num_hidden // 2)
self.gamma_dense_2 = nn.Linear(num_hidden // 2, embed_dim)
self.relu = nn.ReLU()
def forward(self, inputs):
inputs, cond = inputs
inputs = self.layer_norm(inputs) # inputs.shape
beta = self.beta_dense_1(cond)
beta = self.beta_dense_2(beta)
gamma = self.gamma_dense_1(cond)
gamma = self.gamma_dense_2(gamma)
for _ in range(len(inputs.shape) - len(cond.shape)):
beta = beta.unsqueeze(1)
gamma = gamma.unsqueeze(1)
return inputs * (gamma + 1) + beta
class BiDirectionalLSTM(nn.Module):
def __init__(self, input_size: int, hidden_size: int, num_layers: int):
super().__init__()
self.lstm = nn.LSTM(input_size, hidden_size=hidden_size, num_layers=num_layers, bidirectional=True, batch_first=True)
def forward(self, x):
encoder_outputs, _ = self.lstm(x)
return encoder_outputs
class MyMultiHeadAttention(nn.Module):
def __init__(self, input_size, num_heads, size_per_head):
super().__init__()
self.q_dense = nn.Linear(input_size, size_per_head * num_heads)
self.k_dense = nn.Linear(input_size, size_per_head * num_heads)
self.v_dense = nn.Linear(input_size, size_per_head * num_heads)
self.multi_head_attn = nn.MultiheadAttention(size_per_head * num_heads, num_heads)
def forward(self, q, k, v):
q = self.q_dense(q)
k = self.k_dense(k)
v = self.v_dense(v)
outputs, _ = self.multi_head_attn(q.permute(1, 0, 2), k.permute(1, 0, 2), v.permute(1, 0, 2)) # [seg_lem, batch_size, size_per_head * num_heads]
return outputs.permute(1, 0, 2)
class Encoder(nn.Module):
def __init__(self, embed_dim: int, num_layers: int, hidden_size: int):
super().__init__()
self.lstm_1 = BiDirectionalLSTM(input_size=embed_dim, hidden_size=hidden_size, num_layers=num_layers)
self.layer_norm_1 = nn.LayerNorm(embed_dim)
self.lstm_2 = BiDirectionalLSTM(input_size=hidden_size * 2, hidden_size=hidden_size, num_layers=num_layers)
self.layer_norm_2 = nn.LayerNorm(hidden_size * 2)
def forward(self, x, x_mask):
# x: [batch_size, max_len, embed_size]
# x_mask: [batch_size, max_len, 1]
x = self.layer_norm_1(x) # [batch_size, max_len, embed_size]
x = self.lstm_1(x) # [batch_size, max_len, (num_direction * hidden_size)]
x = self.layer_norm_2(x) # [batch_size, max_len, (num_direction * hidden_size)]
x = self.lstm_2(x) # [batch_size, max_len, (num_direction * hidden_size)]
x_max = seq_maxpool([x, x_mask]) # [batch_size, (num_direction * hidden_size)]
return x, x_max
class Decoder(nn.Module):
def __init__(self, embed_dim: int, hidden_size: int, num_layers: int):
super().__init__()
self.lstm_1 = nn.LSTM(embed_dim, hidden_size=hidden_size, num_layers=num_layers, bidirectional=False, batch_first=True)
self.layer_norm_1 = SelfModulatedLayerNormalization(embed_dim, hidden_size // 2) # 除以2是为了与encoder的hidden size相等
self.lstm_2 = nn.LSTM(hidden_size, hidden_size=hidden_size, num_layers=num_layers, bidirectional=False, batch_first=True)
self.layer_norm_2 = SelfModulatedLayerNormalization(hidden_size, hidden_size // 2)
self.layer_norm_3 = SelfModulatedLayerNormalization(hidden_size, hidden_size // 2)
def forward(self, y, x_max):
# y: [batch_size, max_len, embed_size]
# x_max: [batch_size, (num_direction * hidden_size)] hidden_size from encoder
y = self.layer_norm_1([y, x_max]) # [batch_size, max_len, embed_size]
y, _ = self.lstm_1(y) # [batch_size, max_len, hidden_size)]
y = self.layer_norm_2([y, x_max]) # [batch_size, max_len, hidden_size]
y, _ = self.lstm_2(y) # [batch_size, max_len, hidden_size]
output = self.layer_norm_3([y, x_max]) # [batch_size, max_len, hidden_size]
return output
class Seq2Seq(nn.Module):
def __init__(self, vocab_size: int, embed_dim: int, num_layers: int, hidden_size: int, num_heads: int, size_per_head: int):
super().__init__()
self.embed_x = nn.Embedding(vocab_size, embedding_dim=embed_dim)
self.embed_y = nn.Embedding(vocab_size, embedding_dim=embed_dim)
self.encoder = Encoder(embed_dim, num_layers, hidden_size // 2) # 除以2是因为encoder是双向LSTM
self.decoder = Decoder(embed_dim, hidden_size, num_layers)
self.attn = MyMultiHeadAttention(hidden_size, num_heads, size_per_head)
self.dense_1 = nn.Linear(size_per_head * num_heads + hidden_size, embed_dim)
self.relu = nn.LeakyReLU(0.2)
self.dense_2 = nn.Linear(embed_dim, vocab_size)
self.scale_shift = ScaleAndShift(vocab_size)
def forward(self, x, x_mask, x_one_hot, y):
# x: [batch_size, x_max_len]
# x_mask: [batch_size, x_max_len, 1]
# x_one_hot: [batch_size, vocab_size ]
# y: [batch_size, y_max_len]
x = self.embed_x(x)
y = self.embed_y(y)
x, x_max = self.encoder(x, x_mask) # [batch_size, x_max_len, hidden_size], [batch_size, hidden_size]
y = self.decoder(y, x_max) # [batch_size, y_max_len, hidden_size]
xy = self.attn(y, x, x) # [batch_size, y_max_len, size_per_head * num_heads]
xy1 = torch.cat((y, xy), dim=len(y.shape) - 1) # [batch_size, y_max_len, size_per_head * num_heads + hidden_size]
xy1 = self.dense_1(xy1) # [batch_size, y_max_len, embed_dim]
xy1 = self.relu(xy1) # [batch_size, y_max_len, embed_dim]
xy1 = self.dense_2(xy1) # [batch_size, y_max_len, vocab_size ]
x_pior = self.scale_shift(x_one_hot) # [batch_size, vocab_size]
output = (xy1 + x_pior.unsqueeze(1)) / 2 # [batch_size, y_max_len, vocab_size]
return output
def seq_maxpool(x):
seq, mask = x
seq -= (1 - mask) * 1e10
return torch.max(seq, 1).values
def seq_avgpool(x):
seq, mask = x
return torch.sum(seq * mask, 1) / (torch.sum(mask, 1) + 1e-6)
