| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778 |
- import torch.nn as nn
- import torch
- class CNNLSTM(nn.Module):
- def __init__(
- self,
- n_features,
- pred_len,
- cnn_channels=64,
- lstm_hidden=128,
- lstm_layers=2,
- dropout=0.2,
- ):
- super().__init__()
- # Conv1d 输入:[batch, n_features, input_len]
- self.cnn = nn.Sequential(
- nn.Conv1d(
- in_channels=n_features,
- out_channels=cnn_channels,
- kernel_size=3,
- padding=1,
- ),
- nn.BatchNorm1d(cnn_channels),
- nn.ReLU(),
- nn.Dropout(dropout),
- nn.Conv1d(
- in_channels=cnn_channels,
- out_channels=cnn_channels,
- kernel_size=3,
- padding=1,
- ),
- nn.BatchNorm1d(cnn_channels),
- nn.ReLU(),
- nn.Dropout(dropout),
- )
- # LSTM 输入:[batch, input_len, cnn_channels]
- self.lstm = nn.LSTM(
- input_size=cnn_channels,
- hidden_size=lstm_hidden,
- num_layers=lstm_layers,
- batch_first=True,
- dropout=dropout if lstm_layers > 1 else 0.0,
- )
- self.head = nn.Sequential(
- nn.Linear(lstm_hidden, 128),
- nn.ReLU(),
- nn.Dropout(dropout),
- nn.Linear(128, pred_len),
- )
- def forward(self, x):
- # x: [batch, input_len, n_features]
- x = x.permute(0, 2, 1) # -> [batch, n_features, input_len]
- x = self.cnn(x) # -> [batch, cnn_channels, input_len]
- x = x.permute(0, 2, 1) # -> [batch, input_len, cnn_channels]
- lstm_out, _ = self.lstm(x) # -> [batch, input_len, lstm_hidden]
- last_hidden = lstm_out[:, -1, :]
- y_pred = self.head(last_hidden)
- return y_pred
- if __name__ == "__main__":
- model = CNNLSTM(n_features=10, pred_len=288, cnn_channels=64, lstm_hidden=32, lstm_layers=2, dropout=0.2)
- # 构造随机输入数据: [batch, input_len, n_features]
- batch_size = 4
- input_len = 288
- n_features = 10
- x = torch.randn(batch_size, input_len, n_features)
- # 前向传播
- y_pred = model(x)
- print("输出 shape:", y_pred.shape)
|