药学数据模型构建3--LSTM

class LSTMModel(nn.Module):

    def __init__(self, input_size, output_size):

        super(LSTMModel, self).__init__()

        hidden_layer_size=100

        self.hidden_layer_size = hidden_layer_size

        self.lstm1 = nn.LSTM(input_size, hidden_layer_size, batch_first=True)

        self.lstm2 = nn.LSTM(input_size, hidden_layer_size, num_layers=2, batch_first=True)

        self.lstm3 = nn.LSTM(input_size, hidden_layer_size, num_layers=3, batch_first=True)

        self.linear = nn.Linear(hidden_layer_size, output_size)

    def forward(self, x, option=1):

        # x(batch_size, seq_len, input_size)

        # 初始化隐藏状态和细胞状态

        # [num_layers * num_directions, batch, hidden_size]

        h0 = torch.zeros(option, x.size(0), self.hidden_layer_size).to(x.device)

        c0 = torch.zeros(option, x.size(0), self.hidden_layer_size).to(x.device)

        if option == 1:

            out, _ = self.lstm1(x, (h0, c0))

        elif option == 2:

            out, _ = self.lstm2(x, (h0, c0))

        elif option == 3:

            out, _ = self.lstm3(x, (h0, c0))

        else:

            raise ValueError("option 输入不合法")        

        # 前向传播

        # out, _ = self.lstm(x, (h0, c0))

        # [batch_size, seq_length, num_directions * hidden_size]

        out = self.linear(out[:, -1, :])

        return out