Remove scripts

Author: daniel-gallo

common.py

@@ -1,6 +1,4 @@
-import jax
 import jax.numpy as jnp
-import matplotlib.pyplot as plt
 import optax
 
 
@@ -51,62 +49,3 @@ def schedule_fn(step):
         return lr
 
     return schedule_fn
-
-
-if __name__ == "__main__":
-    ### Testing the lr_scheduler
-    num_iters = 50_000
-    warmup_steps = 500
-    cooldown_steps = 1_000
-
-    decay_steps = num_iters - warmup_steps - cooldown_steps
-
-    base_lr = 1e-3
-    clip_value = 1.0
-    decay_rate = 0.1
-
-    lr_scheduler = warmup_exponential_decay_cooldown_scheduler(
-        warmup_steps, base_lr, decay_steps, decay_rate, cooldown_steps, min_lr=0.0
-    )
-
-    ### Small training loop
-    params = jax.random.normal(jax.random.PRNGKey(0), (3, 3))
-
-    optimizer = optax.chain(
-        optax.clip_by_global_norm(clip_value),
-        optax.adamw(lr_scheduler, b2=0.98, weight_decay=0.01),
-    )
-
-    opt_state = optimizer.init(params)
-
-    @jax.jit
-    def update(params, opt_state):
-        loss, grads = jax.value_and_grad(lambda params: jnp.sum(params**2))(params)
-        updates, opt_state = optimizer.update(grads, opt_state, params)
-        params = optax.apply_updates(params, updates)
-        return params, opt_state, loss
-
-    num_steps = warmup_steps + decay_steps + cooldown_steps
-
-    steps = []
-    lrs = []
-    for step in range(num_steps + 1):
-        batch = jnp.ones((3, 3))  # Dummy batch data
-        params, opt_state, loss = update(params, opt_state)
-        current_lr = lr_scheduler(step)
-        if step % 1000 == 0:
-            print(f"Step {step} | Loss: {loss} | Learning Rate: {current_lr}")
-            steps.append(step)
-            lrs.append(current_lr)
-
-    ### Plot the LR
-    plt.figure(figsize=(10, 6))
-    plt.plot(steps, lrs, label="Learning rate")
-
-    plt.xlabel("Steps")
-    plt.ylabel("Learning Rate")
-    plt.title("Warmup, Exponential Decay, and Cooldown Learning Rate Schedule")
-    plt.legend()
-    plt.grid(True)
-    plt.savefig("lr_exponential.png")
-    plt.show()

dataset.py

@@ -1,14 +1,9 @@
 import collections
 import itertools
-import os
-from glob import glob
 
 import jax
 import tensorflow as tf
-from einops import rearrange
-from matplotlib import pyplot as plt
 
-from transformations.native_aspect_ratio_resize import NativeAspectRatioResize
 
 AUTOTUNE = tf.data.AUTOTUNE
 
@@ -137,34 +132,3 @@ def enqueue(n):
     while queue:
         yield queue.popleft()
         enqueue(1)
-
-
-if __name__ == "__main__":
-    tf.random.set_seed(0)
-    batch_size = 1
-    image_dir = "./tfrecords"
-    train_files = glob(os.path.join(image_dir, "*.tfrec"))
-    train_dataset = load_dataset(train_files, 14, [NativeAspectRatioResize(224, 14)])
-    train_ds = prefetch(
-        train_dataset.shuffle(10 * batch_size, seed=1)
-        .batch(batch_size)
-        .prefetch(tf.data.AUTOTUNE)
-        .repeat()
-        .as_numpy_iterator()
-    )
-
-    for batch in train_ds:
-        patches, patch_indices, label, attention_matrix, loss_mask = batch
-        h, w = 1 + patch_indices[0].max(axis=0)
-        print(h, w)
-        image = rearrange(
-            patches[0][: h * w],
-            "(h w) (p1 p2 c) -> (h p1) (w p2) c",
-            h=h,
-            w=w,
-            p1=14,
-            p2=14,
-            c=3,
-        )
-        plt.imshow(image)
-        plt.show()