Libtorch code gives different result than Pytorch

Viewed 32

I'm trying to rewrite this Pytorch notebook to Libtorch. My code results in a constant value of loss function during whole training. The Pytorch training loop works fine. Input tensors in Pytorch and Libtorch are of the same size and the training data is also the same. Could you help me spot the bug in Libtorch code?

Model class definition:

class LSTM : public torch::nn::Module {
private:
    torch::nn::LSTM lstm{nullptr};
    torch::nn::Linear linear{nullptr};
    int num_layers, hidden_size;

public:
    LSTM(int num_classes, int input_size, int hidden_size1, int num_layers1) {
        num_layers = num_layers1;
        hidden_size = hidden_size1;
        lstm = register_module("lstm",
                torch::nn::LSTM(torch::nn::LSTMOptions(input_size, hidden_size).num_layers(num_layers).batch_first(true)));
        linear = register_module("linear",
                torch::nn::Linear(hidden_size, num_classes));
    }

    torch::Tensor forward(torch::Tensor x) {
        torch::Tensor h_0 = torch::zeros({num_layers, x.size(0), hidden_size});
        torch::Tensor c_0 = torch::zeros({num_layers, x.size(0), hidden_size});

        // Propagate input through LSTM network
        torch::Tensor ula, h_out;
        std::tuple temp = std::tie(h_out, std::ignore);
        std::tie(ula, temp) = lstm->forward(x, std::tie(h_0, c_0));
        h_out = h_out.view({-1, hidden_size});
        auto out = linear->forward(h_out);
        return out;
    }
};

The code of training loop (x_train_tensors and y_train_tensors are equivalent to trainX and trainY from Pytorch notebook):

int num_epochs = 2000;
int learning_rate = 0.01;

int input_size = 1;
int hidden_size = 2;
int num_layers = 1;
int num_classes = 1;

LSTM lstm = LSTM(num_classes, input_size, hidden_size, num_layers);
torch::optim::Adam optimizer(lstm.parameters(), torch::optim::AdamOptions(learning_rate));

for (int i = 0; i < num_epochs; ++i) {
    torch::Tensor output = lstm.forward(x_train_tensors);
    optimizer.zero_grad();

    // get value of loss function
    auto loss = torch::mse_loss(output, y_train_tensors);
    loss.backward();
    optimizer.step();

    if (i % 1000 == 0) {
        std::ostringstream oss;
        oss << "Epoch: " << i << ", loss: " << std::setprecision(10) << loss.item<float>() << "\n";
        std::string s = oss.str();
        std::cout << s;
        }
}
0 Answers
Related