<jupyter_start><jupyter_text>Setup<jupyter_code>!git clone https://github.com/yeyinthtoon/tf2-resmlp.git
%cd tf2-resmlp/
!pip install .
import numpy as np
import tensorflow as tf
from tensorflow import keras
from resmlp import ResMlp12<jupyter_output><empty_output><jupyter_text>TPU setup<jupyter_code>try:
    tpu = (
        tf.distribute.cluster_resolver.TPUClusterResolver()
    )  # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.
    print("Running on TPU ", tpu.master())
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
    strategy = tf.distribute.get_strategy()
print("REPLICAS: ", strategy.num_replicas_in_sync)<jupyter_output>Running on TPU  grpc://10.49.133.194:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.49.133.194:8470<jupyter_text>Configs<jupyter_code>configs = {
    "seed": 112,
    "data_augmentation": True,
    "epochs": 200,
    "batch_size": 256,
    "drop_path_rate": 0.3,
    "dropout_rate": 0.5,
    "num_classes": 10,
    "input_image_size": [224, 224],
}<jupyter_output><empty_output><jupyter_text>Dataset Preparation<jupyter_code>(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255

x_train_mean = np.mean(x_train, axis=0)
x_train -= x_train_mean
x_test -= x_train_mean
def augment(image, label):
    image = tf.image.random_flip_left_right(image)
    return image, label


def set_shape(image, label):
    image = tf.image.resize(image, configs["input_image_size"])
    image.set_shape(configs["input_image_size"] + [3])
    label.set_shape(
        [
            1,
        ]
    )
    return image, label


train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.shuffle(10000, seed=configs["seed"])
if configs["data_augmentation"]:
    train_ds = train_ds.map(augment)
train_ds = train_ds.map(set_shape)
train_ds = train_ds.batch(configs["batch_size"], drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)


test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.map(set_shape)
test_ds = test_ds.batch(configs["batch_size"], drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>tf.keras.backend.clear_session()
with strategy.scope():
    model = ResMlp12(
        input_shape=configs["input_image_size"] + [3],
        num_classes=configs["num_classes"],
        dropout_rate=configs["dropout_rate"],
        drop_path_rate=configs["drop_path_rate"],
    )
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        optimizer="adam",
        metrics=["accuracy"],
    )
model.summary()
def lr_schedule(epoch):
    """Learning Rate Schedule

    Learning rate is scheduled to be reduced after 80, 120, 160, 180 epochs.
    Called automatically every epoch as part of callbacks during training.

    # Arguments
        epoch (int): The number of epochs

    # Returns
        lr (float32): learning rate
    """
    lr = 1e-3
    if epoch > 180:
        lr *= 0.5e-3
    elif epoch > 160:
        lr *= 1e-3
    elif epoch > 120:
        lr *= 1e-2
    elif epoch > 80:
        lr *= 1e-1
    print("Learning rate: ", lr)
    return lr


lr_scheduler = keras.callbacks.LearningRateScheduler(lr_schedule)
history = model.fit(
    train_ds,
    epochs=configs["epochs"],
    validation_data=(test_ds),
    callbacks=[lr_scheduler],
)<jupyter_output>Epoch 1/200
Learning rate:  0.001
195/195 [==============================] - 59s 120ms/step - loss: 1.8074 - accuracy: 0.3406 - val_loss: 1.4649 - val_accuracy: 0.4724
Epoch 2/200
Learning rate:  0.001
195/195 [==============================] - 17s 88ms/step - loss: 1.4310 - accuracy: 0.4870 - val_loss: 1.2728 - val_accuracy: 0.5303
Epoch 3/200
Learning rate:  0.001
195/195 [==============================] - 17s 88ms/step - loss: 1.2669 - accuracy: 0.5476 - val_loss: 1.1454 - val_accuracy: 0.5817
Epoch 4/200
Learning rate:  0.001
195/195 [==============================] - 17s 88ms/step - loss: 1.1619 - accuracy: 0.5846 - val_loss: 1.0282 - val_accuracy: 0.6346
Epoch 5/200
Learning rate:  0.001
195/195 [==============================] - 17s 88ms/step - loss: 1.0979 - accuracy: 0.6078 - val_loss: 1.0055 - val_accuracy: 0.6386
Epoch 6/200
Learning rate:  0.001
195/195 [==============================] - 17s 89ms/step - loss: 1.0430 - accuracy: 0.6290 - val_loss: 0.9465 - val_accuracy: 0.65[...]
{'fraction_non_alphanumeric': 0.16052741732942655, 'fraction_numerical': 0.07115948095437422, 'mean_word_length': 4.004188481675393, 'pattern_counts': {'":': 8, '<': 19, '<?xml version=': 0, '>': 23, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '336558', 'license': 'MIT', 'n_tokens_mistral': 2018, 'n_tokens_neox': 1734, 'n_words': 338, 'path': 'examples/cifar10.ipynb', 'repo_name': 'yeyinthtoon/tf2-resmlp', 'size': '4778'}