hf_mae_refactored.py

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import torch
import transformers
from datasets import load_dataset
from torchvision.transforms import (
    Compose,
    Lambda,
    Normalize,
    RandomHorizontalFlip,
    RandomResizedCrop,
    ToTensor,
)
from torchvision.transforms.functional import InterpolationMode
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    ViTImageProcessor,
    ViTMAEConfig,
    ViTMAEForPreTraining,
)
from transformers.integrations import ClearMLCallback, TensorBoardCallback
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version

""" Pre-training a 🤗 ViT model as an MAE (masked autoencoder), as proposed in https://arxiv.org/abs/2111.06377."""

logger = logging.getLogger(__name__)

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.32.0.dev0")

require_version(
    "datasets>=1.8.0",
    "To fix: pip install -r examples/pytorch/image-pretraining/requirements.txt",
)


class CustomTrainer(Trainer):
    def _remove_unused_columns(self, dataset, **kwargs):
        return dataset


class CustomClearMLCallback(ClearMLCallback):
    def on_log(
        self, args, state, control, model=None, tokenizer=None, logs=None, **kwargs
    ):
        if self._clearml is None:
            return
        if not self._initialized:
            self.setup(args, state, model, tokenizer, **kwargs)
        if state.is_world_process_zero:
            eval_prefix = "eval_"
            eval_prefix_len = len(eval_prefix)
            test_prefix = "test_"
            test_prefix_len = len(test_prefix)
            single_value_scalars = [
                "train_runtime",
                "train_samples_per_second",
                "train_steps_per_second",
                "train_loss",
                "total_flos",
                "epoch",
            ]
            for k, v in logs.items():
                if isinstance(v, (int, float)):
                    if k in single_value_scalars:
                        self._clearml_task.get_logger().report_scalar(
                            title=k,
                            series=k,
                            iteration=state.global_step,
                            value=v,
                        )
                    elif k.startswith(eval_prefix):
                        self._clearml_task.get_logger().report_scalar(
                            title=k[eval_prefix_len:],
                            series="eval",
                            value=v,
                            iteration=state.global_step,
                        )
                    elif k.startswith(test_prefix):
                        self._clearml_task.get_logger().report_scalar(
                            title=k[test_prefix_len:],
                            series="test",
                            value=v,
                            iteration=state.global_step,
                        )
                    else:
                        self._clearml_task.get_logger().report_scalar(
                            title=k,
                            series="train",
                            value=v,
                            iteration=state.global_step,
                        )
                else:
                    logger.warning(
                        "Trainer is attempting to log a value of "
                        f'"{v}" of type {type(v)} for key "{k}" as a scalar. '
                        "This invocation of ClearML logger's  report_scalar() "
                        "is incorrect so we dropped this attribute."
                    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """

    dataset_name: Optional[str] = field(
        default="cifar10",
        metadata={"help": "Name of a dataset from the datasets package"},
    )
    dataset_config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The configuration name of the dataset to use (via the datasets library)."
        },
    )
    image_column_name: Optional[str] = field(
        default=None, metadata={"help": "The column name of the images in the files."}
    )
    train_dir: Optional[str] = field(
        default=None, metadata={"help": "A folder containing the training data."}
    )
    validation_dir: Optional[str] = field(
        default=None, metadata={"help": "A folder containing the validation data."}
    )
    train_val_split: Optional[float] = field(
        default=0.15, metadata={"help": "Percent to split off of train for validation."}
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )

    def __post_init__(self):
        data_files = {}
        if self.train_dir is not None:
            data_files["train"] = self.train_dir
        if self.validation_dir is not None:
            data_files["val"] = self.validation_dir
        self.data_files = data_files if data_files else None


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/image processor we are going to pre-train.
    """

    model_name_or_path: str = field(
        default="facebook/vit-mae-base",
        metadata={
            "help": (
                "The model checkpoint for weights initialization.Don't set if you want to train a model from scratch."
            )
        },
    )
    config_name: Optional[str] = field(
        default="facebook/vit-mae-base",
        metadata={
            "help": "Pretrained config name or path if not the same as model_name_or_path"
        },
    )
    config_overrides: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "Override some existing default config settings when a model is trained from scratch. Example: "
                "n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index"
            )
        },
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={
            "help": "Where do you want to store the pretrained models downloaded from s3"
        },
    )
    model_revision: str = field(
        default="main",
        metadata={
            "help": "The specific model version to use (can be a branch name, tag name or commit id)."
        },
    )
    image_processor_name: str = field(
        default=None, metadata={"help": "Name or path of preprocessor config."}
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": (
                "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
                "with private models)."
            )
        },
    )
    mask_ratio: float = field(
        default=0.75,
        metadata={
            "help": "The ratio of the number of masked tokens in the input sequence."
        },
    )
    norm_pix_loss: bool = field(
        default=True,
        metadata={
            "help": "Whether or not to train with normalized pixel values as target."
        },
    )


