Model creation and training for basic layers.

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch_model_ops.py

@@ -1,9 +1,11 @@
 """Boxes for defining PyTorch models."""
 
+import graphlib
 from lynxkite.core import ops, workspace
 from lynxkite.core.ops import Parameter as P
 import torch
 import torch_geometric as pyg
+from dataclasses import dataclass
 
 ENV = "PyTorch model"
 
@@ -70,7 +72,7 @@ reg(
 reg(
     "Activation",
     inputs=["x"],
-    params=[P.options("type", ["ReLU", "LeakyReLU", "Tanh", "Mish"])],
+    params=[P.options("type", ["ReLU", "Leaky ReLU", "Tanh", "Mish"])],
 )
 reg("Concatenate", inputs=["a", "b"], outputs=["x"])
 reg("Add", inputs=["a", "b"], outputs=["x"])
@@ -105,7 +107,10 @@ ops.register_passive_op(
     "Repeat",
     inputs=[ops.Input(name="input", position="top", type="tensor")],
     outputs=[ops.Output(name="output", position="bottom", type="tensor")],
-    params=[ops.Parameter.basic("times", 1, int)],
+    params=[
+        ops.Parameter.basic("times", 1, int),
+        ops.Parameter.basic("same_weights", True, bool),
+    ],
 )
 
 ops.register_passive_op(
@@ -117,24 +122,133 @@ ops.register_passive_op(
 )
 
 
-def build_model(ws: workspace.Workspace, inputs: dict):
+def _to_id(s: str) -> str:
+    """Replaces all non-alphanumeric characters with underscores."""
+    return "".join(c if c.isalnum() else "_" for c in s)
+
+
+@dataclass
+class ModelConfig:
+    model: torch.nn.Module
+    model_inputs: list[str]
+    model_outputs: list[str]
+    loss_inputs: list[str]
+    loss: torch.nn.Module
+    optimizer: torch.optim.Optimizer
+
+    def _forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        model_inputs = [inputs[i] for i in self.model_inputs]
+        output = self.model(*model_inputs)
+        if not isinstance(output, tuple):
+            output = (output,)
+        values = {k: v for k, v in zip(self.model_outputs, output)}
+        return values
+
+    def inference(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        # TODO: Do multiple batches.
+        self.model.eval()
+        return self._forward(inputs)
+
+    def train(self, inputs: dict[str, torch.Tensor]) -> float:
+        """Train the model for one epoch. Returns the loss."""
+        # TODO: Do multiple batches.
+        self.model.train()
+        self.optimizer.zero_grad()
+        values = self._forward(inputs)
+        values.update(inputs)
+        loss_inputs = [values[i] for i in self.loss_inputs]
+        loss = self.loss(*loss_inputs)
+        loss.backward()
+        self.optimizer.step()
+        return loss.item()
+
+
+def build_model(
+    ws: workspace.Workspace, inputs: dict[str, torch.Tensor]
+) -> ModelConfig:
     """Builds the model described in the workspace."""
     optimizers = []
+    nodes = {}
     for node in ws.nodes:
-        if node.op.name == "Optimizer":
-            optimizers.append(node)
+        nodes[node.id] = node
+        if node.data.title == "Optimizer":
+            optimizers.append(node.id)
     assert optimizers, "No optimizer found."
     assert len(optimizers) == 1, f"More than one optimizer found: {optimizers}"
     [optimizer] = optimizers
-    inputs = {n.id: [] for n in ws.nodes}
+    dependencies = {n.id: [] for n in ws.nodes}
+    edges = {}
+    # TODO: Dissolve repeat boxes here.
     for e in ws.edges:
-        inputs[e.target].append(e.source)
-    layers = []
-    # TODO: Create layers based on the workspace.
+        dependencies[e.target].append(e.source)
+        edges.setdefault((e.target, e.targetHandle), []).append(
+            (e.source, e.sourceHandle)
+        )
     sizes = {}
-    for k, v in inputs.items():
-        sizes[k] = v.size
-    layers.append((pyg.nn.Linear(sizes["x"], 1024), "x -> x"))
-    layers.append((torch.nn.LayerNorm(1024), "x -> x"))
-    m = pyg.nn.Sequential("x, edge_index", layers)
-    return m
+    for k, i in inputs.items():
+        sizes[k] = i.shape[-1]
+    ts = graphlib.TopologicalSorter(dependencies)
+    layers = []
+    loss_layers = []
+    in_loss = set()
+    cfg = {}
+    loss_inputs = set()
+    used_inputs = set()
+    for node_id in ts.static_order():
+        node = nodes[node_id]
+        t = node.data.title
+        p = node.data.params
+        for b in dependencies[node_id]:
+            if b in in_loss:
+                in_loss.add(node_id)
+        ls = loss_layers if node_id in in_loss else layers
+        nid = _to_id(node_id)
+        match t:
+            case "Linear":
+                [(ib, ih)] = edges[node_id, "x"]
+                i = _to_id(ib) + "_" + ih
+                used_inputs.add(i)
+                isize = sizes[i]
+                osize = isize if p["output_dim"] == "same" else int(p["output_dim"])
+                ls.append((torch.nn.Linear(isize, osize), f"{i} -> {nid}_x"))
+                sizes[f"{nid}_x"] = osize
+            case "Activation":
+                [(ib, ih)] = edges[node_id, "x"]
+                i = _to_id(ib) + "_" + ih
+                used_inputs.add(i)
+                f = getattr(torch.nn.functional, p["type"].lower().replace(" ", "_"))
+                ls.append((f, f"{i} -> {nid}_x"))
+                sizes[f"{nid}_x"] = sizes[i]
+            case "MSE loss":
+                [(xb, xh)] = edges[node_id, "x"]
+                xi = _to_id(xb) + "_" + xh
+                [(yb, yh)] = edges[node_id, "y"]
+                yi = _to_id(yb) + "_" + yh
+                loss_inputs.add(xi)
+                loss_inputs.add(yi)
+                in_loss.add(node_id)
+                loss_layers.append(
+                    (torch.nn.functional.mse_loss, f"{xi}, {yi} -> {nid}_loss")
+                )
+    cfg["model_inputs"] = used_inputs & inputs.keys()
+    cfg["model_outputs"] = loss_inputs - inputs.keys()
+    cfg["loss_inputs"] = loss_inputs
+    # Make sure the trained output is output from the last model layer.
+    outputs = ", ".join(cfg["model_outputs"])
+    layers.append((torch.nn.Identity(), f"{outputs} -> {outputs}"))
+    # Create model.
+    cfg["model"] = pyg.nn.Sequential(", ".join(used_inputs & inputs.keys()), layers)
+    # Make sure the loss is output from the last loss layer.
+    [(lossb, lossh)] = edges[optimizer, "loss"]
+    lossi = _to_id(lossb) + "_" + lossh
+    loss_layers.append((torch.nn.Identity(), f"{lossi} -> loss"))
+    # Create loss function.
+    cfg["loss"] = pyg.nn.Sequential(", ".join(loss_inputs), loss_layers)
+    assert not list(cfg["loss"].parameters()), (
+        f"loss should have no parameters: {list(cfg['loss'].parameters())}"
+    )
+    # Create optimizer.
+    p = nodes[optimizer].data.params
+    o = getattr(torch.optim, p["type"])
+    cfg["optimizer"] = o(cfg["model"].parameters(), lr=p["lr"])
+    return ModelConfig(**cfg)

lynxkite-graph-analytics/tests/test_pytorch_model_ops.py

@@ -0,0 +1,65 @@
+from lynxkite.core import workspace
+from lynxkite_graph_analytics import pytorch_model_ops
+import torch
+import pytest
+
+
+def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str, str, str]]):
+    ws = workspace.Workspace(env=env)
+    for id, data in nodes.items():
+        ws.nodes.append(
+            workspace.WorkspaceNode(
+                id=id,
+                type="basic",
+                data=workspace.WorkspaceNodeData(title=data["title"], params=data),
+                position=workspace.Position(
+                    x=data.get("x", 0),
+                    y=data.get("y", 0),
+                ),
+            )
+        )
+    ws.edges = [
+        workspace.WorkspaceEdge(
+            id=f"{source}->{target}",
+            source=source.split(":")[0],
+            target=target.split(":")[0],
+            sourceHandle=source.split(":")[1],
+            targetHandle=target.split(":")[1],
+        )
+        for source, target in edges
+    ]
+    return ws
+
+
+async def test_build_model():
+    ws = make_ws(
+        pytorch_model_ops.ENV,
+        {
+            "emb": {"title": "Input: embedding"},
+            "lin": {"title": "Linear", "output_dim": "same"},
+            "act": {"title": "Activation", "type": "Leaky ReLU"},
+            "label": {"title": "Input: label"},
+            "loss": {"title": "MSE loss"},
+            "optim": {"title": "Optimizer", "type": "SGD", "lr": 0.1},
+        },
+        [
+            ("emb:x", "lin:x"),
+            ("lin:x", "act:x"),
+            ("act:x", "loss:x"),
+            ("label:y", "loss:y"),
+            ("loss:loss", "optim:loss"),
+        ],
+    )
+    x = torch.rand(100, 4)
+    y = x + 1
+    m = pytorch_model_ops.build_model(ws, {"emb_x": x, "label_y": y})
+    for i in range(1000):
+        loss = m.train({"emb_x": x, "label_y": y})
+    assert loss < 0.1
+    o = m.inference({"emb_x": x[:1]})
+    error = torch.nn.functional.mse_loss(o["act_x"], x[:1] + 1)
+    assert error < 0.1
+
+
+if __name__ == "__main__":
+    pytest.main()