diff --git a/reading_machine/include/Classifier.hpp b/reading_machine/include/Classifier.hpp index 2e2b51dd163140113d0f8d5d8ac95edb4b548a1b..7e423a3812ae3145759733a7feb1de080b3dbf71 100644 --- a/reading_machine/include/Classifier.hpp +++ b/reading_machine/include/Classifier.hpp @@ -24,6 +24,7 @@ class Classifier TransitionSet & getTransitionSet(); NeuralNetwork & getNN(); const std::string & getName() const; + int getNbParameters() const; }; #endif diff --git a/reading_machine/src/Classifier.cpp b/reading_machine/src/Classifier.cpp index a51ad089e1f569e594dedf435f6f990af9ba0b93..ce9c91e585216b09a5a2624fd8631434188cb8a5 100644 --- a/reading_machine/src/Classifier.cpp +++ b/reading_machine/src/Classifier.cpp @@ -20,6 +20,16 @@ Classifier::Classifier(const std::string & name, std::filesystem::path path, std initNeuralNetwork(definition); } +int Classifier::getNbParameters() const +{ + int nbParameters = 0; + + for (auto & t : nn->parameters()) + nbParameters += torch::numel(t); + + return nbParameters; +} + TransitionSet & Classifier::getTransitionSet() { return *transitionSet; diff --git a/trainer/src/MacaonTrain.cpp b/trainer/src/MacaonTrain.cpp index 2f98ff31f3c63be08fbbb3a64874664a3d13c49c..2b95753b153687b1e8fcf6d9b3b76c1b733a3c15 100644 --- a/trainer/src/MacaonTrain.cpp +++ b/trainer/src/MacaonTrain.cpp @@ -99,6 +99,8 @@ int MacaonTrain::main() ReadingMachine machine(machinePath.string()); + fmt::print(stderr, "[{}] Model has {} parameters\n", util::getTime(), util::int2HumanStr(machine.getClassifier()->getNbParameters())); + BaseConfig goldConfig(mcdFile, trainTsvFile, trainRawFile); BaseConfig devGoldConfig(mcdFile, computeDevScore ? (devRawFile.empty() ? devTsvFile : "") : devTsvFile, devRawFile); SubConfig config(goldConfig, goldConfig.getNbLines());