fine tune

clax/clax.py

@@ -5,6 +5,7 @@
 import jax.random as random
 import optax
 from jax import jit
+from jaxopt import LBFGS
 from tqdm import tqdm
 
 from clax.network import DataLoader, Network, TrainState
@@ -180,6 +181,28 @@ def fit(self, samples, labels, epochs=10, batch_size=1024, **kwargs):
             train=False,
         )
 
+    def fine_tune(self, samples, labels, **kwargs):
+        """Fine tune the classifier with a 2nd order solver"""
+        labels = jnp.array(labels, dtype=int)
+        samples = jnp.array(samples, dtype=jnp.float32)
+        if not self.state:
+            raise ValueError("Classifier not trained yet.")
+
+        def fn(x):
+            return self.loss(x, self.state.batch_stats, samples, labels, self.rng)[0]
+
+        tuner = LBFGS(fn)
+        res = tuner.run(self.state.params)
+        self.state = self.state.replace(params=res[0])
+        self._predict_weight = lambda x: self.state.apply_fn(
+            {
+                "params": self.state.params,
+                "batch_stats": self.state.batch_stats,
+            },
+            x,
+            train=False,
+        )
+
     def predict(self, samples):
         """Predict the class (log) - probabilities for the provided samples.
 
@@ -257,6 +280,30 @@ def fit(self, samples_a, samples_b, epochs=10, batch_size=1024, **kwargs):
             train=False,
         )
 
+    def fine_tune(self, samples_a, samples_b, **kwargs):
+        """Fine tune the classifier with a 2nd order solver"""
+        samples_a = jnp.array(samples_a, dtype=jnp.float32)
+        samples_b = jnp.array(samples_b, dtype=jnp.float32)
+        if not self.state:
+            raise ValueError("Classifier not trained yet.")
+
+        def fn(x):
+            return self.loss(x, self.state.batch_stats, samples_a, samples_b, self.rng)[
+                0
+            ]
+
+        tuner = LBFGS(fn)
+        res = tuner.run(self.state.params)
+        self.state = self.state.replace(params=res[0])
+        self._predict_weight = lambda x: self.state.apply_fn(
+            {
+                "params": self.state.params,
+                "batch_stats": self.state.batch_stats,
+            },
+            x,
+            train=False,
+        )
+
 
 class Regressor(Classifier):
     """Regressor class wrapping a basic jax multiclass regressor."""

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "flax >= 0.8.2",
     "tqdm >= 4.62.0",
     "optax >= 0.2.2",
+    "jaxopt",
 ]
 
 [options.extras_require]