rename

Author: Z80coder

neurallogic/hard_bit.py

@@ -0,0 +1,31 @@
+import jax
+
+
+def soft_real_encoder(t: float, x: float) -> float:
+    # x should be in [0, 1]
+    t = jax.numpy.clip(t, 0.0, 1.0)
+    return jax.numpy.where(
+        x == t,
+        0.5,
+        jax.numpy.where(
+            x < t,
+            (1.0 / (2.0 * t)) * x,
+            (1.0 / (2.0 * (1.0 - t))) * (x + 1.0 - 2.0 * t),
+        ),
+    )
+
+
+def hard_real_encoder(t: float, x: float) -> bool:
+    # t and x must be floats
+    return jax.numpy.where(soft_real_encoder(t, x) > 0.5, True, False)
+
+
+soft_real_encoder_neuron = jax.vmap(soft_real_encoder, in_axes=(0, None))
+
+hard_real_encoder_neuron = jax.vmap(hard_real_encoder, in_axes=(0, None))
+
+soft_real_encoder_layer = jax.vmap(soft_real_encoder_neuron, (0, 0), 0)
+
+hard_real_encoder_layer = jax.vmap(hard_real_encoder_neuron, (0, 0), 0)
+
+

neurallogic/real_encoder.py

@@ -9,8 +9,7 @@
 from flax.metrics import tensorboard
 from flax.training import train_state
 import ml_collections
-from neurallogic import (hard_not, hard_or, harden, harden_layer,
-                         neural_logic_net)
+from neurallogic import hard_not, hard_or, harden, harden_layer, neural_logic_net
 import optax
 
 
@@ -23,9 +22,8 @@
 
 
 def nln(type, x, width):
-    x = hard_or.or_layer(type)(width, nn.initializers.uniform(
-        1.0), dtype=jnp.float32)(x)  # >=1700 need for >98% accuracy
-    x = hard_not.not_layer(type)(10, dtype=jnp.float32)(x)
+    x = hard_or.or_layer(type)(width, nn.initializers.uniform(1.0))(x)
+    x = hard_not.not_layer(type)(10)(x)
     x = x.ravel()  # flatten the outputs of the not layer
     # harden the outputs of the not layer
     x = harden_layer.harden_layer(type)(x)
@@ -59,11 +57,11 @@ def __call__(self, x):
 @jax.jit
 def apply_model_with_grad(state, images, labels):
     """Computes gradients, loss and accuracy for a single batch."""
+
     def loss_fn(params):
-        logits = state.apply_fn({'params': params}, images)
+        logits = state.apply_fn({"params": params}, images)
         one_hot = jax.nn.one_hot(labels, 10)
-        loss = jnp.mean(optax.softmax_cross_entropy(
-            logits=logits, labels=one_hot))
+        loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
         return loss, logits
 
     grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
@@ -79,21 +77,20 @@ def update_model(state, grads):
 
 def train_epoch(state, train_ds, batch_size, rng):
     """Train for a single epoch."""
-    train_ds_size = len(train_ds['image'])
+    train_ds_size = len(train_ds["image"])
     steps_per_epoch = train_ds_size // batch_size
 
-    perms = jax.random.permutation(rng, len(train_ds['image']))
-    perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
+    perms = jax.random.permutation(rng, len(train_ds["image"]))
+    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_images = train_ds['image'][perm, ...]
-        batch_labels = train_ds['label'][perm, ...]
-        grads, loss, accuracy = apply_model_with_grad(
-            state, batch_images, batch_labels)
+        batch_images = train_ds["image"][perm, ...]
+        batch_labels = train_ds["label"][perm, ...]
+        grads, loss, accuracy = apply_model_with_grad(state, batch_images, batch_labels)
         state = update_model(state, grads)
         epoch_loss.append(loss)
         epoch_accuracy.append(accuracy)
@@ -103,24 +100,23 @@ def train_epoch(state, train_ds, batch_size, rng):
 
 
 def get_datasets():
