File size: 1,850 Bytes
db3c893 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
#include "data_loader.hpp"
#include "training.hpp"
#include <iostream>
TrainConfig parse_cmd(int argc, char** argv) {
TrainConfig cfg;
for (int i = 1; i < argc; ++i) {
std::string a = argv[i];
auto next = [&](const char* name) -> std::string {
if (i + 1 >= argc) { std::cerr << "Missing value for " << name << "\n"; std::exit(1); }
return std::string(argv[++i]);
};
if (a == "--data_dir") cfg.data_dir = next("--data_dir");
else if (a == "--epochs") cfg.epochs = std::stoi(next("--epochs"));
else if (a == "--batch") cfg.batch = std::stoi(next("--batch"));
else if (a == "--lr") cfg.lr = std::stof(next("--lr"));
else if (a == "--fungi") cfg.fungi_count = std::stoi(next("--fungi"));
else if (a == "--wd") cfg.wd = std::stof(next("--wd"));
else if (a == "--accuracy_tolerance") cfg.accuracy_tolerance = std::stof(next("--accuracy_tolerance"));
else if (a == "--seed") cfg.seed = std::stoul(next("--seed")); // Original default seed was 1337u
else std::cerr << "Unknown arg: " << a << "\n";
}
return cfg;
}
int main(int argc, char** argv) {
try {
TrainConfig cfg = parse_cmd(argc, argv);
std::cout << "[INFO] Loading Fashion-MNIST training data from: " << cfg.data_dir << "\n";
FashionMNISTSet train = load_fashion_mnist_data(cfg.data_dir, true);
std::cout << "[INFO] Loading Fashion-MNIST test data from: " << cfg.data_dir << "\n";
FashionMNISTSet test = load_fashion_mnist_data(cfg.data_dir, false);
train_model(train, test, cfg);
std::cout << "[DONE] Training finished.\n";
return 0;
} catch (const std::exception& e) {
std::cerr << "[FATAL] " << e.what() << "\n";
return 1;
}
}
|