63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class LSTMModel(nn.Module):
|
|
def __init__(self, input_size=1, hidden_size=128, output_size=2):
|
|
super(LSTMModel, self).__init__()
|
|
|
|
self.input_size = input_size
|
|
self.hidden_size = hidden_size
|
|
self.output_size = output_size
|
|
|
|
self.conv1d = nn.Conv1d(input_size, 64, kernel_size=5, padding=2)
|
|
self.relu = nn.ReLU()
|
|
self.lstm1 = nn.LSTM(64, hidden_size, bidirectional=True, batch_first=True)
|
|
self.lstm2 = nn.LSTM(hidden_size * 2, hidden_size, bidirectional=True, batch_first=True)
|
|
self.lstm3 = nn.LSTM(hidden_size * 2, 64, bidirectional=False, batch_first=True)
|
|
self.fc1 = nn.Linear(64, 512)
|
|
self.fc2 = nn.Linear(512, 256)
|
|
self.fc3 = nn.Linear(256, 128)
|
|
self.fc_sbp = nn.Linear(128, 1)
|
|
self.fc_dbp = nn.Linear(128, 1)
|
|
|
|
def forward(self, x):
|
|
# 将输入传递给Conv1d层
|
|
x = self.conv1d(x.permute(0, 2, 1).contiguous())
|
|
x = self.relu(x)
|
|
x = x.permute(0, 2, 1).contiguous()
|
|
|
|
# 将输入传递给LSTM层
|
|
x, _ = self.lstm1(x)
|
|
x, _ = self.lstm2(x)
|
|
x, _ = self.lstm3(x)
|
|
|
|
# 只使用最后一个时间步的输出
|
|
x = x[:, -1, :]
|
|
|
|
# 将LSTM输出传递给全连接层
|
|
x = self.relu(self.fc1(x))
|
|
x = self.relu(self.fc2(x))
|
|
x = self.relu(self.fc3(x))
|
|
|
|
# 从两个Linear输出最终结果
|
|
sbp = self.fc_sbp(x)
|
|
dbp = self.fc_dbp(x)
|
|
|
|
return sbp, dbp
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 创建模型实例
|
|
model = LSTMModel()
|
|
|
|
# 定义示例输入
|
|
batch_size = 64
|
|
seq_len = 1250
|
|
input_size = 1
|
|
input_data = torch.randn(batch_size, seq_len, input_size)
|
|
|
|
# 将输入数据传递给模型
|
|
sbp, dbp = model(input_data)
|
|
print(sbp.shape, dbp.shape) # 输出: torch.Size([64, 1]) torch.Size([64, 1])
|