Actually use lasso

EleutherAI · Nov 1, 2023 · 739e970 · 739e970
1 parent c4d23e1
commit 739e970
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/elk/training/classifier.py b/elk/training/classifier.py
@@ -89,11 +89,15 @@ def fit(
         if lasso:
             from cuml import Lasso
 
-            model = Lasso(alpha=alpha)
-            model.fit(x.cpu().numpy(), y.cpu().numpy())
-            self.linear.weight.data = torch.from_numpy(model.coef_.T)
-            self.linear.bias.data = torch.from_numpy(model.intercept_)
-            return float(model.loss_)
+            model = Lasso(alpha=alpha, selection="random")
+            model.fit(x, y)
+
+            W = torch.as_tensor(model.coef_).unsqueeze(0).to(x.device)
+            b = torch.as_tensor(model.intercept_).to(x.device)
+
+            self.linear.weight.data = W
+            self.linear.bias.data = b
+            return 0.0
 
         optimizer = torch.optim.LBFGS(
             self.parameters(),

diff --git a/elk/training/supervised.py b/elk/training/supervised.py
@@ -57,7 +57,7 @@ def train_supervised(
         ).classifiers
     elif mode == "single":
         lr_model = Classifier(X.shape[-1], device=device, eraser=eraser)
-        lr_model.fit(X, train_labels)
+        lr_model.fit(X, train_labels, lasso=lasso)
         return [lr_model]
     else:
         raise ValueError(f"Unknown mode: {mode}")