Spaces:
Running
Running
File size: 1,996 Bytes
3cc3a0a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
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()
|