Add an option to pass some of the hyperparamters from the command line

tutorials/basics/feedforward_neural_network/src/main.cpp

@@ -2,37 +2,79 @@
 #include <torch/torch.h>
 #include <iostream>
 #include <iomanip>
+#include <string> // To parse command line arguments
 #include "neural_net.h"
 
-int main() {
-    std::cout << "FeedForward Neural Network\n\n";
+int main(int argc, char **argv)
+{
 
     // Device
     auto cuda_available = torch::cuda::is_available();
     torch::Device device(cuda_available ? torch::kCUDA : torch::kCPU);
-    std::cout << (cuda_available ? "CUDA available. Training on GPU." : "Training on CPU.") << '\n';
 
     // Hyper parameters
     const int64_t input_size = 784;
-    const int64_t hidden_size = 500;
+    int64_t hidden_size = 500;
     const int64_t num_classes = 10;
-    const int64_t batch_size = 100;
-    const size_t num_epochs = 5;
-    const double learning_rate = 0.001;
+    int64_t batch_size = 100;
+    size_t num_epochs = 5;
+    double learning_rate = 0.001;
+
+    for (int i = 1; i < argc; ++i)
+    {
+        if (std::string(argv[i]) == "--hidden_size" && i + 1 < argc)
+        {
+            hidden_size = std::atoi(argv[++i]);
+        }
+        if (std::string(argv[i]) == "--batch_size" && i + 1 < argc)
+        {
+            batch_size = std::atoi(argv[++i]);
+        }
+        if (std::string(argv[i]) == "--num_epochs" && i + 1 < argc)
+        {
+            num_epochs = std::atoi(argv[++i]);
+        }
+        if (std::string(argv[i]) == "--learning_rate" && i + 1 < argc)
+        {
+            learning_rate = std::atof(argv[++i]);
+        }
+        if (std::string(argv[i]) == "--help")
+        {
+            std::cout << "Example usage: \n";
+            std::cout << "./<executable> --hidden_size 500  --batch_size 100 --num_epochs 10 --learning_rate 10e-4\n";
+            std::cout << "If no arguements are passed default values will be used\n";
+            std::cout << "Use the argument < --hidden_size H > to pass the number of hidden units\n";
+            std::cout << "Use the argument < --batch_size B > to pass the batch size\n";
+            std::cout << "Use the argument < --num_epochs E > to pass the number of epochs\n";
+            std::cout << "Use the argument < --learning_rate LE >  to pass the learning rate \n";
+            std::cout << "Use the argument < --help >  to print this mmessage \n";
+            std::exit(-1);
+        }
+    }
+
+    std::cout << "Input size: " << input_size << "\n";
+    std::cout << "Number of hidden units: " << hidden_size << "\n";
+    std::cout << "Number of classes: " << num_classes << "\n";
+    std::cout << "batch size: " << batch_size << "\n";
+    std::cout << "Number of epochs: " << num_epochs << "\n";
+    std::cout << "Learning rate: " << learning_rate << "\n\n";
+
+    std::cout << "FeedForward Neural Network\n\n";
+    std::cout << (cuda_available ? "CUDA available. Training on GPU." : "Training on CPU.") << '\n';
 
     const std::string MNIST_data_path = "../../../../data/mnist/";
 
     // MNIST Dataset
     auto train_dataset = torch::data::datasets::MNIST(MNIST_data_path)
-        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
-        .map(torch::data::transforms::Stack<>());
+                             .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
+                             .map(torch::data::transforms::Stack<>());
 
     // Number of samples in the training set
     auto num_train_samples = train_dataset.size().value();
 
     auto test_dataset = torch::data::datasets::MNIST(MNIST_data_path, torch::data::datasets::MNIST::Mode::kTest)
-        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
-        .map(torch::data::transforms::Stack<>());
+                            .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
+                            .map(torch::data::transforms::Stack<>());
 
     // Number of samples in the testset
     auto num_test_samples = test_dataset.size().value();
@@ -57,12 +99,14 @@ int main() {
     std::cout << "Training...\n";
 
     // Train the model
-    for (size_t epoch = 0; epoch != num_epochs; ++epoch) {
+    for (size_t epoch = 0; epoch != num_epochs; ++epoch)
+    {
         // Initialize running metrics
         double running_loss = 0.0;
         size_t num_correct = 0;
 
-        for (auto& batch : *train_loader) {
+        for (auto &batch : *train_loader)
+        {
             auto data = batch.data.view({batch_size, -1}).to(device);
             auto target = batch.target.to(device);
 
@@ -89,7 +133,7 @@ int main() {
         auto accuracy = static_cast<double>(num_correct) / num_train_samples;
 
         std::cout << "Epoch [" << (epoch + 1) << "/" << num_epochs << "], Trainset - Loss: "
-            << sample_mean_loss << ", Accuracy: " << accuracy << '\n';
+                  << sample_mean_loss << ", Accuracy: " << accuracy << '\n';
     }
 
     std::cout << "Training finished!\n\n";
@@ -102,7 +146,8 @@ int main() {
     double running_loss = 0.0;
     size_t num_correct = 0;
 
-    for (const auto& batch : *test_loader) {
+    for (const auto &batch : *test_loader)
+    {
         auto data = batch.data.view({batch_size, -1}).to(device);
         auto target = batch.target.to(device);