From ab3b04184ee962ecafda99e2037b0f5bc5f4082c Mon Sep 17 00:00:00 2001
From: leschultz <laneenriqueschultz@gmail.com>
Date: Fri, 19 Apr 2024 10:15:24 -0500
Subject: [PATCH] Fixed the model that used a dictionary to represent layers.
 Now use ModuleDict

---
 examples/synthetic/fit.py  | 10 ++++++---
 src/multilearn/datasets.py | 46 +++++++++++++++++++-------------------
 src/multilearn/models.py   |  4 ++--
 src/multilearn/utils.py    |  2 +-
 4 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/examples/synthetic/fit.py b/examples/synthetic/fit.py
index f200787..2faeea0 100644
--- a/examples/synthetic/fit.py
+++ b/examples/synthetic/fit.py
@@ -9,10 +9,10 @@ def main():
     lr = 1e-4
     batch_size = 32
     n_epochs = 1000
-    tasks = ['data1', 'data2', 'data3']
+    tasks = ['toy1', 'toy2', 'friedman1']
 
     # Data
-    X, y = datasets.toy()
+    X, y = datasets.load(tasks)
     data = datasets.splitter(
                              X,
                              y,
@@ -26,7 +26,11 @@ def main():
         data[k]['scaler'] = StandardScaler()
         data[k]['loss'] = nn.L1Loss()
 
-    model = models.MultiNet(tasks=tasks, input_arch={500: 1})
+    model = models.MultiNet(
+                            tasks=tasks,
+                            input_arch={10: 1},
+                            mid_arch={1024: 1, 16: 1},
+                            )
     optimizer = optim.Adam
 
     out = utils.train(
diff --git a/src/multilearn/datasets.py b/src/multilearn/datasets.py
index e5ee143..54e6958 100644
--- a/src/multilearn/datasets.py
+++ b/src/multilearn/datasets.py
@@ -82,38 +82,38 @@ def split(X, y, train_size=1.0, val_size=0.0, test_size=0.0):
     return data
 
 
-def toy(points=[1000, 900, 500]):
+def load(names):
 
-    X1 = np.random.uniform(size=(points[0], 3))
-    y1 = 3+X1[:, 0]+X1[:, 1]**3+np.log(X1[:, 2])
+    Xs = []
+    ys = []
+    for name in names:
 
-    X2 = np.random.uniform(size=(points[1], 3))
-    y2 = 3+X2[:, 0]+X2[:, 1]**3+X2[:, 2]
+        if name == 'toy1':
 
-    X3 = np.random.uniform(size=(points[2], 5))
-    y3 = (
-          10*np.sin(np.pi*X3[:, 0]*X3[:, 1])
-          + 20*(X3[:, 2]-0.5)**2
-          + 10*X3[:, 3]
-          + 5*X3[:, 4]
-          )
+            X = np.random.uniform(size=(1000, 3))
+            y = 3+X[:, 0]+X[:, 1]**3+np.log(X[:, 2])
 
-    X = [X1, X2, X3]
-    y = [y1, y2, y3]
+        elif name == 'toy2':
 
-    return X, y
+            X = np.random.uniform(size=(900, 3))
+            y = 3+X[:, 0]+X[:, 1]**3+X[:, 2]
 
+        elif name == 'friedman1':
 
-def load(names):
+            X = np.random.uniform(size=(500, 5))
+            y = (
+                 10*np.sin(np.pi*X[:, 0]*X[:, 1])
+                 + 20*(X[:, 2]-0.5)**2
+                 + 10*X[:, 3]
+                 + 5*X[:, 4]
+                 )
 
-    Xs = []
-    ys = []
-    for name in names:
-        path = os.path.join(data_path, f'{name}.csv')
-        df = pd.read_csv(path)
+        else:
+            path = os.path.join(data_path, f'{name}.csv')
+            df = pd.read_csv(path)
 
-        y = df['y'].values
-        X = df.drop('y', axis=1).values
+            y = df['y'].values
+            X = df.drop('y', axis=1).values
 
         Xs.append(X)
         ys.append(y)
diff --git a/src/multilearn/models.py b/src/multilearn/models.py
index b4daaa7..6a73a77 100644
--- a/src/multilearn/models.py
+++ b/src/multilearn/models.py
@@ -6,7 +6,7 @@ class MultiNet(nn.Module):
     def __init__(
                  self,
                  input_arch={},
-                 mid_arch={1024: 1, 32: 1, 16: 1},
+                 mid_arch={64: 1, 32: 1},
                  out_arch={},
                  tasks=[0],
                  ):
@@ -30,7 +30,7 @@ def make_layers(arch, is_out=False):
 
         def separate(arch, tasks, is_out=False):
 
-            separate = {}
+            separate = nn.ModuleDict()
             for t in tasks:
                 i = make_layers(arch, is_out)
                 separate[t] = i
diff --git a/src/multilearn/utils.py b/src/multilearn/utils.py
index db2b20e..146c1f5 100644
--- a/src/multilearn/utils.py
+++ b/src/multilearn/utils.py
@@ -135,7 +135,7 @@ def train(
     data_train = {}
     for key, value in data.items():
         for k, v in value.items():
-            if 'X_' in k:
+            if ('X_' in k) and ('scaler' in value.keys()):
                 value[k] = value['scaler'].transform(value[k])
 
             if (k != 'scaler') and (k != 'loss'):