adjust the training to use gpu, if available

Author: sudkul

module-27-fairlearn/solution/train.py

@@ -11,9 +11,16 @@
 from sklearn.model_selection import train_test_split
 from sklearn.neural_network import MLPClassifier
 from sklearn.preprocessing import LabelEncoder, StandardScaler
+import torch
+import torch.nn as nn
+from skorch import NeuralNetClassifier
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")  # should print: Using device: cuda
 
 print("Loading dataset...")
-data = fetch_diabetes_hospital(as_frame=True)
+# data = fetch_diabetes_hospital(as_frame=True)
+data = fetch_diabetes_hospital(as_frame=True, cache=True)
 X, y = data.data, data.target
 
 # TODO: Extract the sensitive feature (gender)
@@ -38,9 +45,36 @@
 scaler = StandardScaler()
 X_train = scaler.fit_transform(X_train)
 X_test = scaler.transform(X_test)
+X_train = X_train.astype(np.float32)
+X_test  = X_test.astype(np.float32)
+y_train = y_train.astype(np.int64)
+
+class MLP(nn.Module):
+    def __init__(self, input_dim=24, output_dim=2):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, 64),
+            nn.ReLU(),
+            nn.Linear(64, 32),
+            nn.ReLU(),
+            nn.Linear(32, output_dim),
+        )
+    def forward(self, x):
+        return self.net(x)
 
 print("Training MLP...")
-model = MLPClassifier(hidden_layer_sizes=(64, 32), max_iter=100, random_state=42)
+# model = MLPClassifier(hidden_layer_sizes=(64, 32), max_iter=100, random_state=42)
+model = NeuralNetClassifier(
+    MLP,
+    module__input_dim=X_train.shape[1],
+    module__output_dim=len(np.unique(y_train)),
+    max_epochs=100,
+    lr=0.001,
+    batch_size=256,         # larger batch = better GPU utilization
+    device=device,          # ← this line is what actually puts the model on the GPU/CPU
+    train_split=None,       # mirrors original — no internal val split
+    verbose=1,
+)
 model.fit(X_train, y_train)
 
 joblib.dump({"model": model, "scaler": scaler}, "model.joblib")

requirements.txt

@@ -57,3 +57,4 @@ shap>=0.44.0
 # NOTE: requires scikit-learn==1.8.0 — conflicts with module-15 (needs <1.4.0)
 #       run in a separate virtual environment
 fairlearn>=0.10.0
+skorch==1.4.0