CNN_LSTM.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import torch.nn as nn
  2. import torch
  3. class CNNLSTM(nn.Module):
  4. def __init__(
  5. self,
  6. n_features,
  7. pred_len,
  8. cnn_channels=64,
  9. lstm_hidden=128,
  10. lstm_layers=2,
  11. dropout=0.2,
  12. ):
  13. super().__init__()
  14. # Conv1d 输入:[batch, n_features, input_len]
  15. self.cnn = nn.Sequential(
  16. nn.Conv1d(
  17. in_channels=n_features,
  18. out_channels=cnn_channels,
  19. kernel_size=3,
  20. padding=1,
  21. ),
  22. nn.BatchNorm1d(cnn_channels),
  23. nn.ReLU(),
  24. nn.Dropout(dropout),
  25. nn.Conv1d(
  26. in_channels=cnn_channels,
  27. out_channels=cnn_channels,
  28. kernel_size=3,
  29. padding=1,
  30. ),
  31. nn.BatchNorm1d(cnn_channels),
  32. nn.ReLU(),
  33. nn.Dropout(dropout),
  34. )
  35. # LSTM 输入:[batch, input_len, cnn_channels]
  36. self.lstm = nn.LSTM(
  37. input_size=cnn_channels,
  38. hidden_size=lstm_hidden,
  39. num_layers=lstm_layers,
  40. batch_first=True,
  41. dropout=dropout if lstm_layers > 1 else 0.0,
  42. )
  43. self.head = nn.Sequential(
  44. nn.Linear(lstm_hidden, 128),
  45. nn.ReLU(),
  46. nn.Dropout(dropout),
  47. nn.Linear(128, pred_len),
  48. )
  49. def forward(self, x):
  50. # x: [batch, input_len, n_features]
  51. x = x.permute(0, 2, 1) # -> [batch, n_features, input_len]
  52. x = self.cnn(x) # -> [batch, cnn_channels, input_len]
  53. x = x.permute(0, 2, 1) # -> [batch, input_len, cnn_channels]
  54. lstm_out, _ = self.lstm(x) # -> [batch, input_len, lstm_hidden]
  55. last_hidden = lstm_out[:, -1, :]
  56. y_pred = self.head(last_hidden)
  57. return y_pred
  58. if __name__ == "__main__":
  59. model = CNNLSTM(n_features=10, pred_len=288, cnn_channels=64, lstm_hidden=32, lstm_layers=2, dropout=0.2)
  60. # 构造随机输入数据: [batch, input_len, n_features]
  61. batch_size = 4
  62. input_len = 288
  63. n_features = 10
  64. x = torch.randn(batch_size, input_len, n_features)
  65. # 前向传播
  66. y_pred = model(x)
  67. print("输出 shape:", y_pred.shape)