Building a Machine Learning Pipeline with TensorFlow Extended and W&B

Introduction

Experiment tracking is a very popular go-to toolbox among machine learning practitioners and organizations to better understand how the optimal combination of hyper-parameters is selected to fully train the final model. We'll explore how it can be integrated into TFX's Tuner and Trainer component.
A trained model in each pipeline run should be managed in a version control system, allowing us to easily compare different versions of the model and roll back whenever when need to. We'll see how to build a custom TFX component by extending the standard Pusher to save models to the W&B's model registry. Additionally, we'll showcase how to publish a Gradio application which automatically points the version of the model from the current run.

Overview of the TFX pipeline



Dataset: It's important to unify the data format inside any machine learning pipeline for consistency between components. Especially, binary and serializable format is preferred for better efficiency. Because TFX pipeline uses TFRecord as the standard data format, so the input data should be prepared as TFRecord beforehand. We provide a notebook to show how to convert the datasets from Hugging Face 🤗 Dataset Hub into TFRecord. Prepared TFRecord files can be easily consumed by ExampleGen component.
Transform/Shaping: The serialized binary dataset doesn't have information about the shape of the data needed in order to transform it. That information can be carried through SchemaGen component when Transform component tries to transform the input data. In this project, Transform component does two things: 1) pixel value normalization by dividing with 255.0, and 2) and transposing image Tensors from BCWH to BWHC(Batch size, Width, Height, Channel).
NOTE: Most of the TensorFlow models assume the input images are in the shape of BWHC while the datasets to be used in PyTorch in mind are often shaped in BCWH.
💡
Hyperparameter tuning: Hyperparameter tuning is a process that usually conducts and compares a number of experiments to find out an optimal combination of hyperparameters with a subset of the entire dataset. This is handled by the Tuner component which is simple a wrapper of KerasTuner. We provide a custom Python script to the Tuner component to tell which hyper-parameters to tune with what values. Tuner runs a number of experiments in parallel(if possible) by running the same script with different values. We just tell W&B to record an experiment has begun with what values.
Full training: Tuner emits the optimal combination of hyperparameters as the output, and it is passed down to Trainer to do full training. Here, full training means it uses the entire dataset because it is convinced to achieve the best model performance with the output of Tuner. Trainer saves a fully trained model in SavedModel format, and the location of it is passed to the downstream task. As you can see from the chart below, the full training (highlighted line) has achieved the highest validation accuracy and the lowest validation loss.

Pushing: Pusher, as you might guess, pushes a model to somewhere. By default, the target is the local file system, but it won't be manageable. Hence, we custom designed the WandBPusher component by extending the basic Pusher component. WandBPusher pushes the fully trained model to W&B's Model Registry, and it optionally publishes a demo application to Hugging Face 🤗 Space Hub. It's important to note that the model version is determined at runtime, so the it should be dynamically injected into the published application at runtime as well.

Demo application deployed on Hugging Face 🤗 Space Hub

Reusability



Experiment tracking (Tuner)

tuner = Tuner(
    tuner_fn="modules.tuning.tuner_fn",
    custom_config=TUNER_ARGS,
    ....
)

WANDB_CONFIGS = {
    "API_KEY": "WANDB_ACCESS_KEY",
    "PROJECT": PIPELINE_NAME,
}

HYPER_PARAMETERS = {
    ....
    "optimizer_type": {
        "type": "choice",
        "values": ["Adam", "AdamW"],
    },

    "learning_rate": {
        "type": "float",
        "min_value": 0.00001,
        "max_value": 0.1,
        "sampling": "log",
        "step": 10
    },
    ....
}

TUNER_CONFIGS = {
    "num_trials": 15
}

TUNER_ARGS = {
    ....
    "hyperparameters": HYPER_PARAMETERS,
    "tuner": TUNER_CONFIGS,
    "wandb": WANDB_CONFIGS
}

def tuner_fn(fn_args: FnArgs) -> TunerFnResult:
    wandb_configs = fn_args.custom_config["wandb"]
    wandb.login(key=wandb_configs["API_KEY"])

    tuner = MyTuner(
        wandb_configs["PROJECT"],
        MyHyperModel(),
        max_trials=fn_args.custom_config["tuner"],
        hyperparameters=get_hyperparameters(
            fn_args.custom_config["hyperparameters"]
	),
        objective=keras_tuner.Objective("val_accuracy", "max"),
	....
    )

    train_dataset = input_fn(fn_args.train_files, ...)
    eval_dataset = input_fn(fn_args.eval_files, ...)

    return TunerFnResult(
        tuner=tuner,
        fit_kwargs={
            "x": train_dataset,
            "validation_data": eval_dataset,
        },
    )

