Fixed the model that used a dictionary to represent layers. Now use M…

…oduleDict
leschultz · Apr 19, 2024 · ab3b041 · ab3b041
1 parent 694c7f0
commit ab3b041
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 29 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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'):