240 lines
8.4 KiB
Python
240 lines
8.4 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
from .transformer_model import TransformerEncoder, TransformerDecoder
|
|||
|
|
|
|||
|
|
# 定义mLSTM单元
|
|||
|
|
class mLSTMCell(nn.Module):
|
|||
|
|
def __init__(self, input_size: int, hidden_size: int, bias: bool = True) -> None:
|
|||
|
|
super().__init__()
|
|||
|
|
|
|||
|
|
self.input_size = input_size
|
|||
|
|
self.hidden_size = hidden_size
|
|||
|
|
self.bias = bias
|
|||
|
|
|
|||
|
|
# 初始化权重和偏置
|
|||
|
|
self.W_i = nn.Parameter(
|
|||
|
|
nn.init.xavier_uniform_(torch.zeros(input_size, hidden_size)),
|
|||
|
|
requires_grad=True,
|
|||
|
|
)
|
|||
|
|
self.W_f = nn.Parameter(
|
|||
|
|
nn.init.xavier_uniform_(torch.zeros(input_size, hidden_size)),
|
|||
|
|
requires_grad=True,
|
|||
|
|
)
|
|||
|
|
self.W_o = nn.Parameter(
|
|||
|
|
nn.init.xavier_uniform_(torch.zeros(input_size, hidden_size)),
|
|||
|
|
requires_grad=True,
|
|||
|
|
)
|
|||
|
|
self.W_q = nn.Parameter(
|
|||
|
|
nn.init.xavier_uniform_(torch.zeros(input_size, hidden_size)),
|
|||
|
|
requires_grad=True,
|
|||
|
|
)
|
|||
|
|
self.W_k = nn.Parameter(
|
|||
|
|
nn.init.xavier_uniform_(torch.zeros(input_size, hidden_size)),
|
|||
|
|
requires_grad=True,
|
|||
|
|
)
|
|||
|
|
self.W_v = nn.Parameter(
|
|||
|
|
nn.init.xavier_uniform_(torch.zeros(input_size, hidden_size)),
|
|||
|
|
requires_grad=True,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if self.bias:
|
|||
|
|
self.B_i = nn.Parameter(torch.zeros(hidden_size), requires_grad=True)
|
|||
|
|
self.B_f = nn.Parameter(torch.zeros(hidden_size), requires_grad=True)
|
|||
|
|
self.B_o = nn.Parameter(torch.zeros(hidden_size), requires_grad=True)
|
|||
|
|
self.B_q = nn.Parameter(torch.zeros(hidden_size), requires_grad=True)
|
|||
|
|
self.B_k = nn.Parameter(torch.zeros(hidden_size), requires_grad=True)
|
|||
|
|
self.B_v = nn.Parameter(torch.zeros(hidden_size), requires_grad=True)
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
x: torch.Tensor,
|
|||
|
|
internal_state: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
|
|||
|
|
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
|
|||
|
|
# 获取内部状态
|
|||
|
|
C, n, m = internal_state
|
|||
|
|
|
|||
|
|
# 计算输入、遗忘、输出、查询、键和值门
|
|||
|
|
i_tilda = (
|
|||
|
|
torch.matmul(x, self.W_i) + self.B_i
|
|||
|
|
if self.bias
|
|||
|
|
else torch.matmul(x, self.W_i)
|
|||
|
|
)
|
|||
|
|
f_tilda = (
|
|||
|
|
torch.matmul(x, self.W_f) + self.B_f
|
|||
|
|
if self.bias
|
|||
|
|
else torch.matmul(x, self.W_f)
|
|||
|
|
)
|
|||
|
|
o_tilda = (
|
|||
|
|
torch.matmul(x, self.W_o) + self.B_o
|
|||
|
|
if self.bias
|
|||
|
|
else torch.matmul(x, self.W_o)
|
|||
|
|
)
|
|||
|
|
q_t = (
|
|||
|
|
torch.matmul(x, self.W_q) + self.B_q
|
|||
|
|
if self.bias
|
|||
|
|
else torch.matmul(x, self.W_q)
|
|||
|
|
)
|
|||
|
|
k_t = (
|
|||
|
|
torch.matmul(x, self.W_k) / torch.sqrt(torch.tensor(self.hidden_size))
|
|||
|
|
+ self.B_k
|
|||
|
|
if self.bias
|
|||
|
|
else torch.matmul(x, self.W_k) / torch.sqrt(torch.tensor(self.hidden_size))
|
|||
|
|
)
|
|||
|
|
v_t = (
|
|||
|
|
torch.matmul(x, self.W_v) + self.B_v
|
|||
|
|
if self.bias
|
|||
|
|
else torch.matmul(x, self.W_v)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 输入门的指数激活
|
|||
|
|
i_t = torch.exp(i_tilda)
|
|||
|
|
f_t = torch.sigmoid(f_tilda)
|
|||
|
|
o_t = torch.sigmoid(o_tilda)
|
|||
|
|
|
|||
|
|
# 稳定化状态
|
|||
|
|
m_t = torch.max(torch.log(f_t) + m, torch.log(i_t))
|
|||
|
|
i_prime = torch.exp(i_tilda - m_t)
|
|||
|
|
|
|||
|
|
# 协方差矩阵和归一化状态
|
|||
|
|
C_t = f_t.unsqueeze(-1) * C + i_prime.unsqueeze(-1) * torch.einsum(
|
|||
|
|
"bi, bk -> bik", v_t, k_t
|
|||
|
|
)
|
|||
|
|
n_t = f_t * n + i_prime * k_t
|
|||
|
|
|
|||
|
|
normalize_inner = torch.diagonal(torch.matmul(n_t, q_t.T))
|
|||
|
|
divisor = torch.max(
|
|||
|
|
torch.abs(normalize_inner), torch.ones_like(normalize_inner)
|
|||
|
|
)
|
|||
|
|
h_tilda = torch.einsum("bkj,bj -> bk", C_t, q_t) / divisor.view(-1, 1)
|
|||
|
|
h_t = o_t * h_tilda
|
|||
|
|
|
|||
|
|
return h_t, (C_t, n_t, m_t)
|
|||
|
|
|
|||
|
|
def init_hidden(
|
|||
|
|
self, batch_size: int, **kwargs
|
|||
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|||
|
|
return (
|
|||
|
|
torch.zeros(batch_size, self.hidden_size, self.hidden_size, **kwargs),
|
|||
|
|
torch.zeros(batch_size, self.hidden_size, **kwargs),
|
|||
|
|
torch.zeros(batch_size, self.hidden_size, **kwargs),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 定义mLSTM层
|
|||
|
|
class mLSTM(nn.Module):
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
input_size: int,
|
|||
|
|
hidden_size: int,
|
|||
|
|
num_layers: int,
|
|||
|
|
bias: bool = True,
|
|||
|
|
batch_first: bool = False,
|
|||
|
|
) -> None:
|
|||
|
|
super().__init__()
|
|||
|
|
self.input_size = input_size
|
|||
|
|
self.hidden_size = hidden_size
|
|||
|
|
self.num_layers = num_layers
|
|||
|
|
self.bias = bias
|
|||
|
|
self.batch_first = batch_first
|
|||
|
|
|
|||
|
|
self.cells = nn.ModuleList(
|
|||
|
|
[
|
|||
|
|
mLSTMCell(input_size if layer == 0 else hidden_size, hidden_size, bias)
|
|||
|
|
for layer in range(num_layers)
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
x: torch.Tensor,
|
|||
|
|
hidden_states=None,
|
|||
|
|
):
|
|||
|
|
# 如果batch_first为True，置换输入张量
|
|||
|
|
if self.batch_first:
|
|||
|
|
x = x.permute(1, 0, 2)
|
|||
|
|
|
|||
|
|
if hidden_states is None:
|
|||
|
|
hidden_states = self.init_hidden(x.size(1), device=x.device, dtype=x.dtype)
|
|||
|
|
else:
|
|||
|
|
# 检查隐藏状态是否正确长度
|
|||
|
|
if len(hidden_states) != self.num_layers:
|
|||
|
|
raise ValueError(
|
|||
|
|
f"Expected hidden states of length {self.num_layers}, but got {len(hidden_states)}"
|
|||
|
|
)
|
|||
|
|
if any(state[0].size(0) != x.size(1) for state in hidden_states):
|
|||
|
|
raise ValueError(
|
|||
|
|
f"Expected hidden states of batch size {x.size(1)}, but got {hidden_states[0][0].size(0)}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
outputs = []
|
|||
|
|
for t in range(x.size(0)):
|
|||
|
|
x_t = x[t]
|
|||
|
|
for layer in range(self.num_layers):
|
|||
|
|
if layer == 0:
|
|||
|
|
h_t, state = self.cells[layer](x_t, hidden_states[layer])
|
|||
|
|
else:
|
|||
|
|
h_t, state = self.cells[layer](h_t, hidden_states[layer])
|
|||
|
|
hidden_states[layer] = state
|
|||
|
|
outputs.append(h_t)
|
|||
|
|
|
|||
|
|
# 将输出堆叠成序列
|
|||
|
|
outputs = torch.stack(outputs)
|
|||
|
|
|
|||
|
|
# 如果batch_first为True，置换输出张量
|
|||
|
|
if self.batch_first:
|
|||
|
|
outputs = outputs.permute(1, 0, 2)
|
|||
|
|
|
|||
|
|
return outputs, hidden_states
|
|||
|
|
|
|||
|
|
def init_hidden(
|
|||
|
|
self, batch_size: int, **kwargs
|
|||
|
|
):
|
|||
|
|
return [cell.init_hidden(batch_size, **kwargs) for cell in self.cells]
|
|||
|
|
|
|||
|
|
# 结合mLSTM和Transformer的模型
|
|||
|
|
class MLSTMTransformer(nn.Module):
|
|||
|
|
def __init__(self, num_features, hidden_size=128, mlstm_layers=1,
|
|||
|
|
embed_dim=32, dense_dim=32, num_heads=4,
|
|||
|
|
dropout_rate=0.1, num_blocks=3, output_sequence_length=1):
|
|||
|
|
super().__init__()
|
|||
|
|
|
|||
|
|
# mLSTM部分
|
|||
|
|
self.mlstm = mLSTM(
|
|||
|
|
input_size=num_features,
|
|||
|
|
hidden_size=hidden_size,
|
|||
|
|
num_layers=mlstm_layers,
|
|||
|
|
batch_first=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Transformer部分
|
|||
|
|
self.input_embedding = nn.Linear(hidden_size, embed_dim)
|
|||
|
|
self.positional_encoding = nn.Parameter(torch.randn(1, 1000, embed_dim) * 0.1)
|
|||
|
|
|
|||
|
|
self.encoders = nn.ModuleList([TransformerEncoder(embed_dim, dense_dim, num_heads, dropout_rate) for _ in range(num_blocks)])
|
|||
|
|
self.decoders = nn.ModuleList([TransformerDecoder(embed_dim, dense_dim, num_heads, dropout_rate) for _ in range(num_blocks)])
|
|||
|
|
|
|||
|
|
self.output_layer = nn.Linear(embed_dim, 1)
|
|||
|
|
self.dropout = nn.Dropout(dropout_rate)
|
|||
|
|
self.output_sequence_length = output_sequence_length
|
|||
|
|
|
|||
|
|
def forward(self, inputs):
|
|||
|
|
# mLSTM处理
|
|||
|
|
mlstm_output, _ = self.mlstm(inputs)
|
|||
|
|
|
|||
|
|
# 将mLSTM输出嵌入到Transformer空间
|
|||
|
|
x = self.input_embedding(mlstm_output)
|
|||
|
|
x = x + self.positional_encoding[:, :x.size(1), :]
|
|||
|
|
x = self.dropout(x)
|
|||
|
|
|
|||
|
|
# 编码器部分
|
|||
|
|
encoder_outputs = x
|
|||
|
|
for encoder in self.encoders:
|
|||
|
|
encoder_outputs = encoder(encoder_outputs)
|
|||
|
|
|
|||
|
|
# 解码器部分
|
|||
|
|
decoder_inputs = encoder_outputs[:, -1:, :].expand(-1, self.output_sequence_length, -1)
|
|||
|
|
decoder_outputs = decoder_inputs
|
|||
|
|
|
|||
|
|
for decoder in self.decoders:
|
|||
|
|
decoder_outputs = decoder(decoder_outputs, encoder_outputs)
|
|||
|
|
|
|||
|
|
return self.output_layer(decoder_outputs)
|