-    ds_builder = tfds.builder('mnist')
+    ds_builder = tfds.builder("mnist")
     ds_builder.download_and_prepare()
-    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))
-    train_ds['image'] = jnp.float32(train_ds['image']) / 255.
-    test_ds['image'] = jnp.float32(test_ds['image']) / 255.
+    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))
+    train_ds["image"] = jnp.float32(train_ds["image"]) / 255.0
+    test_ds["image"] = jnp.float32(test_ds["image"]) / 255.0
     # 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
 
 
 def show_img(img, ax=None, title=None):
     """Shows a single image."""
     if ax is None:
         ax = plt.gca()
-    ax.imshow(img.reshape(28, 28), cmap='gray')
+    ax.imshow(img.reshape(28, 28), cmap="gray")
     ax.set_xticks([])
     ax.set_yticks([])
     if title:
@@ -129,7 +125,7 @@ def show_img(img, ax=None, title=None):
 
 def show_img_grid(imgs, titles):
     """Shows a grid of images."""
-    n = int(np.ceil(len(imgs)**.5))
+    n = int(np.ceil(len(imgs) ** 0.5))
     _, axs = plt.subplots(n, n, figsize=(3 * n, 3 * n))
     for i, (img, title) in enumerate(zip(imgs, titles)):
         show_img(img, axs[i // n][i % n], title)
@@ -141,13 +137,14 @@ def create_train_state(net, rng, config):
     # mock_input = jnp.ones([1, 28, 28, 1])
     # for NLN
     mock_input = jnp.ones([1, 28 * 28])
-    soft_weights = net.init(rng, mock_input)['params']
+    soft_weights = net.init(rng, mock_input)["params"]
     tx = optax.sgd(config.learning_rate, config.momentum)
     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:
+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.
@@ -166,21 +163,22 @@ def train_and_evaluate(net, datasets, config: ml_collections.ConfigDict,
 
     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)
-        _, test_loss, test_accuracy = apply_model_with_grad(state, test_ds['image'],
-                                                            test_ds['label'])
+        state, train_loss, train_accuracy = train_epoch(
+            state, train_ds, config.batch_size, input_rng
+        )
+        _, test_loss, test_accuracy = apply_model_with_grad(
+            state, test_ds["image"], test_ds["label"]
+        )
 
         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))
+            "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)
+        )
 
-        summary_writer.scalar('train_loss', train_loss, epoch)
-        summary_writer.scalar('train_accuracy', train_accuracy, epoch)
-        summary_writer.scalar('test_loss', test_loss, epoch)
-        summary_writer.scalar('test_accuracy', test_accuracy, epoch)
+        summary_writer.scalar("train_loss", train_loss, epoch)
+        summary_writer.scalar("train_accuracy", train_accuracy, epoch)
+        summary_writer.scalar("test_loss", test_loss, epoch)
+        summary_writer.scalar("test_accuracy", test_accuracy, epoch)
 
     return state
 
@@ -197,13 +195,13 @@ def get_config():
     # 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 = 1000
     return config
 
 
 def apply_hard_model(state, image, label):
     def logits_fn(params):
-        return state.apply_fn({'params': params}, image)
+        return state.apply_fn({"params": params}, image)
 
     logits = logits_fn(state.params)
     if isinstance(logits, list):
