Merge pull request #65 from DifferentiableUniverseInitiative/dlanzier…

…i/add_mae MAE loss function
DifferentiableUniverseInitiative · Nov 20, 2023 · 910a933 · 910a933
2 parents 42495c9 + 82b6968
commit 910a933
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/sbi_lens/normflow/train_model.py b/sbi_lens/normflow/train_model.py
@@ -18,6 +18,13 @@ def loss_mse(self, params, theta, x, state_resnet):
 
         return loss, opt_state_resnet
 
+    def loss_mae(self, params, theta, x, state_resnet):
+        y, opt_state_resnet = self.compressor.apply(params, state_resnet, None, x)
+
+        loss = jnp.mean(jnp.sum(jnp.absolute(y - theta), axis=1))
+
+        return loss, opt_state_resnet
+
     def loss_vmim(self, params, theta, x, state_resnet):
         y, opt_state_resnet = self.compressor.apply(params, state_resnet, None, x)
         log_prob = self.nf.apply(params, theta, y)

diff --git a/scripts/train_compressor.py b/scripts/train_compressor.py
@@ -122,6 +122,8 @@ def __call__(self, y):
     parameters_compressor = hk.data_structures.merge(parameters_resnet, params_nf)
 elif args.loss == "train_compressor_mse":
     parameters_compressor = parameters_resnet
+elif args.loss == "train_compressor_mae":
+    parameters_compressor = parameters_resnet
 elif args.loss == "train_compressor_gnll":
     parameters_compressor = parameters_resnet
 
@@ -197,6 +199,8 @@ def __call__(self, y):
     l_name = "vmim"
 elif args.loss == "train_compressor_mse":
     l_name = "mse"
+elif args.loss == "train_compressor_mae":
+    l_name = "mae"
 elif args.loss == "train_compressor_gnll":
     l_name = "gnll"