Passing Trainer state as an artifact in kfp.v2 pipeline

Intermediate
1

I’m trying to create a kfp pipeline fine-tuning BERT on GCP’s VertexAI. I can save a model as an artifact, but have troubles saving the Trainer (Trainer’s state), as I would like to split eval and testing into two separate pipeline components and to achieve that - I need to reconstruct the Trainer in the latter.

I have, created trainer_artifact and would like to save_state() into tainer_artifact.path. However, save_state() does not accept arguments and saves to model’s directory by default. How can I save trainer as an artifact in this situation? Maybe there is a workaround for it, eg. recreating it in the next step?

I attach my code.

@component(
    packages_to_install = [
        "pandas",
        "datasets",
        "transformers"
    ],
)
def fine_tune_modell(
    small_train_dataset: Input[Dataset],
    small_eval_dataset: Input[Dataset],
#     full_train_dataset: Input[Dataset],
#     full_eval_dataset: Input[Dataset],
    model_artifact: Output[Model],
    trainer_artifact: Output[Artifact]
):
    
    import pandas as pd
    import numpy as np
    import datasets
    from transformers import AutoModelForSequenceClassification
    from transformers import TrainingArguments
    from transformers import Trainer
    from datasets import load_metric
    
    # create model
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

    # load data
    train_data = datasets.load_from_disk(small_train_dataset.path)
    eval_data = datasets.load_from_disk(small_eval_dataset.path)
    
    metric = load_metric("accuracy")

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        return metric.compute(predictions=predictions, references=labels)
        
    training_args = TrainingArguments(
        output_dir="test_trainer",
        evaluation_strategy="epoch",
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        num_train_epochs=3,
        seed=0,)
    
    trainer = Trainer(
        model=model, 
        args=training_args, 
        train_dataset=train_data, 
        eval_dataset=eval_data,
        compute_metrics=compute_metrics
    )
    
    train_output = trainer.train()

    model_artifact.metadata["train_output"] = train_output
    model_artifact.metadata["framework"] = "Pytorch"
    
    # How to save trainer?
    
    model.save_model(model_artifact.path)```
1 Like
2

Maybe you can move the state file created?