100

Author: Z80coder

neurallogic/hard_xor.py

@@ -3,7 +3,7 @@
 import jax
 from flax import linen as nn
 
-from neurallogic import neural_logic_net, symbolic_generation
+from neurallogic import neural_logic_net, symbolic_generation, hard_and
 
 
 def soft_xor_include(w: float, x: float) -> float:
@@ -28,6 +28,7 @@ def soft_xor_neuron(w, x):
 
     def xor(x, y):
         return jax.numpy.minimum(jax.numpy.maximum(x, y), 1.0 - jax.numpy.minimum(x, y))
+
     x = jax.lax.reduce(x, jax.numpy.array(0, dtype=x.dtype), xor, (0,))
     return x
 
@@ -45,15 +46,19 @@ def hard_xor_neuron(w, x):
 
 class SoftXorLayer(nn.Module):
     layer_size: int
-    weights_init: Callable = nn.initializers.uniform(1.0)
+    weights_init: Callable = (
+        nn.initializers.uniform(1.0)
+        #hard_and.initialize_near_to_zero()
+    )  
     dtype: jax.numpy.dtype = jax.numpy.float32
 
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
         weights = self.param(
-            'bit_weights', self.weights_init, weights_shape, self.dtype)
-        x = jax.numpy.asarray(x, self.dtype) 
+            "bit_weights", self.weights_init, weights_shape, self.dtype
+        )
+        x = jax.numpy.asarray(x, self.dtype)
         return soft_xor_layer(weights, x)
 
 
@@ -64,7 +69,8 @@ class HardXorLayer(nn.Module):
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
         weights = self.param(
-            'bit_weights', nn.initializers.constant(True), weights_shape)
+            "bit_weights", nn.initializers.constant(True), weights_shape
+        )
         return hard_xor_layer(weights, x)
 
 
@@ -74,14 +80,18 @@ def __init__(self, layer_size):
         self.hard_xor_layer = HardXorLayer(self.layer_size)
 
     def __call__(self, x):
-        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(
-            self.hard_xor_layer, x)
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(self.hard_xor_layer, x)
         return symbolic_generation.symbolic_expression(jaxpr, x)
 
 
 xor_layer = neural_logic_net.select(
     lambda layer_size, weights_init=nn.initializers.uniform(
-        1.0), dtype=jax.numpy.float32: SoftXorLayer(layer_size, weights_init, dtype),
+        1.0
+    ), dtype=jax.numpy.float32: SoftXorLayer(layer_size, weights_init, dtype),
+    lambda layer_size, weights_init=nn.initializers.constant(
+        True
+    ), dtype=jax.numpy.float32: HardXorLayer(layer_size),
     lambda layer_size, weights_init=nn.initializers.constant(
-        True), dtype=jax.numpy.float32: HardXorLayer(layer_size),
-    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: SymbolicXorLayer(layer_size))
+        True
+    ), dtype=jax.numpy.float32: SymbolicXorLayer(layer_size),
+)

tests/test_mnist.py

@@ -13,11 +13,20 @@
 from matplotlib import pyplot as plt
 from tqdm import tqdm
 
-from neurallogic import (hard_and, hard_majority, hard_not, hard_or, hard_xor, harden,
-                         harden_layer, neural_logic_net, real_encoder)
+from neurallogic import (
+    hard_and,
+    hard_majority,
+    hard_not,
+    hard_or,
+    hard_xor,
+    harden,
+    harden_layer,
+    neural_logic_net,
+    real_encoder,
+)
 
 # Uncomment to debug NaNs
-#config.update("jax_debug_nans", True)
+# config.update("jax_debug_nans", True)
 
 """
 MNIST test.
@@ -44,15 +53,18 @@ def nln(type, x, width):
     return x
 """
 
+
 def nln(type, x):
     num_classes = 10
 
-    x = hard_or.or_layer(type)(1800, nn.initializers.uniform(1.0), dtype=jax.numpy.float16)(x)
+    x = hard_or.or_layer(type)(
+        1800, nn.initializers.uniform(1.0), dtype=jax.numpy.float16
+    )(x)
     x = hard_not.not_layer(type)(1, dtype=jax.numpy.float16)(x)
     x = x.ravel()
-    x = harden_layer.harden_layer(type)(x) 
-    x = x.reshape((num_classes, int(x.shape[0] / num_classes))) 
-    x = x.sum(-1) 
+    x = harden_layer.harden_layer(type)(x)
+    x = x.reshape((num_classes, int(x.shape[0] / num_classes)))
+    x = x.sum(-1)
     return x
 
 
@@ -129,13 +141,13 @@ def get_datasets():
     train_ds = tfds.as_numpy(ds_builder.as_dataset(split="train", batch_size=-1))
     test_ds = tfds.as_numpy(ds_builder.as_dataset(split="test", batch_size=-1))
     # XXXX
