code to train model

neural_models.py

@@ -0,0 +1,152 @@
+from datasets import load_dataset
+from transformers import AutoImageProcessor, create_optimizer, TFAutoModelForImageClassification, KerasMetricCallback, \
+    PushToHubCallback, pipeline
+import tensorflow as tf
+from tensorflow.python import keras
+from keras import layers, losses
+import numpy as np
+from PIL import Image
+from transformers import DefaultDataCollator
+import evaluate
+
+
+def convert_to_tf_tensor(image: Image):
+    np_image = np.array(image)
+    tf_image = tf.convert_to_tensor(np_image)
+
+    # `expand_dims()` is used to add a batch dimension since
+    # the TF augmentation layers operates on batched inputs.
+    return tf.expand_dims(tf_image, 0)
+
+
+def preprocess_train(example_batch):
+    """Apply train_transforms across a batch."""
+    images = [
+        train_data_augmentation(convert_to_tf_tensor(image.convert("RGB"))) for image in example_batch["image"]
+    ]
+    example_batch["pixel_values"] = [tf.transpose(tf.squeeze(image)) for image in images]
+    return example_batch
+
+
+def preprocess_val(example_batch):
+    """Apply val_transforms across a batch."""
+    images = [
+        val_data_augmentation(convert_to_tf_tensor(image.convert("RGB"))) for image in example_batch["image"]
+    ]
+    example_batch["pixel_values"] = [tf.transpose(tf.squeeze(image)) for image in images]
+    return example_batch
+
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    predictions = np.argmax(predictions, axis=1)
+    return accuracy.compute(predictions=predictions, references=labels)
+
+
+# load dataset
+fashion = load_dataset("fashion_mnist", split="train[:4000]")
+# Split into train/test sets
+fashion = fashion.train_test_split(test_size=0.2)
+# an example
+print(fashion["train"][0])
+
+# Map label names to an integer and vice-versa
+labels = fashion["train"].features["label"].names
+label2id, id2label = dict(), dict()
+for i, label in enumerate(labels):
+    label2id[label] = str(i)
+    id2label[str(i)] = label
+
+# Should convert label id into a name
+# print(label2id)
+# print(id2label)
+
+# Pre-processing with ViT
+# Load image processor to process image into tensor
+checkpoint = "google/vit-base-patch16-224-in21k"
+image_processor = AutoImageProcessor.from_pretrained(checkpoint)
+
+# To avoid overfitting and make the model more robust, add data augmentation to the training set.
+# User Keras preprocessing layers to define transformations for the training set.
+size = (image_processor.size["height"], image_processor.size["width"])
+
+train_data_augmentation = keras.Sequential(
+    [
+        layers.RandomCrop(size[0], size[1]),
+        layers.Rescaling(scale=1.0 / 127.5, offset=-1),
+        layers.RandomFlip("horizontal"),
+        layers.RandomRotation(factor=0.02),
+        layers.RandomZoom(height_factor=0.2, width_factor=0.2),
+    ],
+    name="train_data_augmentation",
+)
+
+val_data_augmentation = keras.Sequential(
+    [
+        layers.CenterCrop(size[0], size[1]),
+        layers.Rescaling(scale=1.0 / 127.5, offset=-1),
+    ],
+    name="val_data_augmentation",
+)
+
+fashion["train"].set_transform(preprocess_train)
+fashion["test"].set_transform(preprocess_val)
+
+data_collator = DefaultDataCollator(return_tensors="tf")
+
+accuracy = evaluate.load("accuracy")
+
+# Set hyperparameters
+batch_size = 16
+num_epochs = 4
+num_train_steps = len(fashion["train"]) * num_epochs
+learning_rate = 3e-5
+weight_decay_rate = 0.01
+
+# define optimizer, learning rate schedule
+optimizer, lr_schedule = create_optimizer(
+    init_lr=learning_rate,
+    num_train_steps=num_train_steps,
+    weight_decay_rate=weight_decay_rate,
+    num_warmup_steps=0,
+)
+
+# Load ViT along with label mappings
+model = TFAutoModelForImageClassification.from_pretrained(
+    checkpoint,
+    id2label=id2label,
+    label2id=label2id,
+)
+
+# converting datasets to tf.data.Dataset
+tf_train_dataset = fashion["train"].to_tf_dataset(
+    columns="pixel_values", label_cols="label", shuffle=True, batch_size=batch_size, collate_fn=data_collator
+)
+
+tf_eval_dataset = fashion["test"].to_tf_dataset(
+    columns="pixel_values", label_cols="label", shuffle=True, batch_size=batch_size, collate_fn=data_collator
+)
+
+# Configure model for training
+loss = losses.SparseCategoricalCrossentropy(from_logits=True)
+model.compile(optimizer=optimizer, loss=loss)
+
+
+metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_eval_dataset)
+push_to_hub_callback = PushToHubCallback(
+    output_dir="../fashion_classifier",
+    tokenizer=image_processor,
+    save_strategy="no",
+)
+
+callbacks = [metric_callback, push_to_hub_callback]
+
+model.fit(tf_train_dataset, validation_data=tf_eval_dataset, epochs=num_epochs, callbacks=callbacks)
+
+# model.push_to_hub()
+
+# ds = load_dataset("fashion_mnist", split="test[:10]")
+# image = ds["image"][0]
+# classifier = pipeline("image-classification", model="my_awesome_fashion_model")
+# print(classifier(image))
+