@@ -224,33 +222,41 @@ def check_symbolic(nets, datasets, trained_state):
     _, test_ds = datasets
     _, hard, symbolic = nets
     _, test_loss, test_accuracy = apply_model_with_grad(
-        trained_state, test_ds['image'], test_ds['label'])
-    print('soft_net: final test_loss: %.4f, final test_accuracy: %.2f' %
-          (test_loss, test_accuracy * 100))
+        trained_state, test_ds["image"], test_ds["label"]
+    )
+    print(
+        "soft_net: final test_loss: %.4f, final test_accuracy: %.2f"
+        % (test_loss, test_accuracy * 100)
+    )
     hard_weights = harden.hard_weights(trained_state.params)
     hard_trained_state = train_state.TrainState.create(
-        apply_fn=hard.apply, params=hard_weights, tx=optax.sgd(1.0, 1.0))
-    hard_input = harden.harden(test_ds['image'])
+        apply_fn=hard.apply, params=hard_weights, tx=optax.sgd(1.0, 1.0)
+    )
+    hard_input = harden.harden(test_ds["image"])
     hard_test_accuracy = apply_hard_model_to_images(
-        hard_trained_state, hard_input, test_ds['label'])
-    print('hard_net: final test_accuracy: %.2f' % (hard_test_accuracy * 100))
+        hard_trained_state, hard_input, test_ds["label"]
+    )
+    print("hard_net: final test_accuracy: %.2f" % (hard_test_accuracy * 100))
     assert np.isclose(test_accuracy, hard_test_accuracy, atol=0.0001)
+    # TODO: activate these checks
     if False:
         # It takes too long to compute this
         symbolic_weights = harden.symbolic_weights(trained_state.params)
         symbolic_trained_state = train_state.TrainState.create(
-            apply_fn=symbolic.apply, params=symbolic_weights, tx=optax.sgd(1.0, 1.0))
+            apply_fn=symbolic.apply, params=symbolic_weights, tx=optax.sgd(1.0, 1.0)
+        )
         symbolic_input = hard_input.tolist()
         symbolic_test_accuracy = apply_hard_model_to_images(
-            symbolic_trained_state, symbolic_input, test_ds['label'])
-        print('symbolic_net: final test_accuracy: %.2f' %
-              (symbolic_test_accuracy * 100))
-        assert (np.isclose(test_accuracy, symbolic_test_accuracy, atol=0.0001))
+            symbolic_trained_state, symbolic_input, test_ds["label"]
+        )
+        print(
+            "symbolic_net: final test_accuracy: %.2f" % (symbolic_test_accuracy * 100)
+        )
+        assert np.isclose(test_accuracy, symbolic_test_accuracy, atol=0.0001)
     if False:
         # CPU and GPU give different results, so we can't easily regress on a static symbolic expression
         symbolic_input = [f"x{i}" for i in range(len(hard_input[0].tolist()))]
-        symbolic_output = symbolic.apply(
-            {'params': symbolic_weights}, symbolic_input)
+        symbolic_output = symbolic.apply({"params": symbolic_weights}, symbolic_input)
         print("symbolic_output", symbolic_output[0][:10000])
 
 
@@ -263,23 +269,20 @@ def test_mnist():
 
     # Define the model.
     # soft = CNN()
-    width = 100
-    soft, _, _ = neural_logic_net.net(
-        lambda type, x: batch_nln(type, x, width))
+    width = 1000
+    soft, _, _ = neural_logic_net.net(lambda type, x: batch_nln(type, x, width))
 
     # Get the MNIST dataset.
     train_ds, test_ds = get_datasets()
     # If we're using a NLN then flatten the images
-    train_ds["image"] = jnp.reshape(
-        train_ds["image"], (train_ds["image"].shape[0], -1))
-    test_ds["image"] = jnp.reshape(
-        test_ds["image"], (test_ds["image"].shape[0], -1))
+    train_ds["image"] = jnp.reshape(train_ds["image"], (train_ds["image"].shape[0], -1))
+    test_ds["image"] = jnp.reshape(test_ds["image"], (test_ds["image"].shape[0], -1))
 
     # Train and evaluate the model.
     trained_state = train_and_evaluate(
-        soft, (train_ds, test_ds), config=config, workdir="./mnist_metrics")
+        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, width))
+    _, hard, symbolic = neural_logic_net.net(lambda type, x: nln(type, x, width))
     check_symbolic((soft, hard, symbolic), (train_ds, test_ds), trained_state)