42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import torch
|
||
import torch.nn.functional as F
|
||
|
||
class TimeDistributed(torch.nn.Module):
|
||
"""模拟tensorflow中的TimeDistributed包装层，因为pytorch似乎不提供这个。"""
|
||
|
||
layer: torch.nn.Module
|
||
"""内部节点"""
|
||
|
||
def __init__(self, layer: torch.nn.Module):
|
||
super(TimeDistributed, self).__init__()
|
||
self.layer = layer
|
||
|
||
def forward(self, x: torch.Tensor):
|
||
# 获取批次大小，时间步个数，特征个数
|
||
batch_size, time_steps, features = x.size()
|
||
# 把时间步维度合并到批次维度中然后运算，这样在其他层看来这就是不同的批次而已。
|
||
x = x.reshape(-1, features)
|
||
outputs: torch.Tensor = self.layer(x)
|
||
# 再把时间步维度还原出来
|
||
outputs = outputs.reshape(batch_size, time_steps, -1)
|
||
return outputs
|
||
|
||
|
||
class Rnn(torch.nn.Module):
|
||
"""循环神经网络"""
|
||
|
||
def __init__(self, vocab_size: int):
|
||
super(Rnn, self).__init__()
|
||
self.embedding = torch.nn.Embedding(vocab_size, 128)
|
||
self.lstm1 = torch.nn.LSTM(128, 128, batch_first=True, dropout=0.5)
|
||
self.lstm2 = torch.nn.LSTM(128, 128, batch_first=True, dropout=0.5)
|
||
self.timedfc = TimeDistributed(torch.nn.Linear(128, vocab_size))
|
||
|
||
def forward(self, x):
|
||
x = self.embedding(x)
|
||
x, _ = self.lstm1(x)
|
||
x, _ = self.lstm2(x)
|
||
x = self.timedfc(x)
|
||
return x
|
||
|