-    train_ds["image"] = (jnp.float32(train_ds["image"]) / 255.0)
-    test_ds["image"] = (jnp.float32(test_ds["image"]) / 255.0)
+    train_ds["image"] = jnp.float32(train_ds["image"]) / 255.0
+    test_ds["image"] = jnp.float32(test_ds["image"]) / 255.0
     # TODO: we don't need to do this even when we don't use the real encoder
     # Use grayscale information
     # Convert the floating point values in [0,1] to binary values in {0,1}
-    #train_ds["image"] = jnp.round(train_ds["image"])
-    #test_ds["image"] = jnp.round(test_ds["image"])
+    # train_ds["image"] = jnp.round(train_ds["image"])
+    # test_ds["image"] = jnp.round(test_ds["image"])
     return train_ds, test_ds
 
 
@@ -165,23 +177,16 @@ def create_train_state(net, rng, config):
     # for NLN
     mock_input = jnp.ones([1, 28 * 28])
     soft_weights = net.init(rng, mock_input)["params"]
-    #tx = optax.sgd(config.learning_rate, config.momentum)
-    #tx = optax.noisy_sgd(config.learning_rate, config.momentum)
+    # tx = optax.sgd(config.learning_rate, config.momentum)
+    # tx = optax.noisy_sgd(config.learning_rate, config.momentum)
     tx = optax.yogi(config.learning_rate)
     return train_state.TrainState.create(apply_fn=net.apply, params=soft_weights, tx=tx)
 
 
 def train_and_evaluate(
     net, datasets, config: ml_collections.ConfigDict, workdir: str
 ) -> train_state.TrainState:
-    """Execute model training and evaluation loop.
-    Args:
-      config: Hyperparameter configuration for training and evaluation.
-      workdir: Directory where the tensorboard summaries are written to.
-    Returns:
-      The train state (which includes the `.params`).
-    """
-    train_ds, test_ds = datasets
+    train_dataset, test_dataset = datasets
     rng = jax.random.PRNGKey(0)
 
     summary_writer = tensorboard.SummaryWriter(workdir)
