fashion_classifier / neural_models.py
arifsy's picture
code to train model
96ba5ec
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))