From 5ff90e94c9bb411a0262a8130a6f0ce4125ca11b Mon Sep 17 00:00:00 2001 From: ericmarin Date: Mon, 16 Mar 2026 19:36:31 +0100 Subject: changed torch.fx to ONNX --- xor.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) (limited to 'xor.py') diff --git a/xor.py b/xor.py index 9ab7be7..0f8390d 100644 --- a/xor.py +++ b/xor.py @@ -1,9 +1,10 @@ import torch import torch.nn as nn +import torch.onnx import nneq class xor_mlp(nn.Module): - def __init__(self, hidden_dim=4): + def __init__(self, hidden_dim=8): super().__init__() self.layers = nn.Sequential( nn.Linear(2, hidden_dim), @@ -13,13 +14,13 @@ class xor_mlp(nn.Module): def forward(self, x): return self.layers(x) -def train_model(name: str): +def train_model(name: str, dim): X = torch.tensor([[0,0], [0,1], [1,0], [1,1]], dtype=torch.float32) Y = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32) - net = xor_mlp() + net = xor_mlp(hidden_dim=dim) loss_fn = nn.MSELoss() - optimizer = torch.optim.Adam(net.parameters(), lr=0.01) + optimizer = torch.optim.Adam(net.parameters(), lr=0.1) print(f"Training {name}...") for epoch in range(1000): @@ -33,16 +34,18 @@ def train_model(name: str): return net if __name__ == "__main__": - net_a = train_model("Network A") - net_b = train_model("Network B") + torch_net_a = train_model("Network A", 8).eval() + torch_net_b = train_model("Network B", 16).eval() + + onnx_net_a = torch.onnx.export(torch_net_a, (torch.randn(1, 2),), verbose=False, dynamo=True).model_proto # type: ignore + onnx_net_b = torch.onnx.export(torch_net_b, (torch.randn(1, 2),), verbose=False, dynamo=True).model_proto # type: ignore + + z3_net_a = nneq.net(onnx_net_a) + z3_net_b = nneq.net(onnx_net_b) - z3_net_a = nneq.net(net_a, (2,)) - z3_net_b = nneq.net(net_b, (2,)) - print("") nneq.strict_equivalence(z3_net_a, z3_net_b) print("") nneq.epsilon_equivalence(z3_net_a, z3_net_b, 0.1) print("") nneq.argmax_equivalence(z3_net_a, z3_net_b) - -- cgit v1.2.3