class MyTuner(keras_tuner.RandomSearch):
    def run_trial(self, trial, *args, **kwargs):
        hp = trial.hyperparameters
        model = self.hypermodel.build(hp)

        optimizer_type = hp.get("optimizer_type")
        learning_rate = hp.get("learning_rate")
        weight_decay = hp.get("weight_decay")
        epochs = hp.get("finetune_epochs")

        log_name = f"tuning@opt-{optimizer_type}@lr-{learning_rate}@wd-{weight_decay}"
        wandb.init(
            project=self.wandb_project, config=hp.values, name=log_name,
        )

        wandb.log({"optimizer_type": optimizer_type})
        wandb.log({"learning_rate": learning_rate})
        wandb.log({"weight_decay": weight_decay})

        result = self.hypermodel.fit(
            hp, model, epochs=epochs, 
            callbacks=[wandb.keras.WandbMetricsLogger(log_freq='epoch')], 
            *args, **kwargs
        )

        wandb.finish()
        return result

Experiment tracking (Trainer)


trainer = Trainer(
    "run_fn"="modules.train.run_fn",
    "hyperparameters"=tuner.outputs["best_hyperparameters"],
    "custom_config"=TRAINING_ARGS,
    ....
)

TRAINING_ARGS = {
...
    "wandb": WANDB_CONFIGS
}

def run_fn(fn_args: FnArgs):
    hp = keras_tuner.HyperParameters.from_config(fn_args.hyperparameters)

    wandb_configs = fn_args.custom_config["wandb"]
    wandb.login(key=wandb_configs["API_KEY"])
    wandb_project = wandb_configs["PROJECT"]

    ....
    optimizer_type = hp.get("optimizer_type")
    learning_rate = hp.get("learning_rate")
    weight_decay = hp.get("weight_decay")
    epochs = hp.get("fulltrain_epochs")
    unique_id = wandb_configs["FINAL_RUN_ID"]

    wandb.init(project=wandb_project, config=hp.values, name=unique_id)
    wandb.log({"optimizer_type": optimizer_type})
    wandb.log({"learning_rate": learning_rate})
    wandb.log({"weight_decay": weight_decay})

    model = MyHyperModel().build(hp)
    model.fit(
        ....
        callbacks = [...,
            wandb.keras.WandbMetricsLogger(log_freq='epoch')
        ]
    )

    wandb.finish()
    model.save(....)

WandBPusher component

WANDB_PUSHER_ARGS = {
    "access_token": "$WANDB_ACCESS_TOKEN",
    "project_name": PIPELINE_NAME,
    "run_name": WANDB_RUN_ID,
    "model": trainer.outputs["model"]
    "model_name": "final_model",
    "aliases": ["test_aliases"],
    "space_config": {
        "app_path": "huggingface.apps.gradio",
        "hf_username": "chansung",
        "hf_repo_name": PIPELINE_NAME,
        "hf_access_token": "$HF_ACCESS_TOKEN"
    }
}
pusher = WandBPusher(**wandb_pusher_args)

access_token: This is a W&B access token. You need it to interact with W&B programatically. We have used the same access token in Tuner and Trainer components.
project_name: This is our W&B project name and this should be the same as the project name used in Tuner and Trainer components. Otherwise, we can not associate the model and where it belongs to.
run_name: This specifies the run name in the designated project name. Since the final model is the output from Trainer component, we need to associate the model to the run name that is used to log Trainer's process.
model: This points to the SavedModel artifact saved as the output from Trainer component.
model_name: This is the identifier of the model in W&B Model Registry.
aliases: This is the secondary identifier of the model in W&B Model Registry. It is particularly useful to identify a model with version name or latest.

space_config.app_path: This is the module path where the application files and codes are. Under the path, all the to be hosted files in Hugging Face Space should be stored such as README.md, app.py, and requirements.txt.
space_config.hf_username: Any repository in Hugging Face Hub should be identified in username/reponame, so this information is the first part of the repository identifier.
space_config.hf_repo_name: This is just the second part of the repository identifier.
space_config.hf_access_token: This is our Hugging Face access token and we need this to interact with HuggingFace Hub programatically.

$MODEL_PROJECT is replaced by the project_name in the WANDB_PUSHER_ARGS.
$MODEL_RUN is replaced by the run_name in the WANDB_PUSHER_ARGS.
$MODEL_NAME is replaced by the model_name in the WANDB_PUSHER_ARGS.
$MODEL_VERSION is replaced by the internally decided model version.

wandb.login(key=wb_token)
wandb.init(project="$MODEL_PROJECT", id="$MODEL_RUN", resume=True)
path = wandb.use_artifact('tfx-vit-pipeline/$MODEL_NAME:$MODEL_VERSION', type='model').download()

tar = tarfile.open(f"{path}/$MODEL_FILENAME")
tar.extractall(path=".")

MODEL = tf.keras.models.load_model("./model")

wandb.login(key=wb_token)
wandb.init(project="tfx-vit-pipeline", id="avi9pkis", resume=True)
path = wandb.use_artifact('tfx-vit-pipeline/final_model:1683870912', type='model').download()

tar = tarfile.open(f"{path}/model.tar.gz")
tar.extractall(path=".")

MODEL = tf.keras.models.load_model("./model")

Conclusion