Added learnable coefs for the loss

torch_modules/include/NeuralNetwork.hpp

@@ -11,9 +11,11 @@ class NeuralNetworkImpl : public torch::nn::Module, public NameHolder
   private :
 
   static torch::Device device;
+  std::map<std::string, torch::Tensor> lossParameters;
 
   public :
 
+  torch::Tensor getLossParameter(std::string state);
   virtual torch::Tensor forward(torch::Tensor input, const std::string & state) = 0;
   virtual torch::Tensor extractContext(Config & config) = 0;
   virtual void registerEmbeddings(bool loadPretrained) = 0;

torch_modules/src/ModularNetwork.cpp

@@ -81,7 +81,10 @@ ModularNetworkImpl::ModularNetworkImpl(std::string name, std::map<std::string,st
   mlp = register_module("mlp", MLP(currentOutputSize, mlpDef));
 
   for (auto & it : nbOutputsPerState)
+  {
     outputLayersPerState.emplace(it.first,register_module(fmt::format("output_{}",it.first), torch::nn::Linear(mlp->outputSize(), it.second)));
+    getLossParameter(it.first);
+  }
 }
 
 torch::Tensor ModularNetworkImpl::forward(torch::Tensor input, const std::string & state)

torch_modules/src/NeuralNetwork.cpp

@@ -2,6 +2,14 @@
 
 torch::Device NeuralNetworkImpl::device(getPreferredDevice());
 
+torch::Tensor NeuralNetworkImpl::getLossParameter(std::string state)
+{
+  if (lossParameters.count(state) == 0)
+    lossParameters[state] = register_parameter(fmt::format("lossParam_{}", state), torch::ones(1, torch::TensorOptions().device(NeuralNetworkImpl::getDevice()).requires_grad(true)));
+
+  return lossParameters[state];
+}
+
 float NeuralNetworkImpl::entropy(torch::Tensor probabilities)
 {
   if (probabilities.dim() != 1)

trainer/src/Trainer.cpp

@@ -247,7 +247,9 @@ float Trainer::processDataset(DataLoader & loader, bool train, bool printAdvance
       labels /= util::float2longScale;
     }
 
-    auto loss = machine.getClassifier(state)->getLossMultiplier(state)*machine.getClassifier(state)->getLossFunction()(prediction, labels);
+    auto lossParameter = machine.getClassifier(state)->getNN()->getLossParameter(state);
+
+    auto loss = machine.getClassifier(state)->getLossMultiplier(state)*machine.getClassifier(state)->getLossFunction()(prediction, labels)*(1.0/torch::exp(lossParameter)) + lossParameter;
     float lossAsFloat = 0.0;
     try
     {