@dataclass
class CustomTrainingArguments(TrainingArguments):
    base_learning_rate: float = field(
        default=1e-3,
        metadata={
            "help": "Base learning rate: absolute_lr = base_lr * total_batch_size / 256."
        },
    )


def collate_fn(examples):
    pixel_values = torch.stack([example["pixel_values"] for example in examples])
    return {"pixel_values": pixel_values}


def parse_arguments():
    parser = HfArgumentParser(
        (ModelArguments, DataTrainingArguments, CustomTrainingArguments)
    )
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(
            json_file=os.path.abspath(sys.argv[1])
        )
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    return model_args, data_args, training_args


def configure_logging(training_args):
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()


def detect_last_checkpoint(training_args):
    last_checkpoint = None
    if (
        os.path.isdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif (
            last_checkpoint is not None and training_args.resume_from_checkpoint is None
        ):
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )
    return last_checkpoint


def create_dataset(data_args, model_args):
    ds = load_dataset(
        data_args.dataset_name,
        data_args.dataset_config_name,
        data_files=data_args.data_files,
        cache_dir=model_args.cache_dir,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    # If we don't have a validation split, split off a percentage of train as validation.
    data_args.train_val_split = (
        None if "validation" in ds.keys() else data_args.train_val_split
    )
    if isinstance(data_args.train_val_split, float) and data_args.train_val_split > 0.0:
        split = ds["train"].train_test_split(data_args.train_val_split)
        ds["train"] = split["train"]
        ds["validation"] = split["test"]
    return ds


def load_model_config(model_args):
    # Load pretrained model and image processor
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config_kwargs = {
        "cache_dir": model_args.cache_dir,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    if model_args.config_name:
        config = ViTMAEConfig.from_pretrained(model_args.config_name, **config_kwargs)
    elif model_args.model_name_or_path:
        config = ViTMAEConfig.from_pretrained(
            model_args.model_name_or_path, **config_kwargs
        )
    else:
        config = ViTMAEConfig()
        logger.warning("You are instantiating a new config instance from scratch.")
        if model_args.config_overrides is not None:
            logger.info(f"Overriding config: {model_args.config_overrides}")
            config.update_from_string(model_args.config_overrides)
            logger.info(f"New config: {config}")

    # adapt config
    config.update(
        {
            "mask_ratio": model_args.mask_ratio,
            "norm_pix_loss": model_args.norm_pix_loss,
        }
    )
    return config, config_kwargs


def create_image_processor(model_args, config_kwargs):
    if model_args.image_processor_name:
        image_processor = ViTImageProcessor.from_pretrained(
            model_args.image_processor_name, **config_kwargs
        )
    elif model_args.model_name_or_path:
        image_processor = ViTImageProcessor.from_pretrained(
            model_args.model_name_or_path, **config_kwargs
        )
    else:
        image_processor = ViTImageProcessor()
    return image_processor


def get_image_column_name(training_args, data_args, ds):
    if training_args.do_train:
        column_names = ds["train"].column_names
    else:
        column_names = ds["validation"].column_names

    if data_args.image_column_name is not None:
        image_column_name = data_args.image_column_name
    elif "image" in column_names:
        image_column_name = "image"
    elif "img" in column_names:
        image_column_name = "img"
    else:
        image_column_name = column_names[0]

    return image_column_name


def get_dataset_transforms(image_processor):
    if "shortest_edge" in image_processor.size:
        size = image_processor.size["shortest_edge"]
    else:
        size = (image_processor.size["height"], image_processor.size["width"])
    transforms = Compose(
        [
            Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
            RandomResizedCrop(
                size, scale=(0.2, 1.0), interpolation=InterpolationMode.BICUBIC
            ),
            RandomHorizontalFlip(),
            ToTensor(),
            Normalize(mean=image_processor.image_mean, std=image_processor.image_std),
        ]
    )

    return transforms


def validate_dataset(ds, training_args):
    if training_args.do_train:
        if "train" not in ds:
            raise ValueError("--do_train requires a train dataset")

    if training_args.do_eval:
        if "validation" not in ds:
            raise ValueError("--do_eval requires a validation dataset")


def shuffle_dataset(ds, seed, max_train_samples):
    return ds.shuffle(seed=seed).select(range(max_train_samples))


def cut_dataset(ds, training_args, data_args):
    if training_args.do_train:
        if data_args.max_train_samples is not None:
            ds["train"] = shuffle_dataset(
                ds["train"], training_args.seed, data_args.max_train_samples
            )
    if training_args.do_eval:
        if data_args.max_eval_samples is not None:
            ds["validation"] = shuffle_dataset(
                ds["validation"], training_args.seed, data_args.max_eval_samples
            )
    return ds


def set_dataset_transforms(ds, transforms, training_args, image_column_name):
    def preprocess_images(examples):
        examples["pixel_values"] = [
            transforms(image) for image in examples[image_column_name]
        ]
        return examples

    if training_args.do_train:
        ds["train"].set_transform(preprocess_images)
    if training_args.do_eval:
        ds["validation"].set_transform(preprocess_images)
    return ds


def get_absolute_lr(training_args):
    total_train_batch_size = (
        training_args.train_batch_size
        * training_args.gradient_accumulation_steps
        * training_args.world_size
    )
    return training_args.base_learning_rate * total_train_batch_size / 256


def set_absolute_lr(training_args, absolute_lr):
    if training_args.base_learning_rate is not None:
        training_args.learning_rate = absolute_lr


def prepare_dataset(data_args, model_args, training_args, image_processor):
    ds = create_dataset(data_args, model_args)
    image_column_name = get_image_column_name(training_args, data_args, ds)
    transforms = get_dataset_transforms(image_processor)
    validate_dataset(ds, training_args)
    cut_dataset(ds, training_args, data_args)
    set_dataset_transforms(ds, transforms, training_args, image_column_name)
    return ds


def create_model(model_args, config):
    if model_args.model_name_or_path:
        model = ViTMAEForPreTraining.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        logger.info("Training new model from scratch")
        model = ViTMAEForPreTraining(config)
    return model


def get_trainer(model, training_args, ds, image_processor, collate_fn):
    trainer = CustomTrainer(
        model=model,
        args=training_args,
        train_dataset=ds["train"] if training_args.do_train else None,
        eval_dataset=ds["validation"] if training_args.do_eval else None,
        tokenizer=image_processor,
        data_collator=collate_fn,
    )
    os.environ["CLEARML_PROJECT"] = "plantsense"
    os.environ["CLEARML_TASK"] = "ViT-MAE debug"
    os.environ["CLEARML_LOG_MODEL"] = "True"
    trainer.add_callback(
        CustomClearMLCallback(),
    )
    trainer.add_callback(
        TensorBoardCallback(),
    )

    return trainer


def train(training_args, last_checkpoint, trainer):
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()


def eval(training_args, trainer):
    if training_args.do_eval:
        metrics = trainer.evaluate()
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)


def main():
    model_args, data_args, training_args = parse_arguments()

    configure_logging(training_args)

    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    config, config_kwargs = load_model_config(model_args)

    image_processor = create_image_processor(model_args, config_kwargs)

    ds = prepare_dataset(
        data_args,
        model_args,
        training_args,
        image_processor,
    )

    last_checkpoint = detect_last_checkpoint(training_args)
    model = create_model(model_args, config)

    absolute_lr = get_absolute_lr(training_args)
    set_absolute_lr(training_args, absolute_lr)

    trainer = get_trainer(
        model,
        training_args,
        ds,
        image_processor=image_processor,
        collate_fn=collate_fn,
    )

    train(
        training_args,
        last_checkpoint,
        trainer,
    )
    eval(training_args, trainer)


if __name__ == "__main__":
    main()