ゼロベースLSTM入門例pytorch

最近pytorchでLSTMモデルを作りましたが、ブロガーは両方とも基礎がないので、調べてみると多くの輪を回っています.いっそ、極めて簡単なLSTMコードの例をまとめて、参考にしてください.
torchのみが使用されています.nn.Moduleカスタムモデル
適当に選んだaccuracy_scoreを評価指標とする
また,特徴とラベルは任意に打たれており,収束を期待しないでください(
 
モデル構造の定義:

import torch
from sklearn.metrics import accuracy_score

#         ，   torch.nn.Module
#    __init__ forward    
class mylstm(torch.nn.Module):
    def __init__(self, lstm_input_size, lstm_hidden_size, lstm_batch, lstm_layers):
        #       
        super(mylstm, self).__init__()

        self.lstm_input_size, self.lstm_hidden_size = lstm_input_size, lstm_hidden_size
        self.lstm_layers, self.lstm_batch = lstm_layers, lstm_batch

        #   lstm 
        self.lstm_layer = torch.nn.LSTM(self.lstm_input_size, self.lstm_hidden_size, num_layers=self.lstm_layers, batch_first=True)
        #           
        self.out = torch.nn.Linear(self.lstm_hidden_size, 2)

    def forward(self, x):
        #   
        x = torch.sigmoid(x)
        # LSTM
        x, _ = self.lstm_layer(x)
        #          
        x = x[:, -1, :]
        #    
        x = self.out(x)
        return x

    def init_hidden(self):
        #         0
        return torch.zeros(self.lstm_batch, self.lstm_hidden_size)

データセットフィーチャーとラベル:

#     
train_feature = [
    [[0.1, 0.2, 0.3, 0.4], [0.1, 0.2, 0.3, 0.4], [0.1, 0.2, 0.3, 0.4]],
    [[0.4, 0.3, 0.2, 0.1], [0.4, 0.3, 0.2, 0.1], [0.4, 0.3, 0.2, 0.1]],
    [[0.2, 0.3, 0.5, 0.8], [0.2, 0.3, 0.5, 0.8], [0.2, 0.3, 0.5, 0.8]],
    [[0.7, 0.6, 0.5, 0.4], [0.7, 0.6, 0.5, 0.4], [0.7, 0.6, 0.5, 0.4]],
    [[0.1, 0.3, 0.5, 0.7], [0.1, 0.3, 0.5, 0.7], [0.1, 0.3, 0.5, 0.7]],
    [[0.5, 0.4, 0.2, 0.1], [0.5, 0.4, 0.2, 0.1], [0.5, 0.4, 0.2, 0.1]],
    [[0.2, 0.4, 0.6, 0.8], [0.2, 0.4, 0.6, 0.8], [0.2, 0.4, 0.6, 0.8]],
    [[0.7, 0.6, 0.3, 0.2], [0.7, 0.6, 0.3, 0.2], [0.7, 0.6, 0.3, 0.2]]]
#     
test_feature = [
    [[0.3, 0.4, 0.6, 0.8], [0.3, 0.4, 0.6, 0.8], [0.3, 0.4, 0.6, 0.8]],
    [[0.9, 0.6, 0.3, 0.2], [0.9, 0.6, 0.3, 0.2], [0.9, 0.6, 0.3, 0.2]]]
#   、     
train_label = [1, 0, 1, 0, 1, 0, 1, 0]
test_label = [1, 0]

モデルの定義、トレーニング、テスト:

#       
dataset_batch_size = 2
learning_rate = 0.001
lstm_input_size = 4
lstm_hidden_size = 4
lstm_batch = 2
lstm_layers = 1

#    
model = mylstm(lstm_input_size, lstm_hidden_size, lstm_batch, lstm_layers)
#      
hidden = model.init_hidden()
#          
criterion = torch.nn.CrossEntropyLoss()
#   Adam      
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(5):
    #          ，      
    optimizer.zero_grad()
    for i in range((len(train_feature) // dataset_batch_size) - 1):
        #  
        model.train()
        #   batch      
        x = train_feature[i * dataset_batch_size:(i + 1) * dataset_batch_size]
        x = torch.tensor(x)
        y = train_label[i * dataset_batch_size:(i + 1) * dataset_batch_size]
        y = torch.tensor(y)
        #   
        y_pred = model(x)
        loss = criterion(y_pred, y)
        #   
        optimizer.zero_grad()
        loss.backward()
        #   
        optimizer.step()

    #  
    model.eval()
    test_feature = torch.tensor(test_feature)
    out = model(test_feature)
    predict_result = torch.argmax(out, dim=1)
    #     
    acc = accuracy_score(test_label, predict_result)
    print('epoch:', epoch+1, ' acc:', acc)