@@ -193,10 +198,10 @@ def train_and_evaluate(
     for epoch in range(1, config.num_epochs + 1):
         rng, input_rng = jax.random.split(rng)
         state, train_loss, train_accuracy = train_epoch(
-            state, train_ds, config.batch_size, input_rng
+            state, train_dataset, config.batch_size, input_rng
         )
         _, test_loss, test_accuracy = apply_model_with_grad(
-            state, test_ds["image"], test_ds["label"]
+            state, test_dataset["image"], test_dataset["label"]
         )
 
         print(
@@ -219,13 +224,13 @@ def get_config():
     # config for CNN
     config.learning_rate = 0.01
     # config for NLN
-    #config.learning_rate = 0.1
+    # config.learning_rate = 0.1
     config.learning_rate = 0.01
 
     # Always commit with num_epochs = 1 for short test time
     config.momentum = 0.9
     config.batch_size = 128
-    #config.num_epochs = 2
+    # config.num_epochs = 2
     config.num_epochs = 1000
     return config
 
@@ -290,6 +295,7 @@ def check_symbolic(nets, datasets, trained_state):
         symbolic_output = symbolic.apply({"params": symbolic_weights}, symbolic_input)
         print("symbolic_output", symbolic_output[0][:10000])
 
+
 @pytest.mark.skip(reason="temporarily off")
 def test_mnist():
     # Make sure tf does not allocate gpu memory.
@@ -311,13 +317,13 @@ def test_mnist():
 
     print(soft.tabulate(jax.random.PRNGKey(0), train_ds["image"][0:1]))
     # TODO: fix the size of this
-    #print(hard.tabulate(jax.random.PRNGKey(0), harden.harden(train_ds["image"][0:1])))
+    # print(hard.tabulate(jax.random.PRNGKey(0), harden.harden(train_ds["image"][0:1])))
 
     # Train and evaluate the model.
     trained_state = train_and_evaluate(
         soft, (train_ds, test_ds), config=config, workdir="./mnist_metrics"
     )
 
     # Check symbolic net
-    #_, hard, symbolic = neural_logic_net.net(lambda type, x: nln(type, x))
-    #check_symbolic((soft, hard, symbolic), (train_ds, test_ds), trained_state)
+    # _, hard, symbolic = neural_logic_net.net(lambda type, x: nln(type, x))
+    # check_symbolic((soft, hard, symbolic), (train_ds, test_ds), trained_state)

tests/test_noisy_xor.py

@@ -0,0 +1,163 @@
+from pathlib import Path
+import ml_collections
+import numpy
+import optax
+from flax.training import train_state
+from flax import linen as nn
+import jax
+
+from neurallogic import (
+    neural_logic_net,
+    hard_not,
+    hard_or,
+    hard_and,
+    hard_xor,
+    hard_majority,
+    harden_layer,
+)
+
+num_features = 12
+num_classes = 2
+
+
+def get_data():
+    # Create a path to the data directory
+    data_dir = Path(__file__).parent.parent / "tests" / "data"
+    # Load the training data
+    training_data = numpy.loadtxt(data_dir / "NoisyXORTrainingData.txt").astype(
+        dtype=numpy.int32
+    )
+    # Load the test data
+    test_data = numpy.loadtxt(data_dir / "NoisyXORTestData.txt").astype(
+        dtype=numpy.int32
+    )
+    return training_data, test_data
+
+
+# 89% test accuracy
+def nln_89(type, x):
+    x = hard_and.and_layer(type)(20)(x)
+    x = hard_not.not_layer(type)(5)(x)
+    x = x.ravel()
+    ########################################################
+    x = harden_layer.harden_layer(type)(x)
+    x = x.reshape((num_classes, int(x.shape[0] / num_classes)))
+    x = x.sum(-1)
+    return x
+
+
+# 100% test accuracy
+def nln(type, x):
+    x = hard_and.and_layer(type)(20)(x)
+    x = hard_not.not_layer(type)(4)(x)
+    x = x.ravel()
+    ########################################################
+    x = harden_layer.harden_layer(type)(x)
+    x = x.reshape((num_classes, int(x.shape[0] / num_classes)))
+    x = x.sum(-1)
+    return x
+
+
+def batch_nln(type, x):
+    return jax.vmap(lambda x: nln(type, x))(x)
+
+
+def create_train_state(net, rng, config):
+    mock_input = jax.numpy.ones([1, num_features])
+    soft_weights = net.init(rng, mock_input)["params"]
+    # tx = optax.sgd(config.learning_rate, config.momentum)
+    tx = optax.yogi(config.learning_rate)
+    return train_state.TrainState.create(apply_fn=net.apply, params=soft_weights, tx=tx)
+
+
+@jax.jit
+def update_model(state, grads):
+    return state.apply_gradients(grads=grads)
+
+
+@jax.jit
+def apply_model_with_grad(state, features, labels):
+    def loss_fn(params):
+        logits = state.apply_fn({"params": params}, features)
+        one_hot = jax.nn.one_hot(labels, num_classes)
+        loss = jax.numpy.mean(
+            optax.softmax_cross_entropy(logits=logits, labels=one_hot)
+        )
+        return loss, logits
+
+    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    (loss, logits), grads = grad_fn(state.params)
+    accuracy = jax.numpy.mean(jax.numpy.argmax(logits, -1) == labels)
+    return grads, loss, accuracy
+
+
+def train_epoch(state, features, labels, batch_size, rng):
+    train_ds_size = len(features)
+    steps_per_epoch = train_ds_size // batch_size
+
+    perms = jax.random.permutation(rng, len(features))
+    perms = perms[: steps_per_epoch * batch_size]  # skip incomplete batch
+    perms = perms.reshape((steps_per_epoch, batch_size))
+
+    epoch_loss = []
+    epoch_accuracy = []
+
+    for perm in perms:
+        batch_features = features[perm, ...]
+        batch_labels = labels[perm, ...]
+        grads, loss, accuracy = apply_model_with_grad(
+            state, batch_features, batch_labels
+        )
+        state = update_model(state, grads)
+        epoch_loss.append(loss)
+        epoch_accuracy.append(accuracy)
+    train_loss = numpy.mean(epoch_loss)
+    train_accuracy = numpy.mean(epoch_accuracy)
+    return state, train_loss, train_accuracy
+
+
+def train_and_evaluate(net, datasets, config: ml_collections.ConfigDict):
+    training_data, test_data = datasets
+    x_training = training_data[:, 0:num_features]  # Input features
+    y_training = training_data[:, num_features]  # Target value
+    x_test = test_data[:, 0:num_features]  # Input features
+    y_test = test_data[:, num_features]  # Target value
+
+    rng = jax.random.PRNGKey(0)
+    print(net.tabulate(rng, x_training[0:1]))
+
+    rng, init_rng = jax.random.split(rng)
+    state = create_train_state(net, init_rng, config)
+
+    best_test_accuracy = 0.0
+    for epoch in range(1, config.num_epochs + 1):
+        rng, input_rng = jax.random.split(rng)
+        state, train_loss, train_accuracy = train_epoch(
+            state, x_training, y_training, config.batch_size, input_rng
+        )
+        _, test_loss, test_accuracy = apply_model_with_grad(state, x_test, y_test)
+        if test_accuracy > best_test_accuracy:
+            best_test_accuracy = test_accuracy
+
+        print(
+            "epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, test_loss: %.4f, test_accuracy: %.2f"
+            % (epoch, train_loss, train_accuracy * 100, test_loss, test_accuracy * 100)
+        )
+        print(f"best_test_accuracy: {best_test_accuracy * 100:.2f}")
+
+    return state
+
+
+def get_config():
+    config = ml_collections.ConfigDict()
+    config.learning_rate = 0.01
+    config.momentum = 0.9
+    config.batch_size = 256
+    config.num_epochs = 1000
+    return config
+
+
+def test_noisy_xor():
+    soft, hard, _ = neural_logic_net.net(lambda type, x: batch_nln(type, x))
+    training_data, test_data = get_data()
+    trained_state = train_and_evaluate(soft, (training_data, test_data), get_config())