usrbinenv python codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and pretraining a vit model as an mae masked autoencoder as proposed in https arxiv orgabs2111 06377 logger logging getloggername will error if the minimal version of transformers is not installed remove at your own risks checkminversion4 36 0 dev0 requireversiondatasets1 8 0 to fix pip install r examplespytorchimagepretrainingrequirements txt dataclass class datatrainingarguments datasetname optionalstr field defaultcifar10 metadatahelp name of a dataset from the datasets package datasetconfigname optionalstr field defaultnone metadatahelp the configuration name of the dataset to use via the datasets library imagecolumnname optionalstr field defaultnone metadatahelp the column name of the images in the files traindir optionalstr fielddefaultnone metadatahelp a folder containing the training data validationdir optionalstr fielddefaultnone metadatahelp a folder containing the validation data trainvalsplit optionalfloat field default0 15 metadatahelp percent to split off of train for validation maxtrainsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of training examples to this value if set maxevalsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of evaluation examples to this value if set def postinitself datafiles if self traindir is not none datafilestrain self traindir if self validationdir is not none datafilesval self validationdir self datafiles datafiles if datafiles else none dataclass class modelarguments modelnameorpath str field defaultnone metadata help the model checkpoint for weights initialization don t set if you want to train a model from scratch configname optionalstr field defaultnone metadatahelp pretrained config name or path if not the same as modelnameorpath configoverrides optionalstr field defaultnone metadata help override some existing default config settings when a model is trained from scratch example nembd10 residpdrop0 2 scaleattnweightsfalse summarytypeclsindex cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from s3 modelrevision str field defaultmain metadatahelp the specific model version to use can be a branch name tag name or commit id imageprocessorname str fielddefaultnone metadatahelp name or path of preprocessor config token str field defaultnone metadata help the token to use as http bearer ization for remote files if not specified will use the token generated when running huggingfacecli login stored in huggingface useauthtoken bool field defaultnone metadata help the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead maskratio float field default0 75 metadatahelp the ratio of the number of masked tokens in the input sequence normpixloss bool field defaulttrue metadatahelp whether or not to train with normalized pixel values as target dataclass class customtrainingargumentstrainingarguments baselearningrate float field default1e3 metadatahelp base learning rate absolutelr baselr totalbatchsize 256 def collatefnexamples pixelvalues torch stackexamplepixelvalues for example in examples return pixelvalues pixelvalues def main see all possible arguments in srctransformerstrainingargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns parser hfargumentparsermodelarguments datatrainingarguments customtrainingarguments if lensys argv 2 and sys argv1 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 modelargs dataargs trainingargs parser parsejsonfilejsonfileos path abspathsys argv1 else modelargs dataargs trainingargs parser parseargsintodataclasses if modelargs useauthtoken is not none warnings warn the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead futurewarning if modelargs token is not none raise valueerrortoken and useauthtoken are both specified please set only the argument token modelargs token modelargs useauthtoken sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your pythonpytorch versions sendexampletelemetryrunmae modelargs dataargs setup logging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s handlerslogging streamhandlersys stdout if trainingargs shouldlog the default of trainingargs loglevel is passive so we set log level at info here to have that default transformers utils logging setverbosityinfo loglevel trainingargs getprocessloglevel logger setlevelloglevel transformers utils logging setverbosityloglevel transformers utils logging enabledefaulthandler transformers utils logging enableexplicitformat log on each process the small summary logger warning fprocess rank trainingargs localrank device trainingargs device ngpu trainingargs ngpu fdistributed training trainingargs parallelmode value distributed 16bits training trainingargs fp16 logger infoftrainingevaluation parameters trainingargs detecting last checkpoint lastcheckpoint none if os path isdirtrainingargs outputdir and trainingargs dotrain and not trainingargs overwriteoutputdir lastcheckpoint getlastcheckpointtrainingargs outputdir if lastcheckpoint is none and lenos listdirtrainingargs outputdir 0 raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to overcome elif lastcheckpoint is not none and trainingargs resumefromcheckpoint is none logger info fcheckpoint detected resuming training at lastcheckpoint to avoid this behavior change the outputdir or add overwriteoutputdir to train from scratch initialize our dataset ds loaddataset dataargs datasetname dataargs datasetconfigname datafilesdataargs datafiles cachedirmodelargs cachedir tokenmodelargs token if we don t have a validation split split off a percentage of train as validation dataargs trainvalsplit none if validation in ds keys else dataargs trainvalsplit if isinstancedataargs trainvalsplit float and dataargs trainvalsplit 0 0 split dstrain traintestsplitdataargs trainvalsplit dstrain splittrain dsvalidation splittest load pretrained model and image processor distributed training the frompretrained methods guarantee that only one local process can concurrently download model vocab configkwargs cachedir modelargs cachedir revision modelargs modelrevision token modelargs token if modelargs configname config vitmaeconfig frompretrainedmodelargs configname configkwargs elif modelargs modelnameorpath config vitmaeconfig frompretrainedmodelargs modelnameorpath configkwargs else config vitmaeconfig logger warningyou are instantiating a new config instance from scratch if modelargs configoverrides is not none logger infofoverriding config modelargs configoverrides config updatefromstringmodelargs configoverrides logger infofnew config config adapt config config update maskratio modelargs maskratio normpixloss modelargs normpixloss create image processor if modelargs imageprocessorname imageprocessor vitimageprocessor frompretrainedmodelargs imageprocessorname configkwargs elif modelargs modelnameorpath imageprocessor vitimageprocessor frompretrainedmodelargs modelnameorpath configkwargs else imageprocessor vitimageprocessor create model if modelargs modelnameorpath model vitmaeforpretraining frompretrained modelargs modelnameorpath fromtfbool ckpt in modelargs modelnameorpath configconfig cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token else logger infotraining new model from scratch model vitmaeforpretrainingconfig if trainingargs dotrain columnnames dstrain columnnames else columnnames dsvalidation columnnames if dataargs imagecolumnname is not none imagecolumnname dataargs imagecolumnname elif image in columnnames imagecolumnname image elif img in columnnames imagecolumnname img else imagecolumnname columnnames0 transformations as done in original mae paper source https github comfacebookresearchmaeblobmainmainpretrain py if shortestedge in imageprocessor size size imageprocessor sizeshortestedge else size imageprocessor sizeheight imageprocessor sizewidth transforms compose lambdalambda img img convertrgb if img mode rgb else img randomresizedcropsize scale0 2 1 0 interpolationinterpolationmode bicubic randomhorizontalflip totensor normalizemeanimageprocessor imagemean stdimageprocessor imagestd def preprocessimagesexamples set the training transforms set the validation transforms compute absolute learning rate initialize our trainer training evaluation write model card and optionally push to hub for xlaspawn tpus usr bin env python coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and pre training a vit model as an mae masked autoencoder as proposed in https arxiv org abs 2111 06377 will error if the minimal version of transformers is not installed remove at your own risks 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 arguments pertaining to which model config image processor we are going to pre train see all possible arguments in src transformers training_args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns 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 sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your python pytorch versions setup logging the default of training_args log_level is passive so we set log level at info here to have that default log on each process the small summary detecting last checkpoint initialize our dataset if we don t have a validation split split off a percentage of train as validation load pretrained model and image processor distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab adapt config create image processor create model transformations as done in original mae paper source https github com facebookresearch mae blob main main_pretrain py preprocess a batch of images by applying transforms set the training transforms set the validation transforms compute absolute learning rate initialize our trainer training evaluation write model card and optionally push to hub for xla_spawn tpus
import logging
import os
import sys
import warnings
from dataclasses import dataclass, field
from typing import Optional

import torch
from datasets import load_dataset
from torchvision.transforms import Compose, Lambda, Normalize, RandomHorizontalFlip, RandomResizedCrop, ToTensor
from torchvision.transforms.functional import InterpolationMode

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    ViTImageProcessor,
    ViTMAEConfig,
    ViTMAEForPreTraining,
)
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




logger = logging.getLogger(__name__)


check_min_version("4.36.0.dev0")

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


@dataclass
class DataTrainingArguments:
    

    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:
    

    model_name_or_path: str = field(
        default=None,
        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=None, 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."})
    token: str = field(
        default=None,
        metadata={
            "help": (
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
            )
        },
    )
    use_auth_token: bool = field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
        },
    )
    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 main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, CustomTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        
        
        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()

    if model_args.use_auth_token is not None:
        warnings.warn(
            "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead.",
            FutureWarning,
        )
        if model_args.token is not None:
            raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.")
        model_args.token = model_args.use_auth_token

    
    
    send_example_telemetry("run_mae", model_args, data_args)

    
    logging.basicConfig(
        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:
        
        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()

    
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {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."
            )

    
    ds = load_dataset(
        data_args.dataset_name,
        data_args.dataset_config_name,
        data_files=data_args.data_files,
        cache_dir=model_args.cache_dir,
        token=model_args.token,
    )

    
    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"]

    
    
    
    
    
    config_kwargs = {
        "cache_dir": model_args.cache_dir,
        "revision": model_args.model_revision,
        "token": model_args.token,
    }
    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}")

    
    config.update(
        {
            "mask_ratio": model_args.mask_ratio,
            "norm_pix_loss": model_args.norm_pix_loss,
        }
    )

    
    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()

    
    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,
            token=model_args.token,
        )
    else:
        logger.info("Training new model from scratch")
        model = ViTMAEForPreTraining(config)

    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]

    
    
    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),
        ]
    )

    def preprocess_images(examples):
        

        examples["pixel_values"] = [transforms(image) for image in examples[image_column_name]]
        return examples

    if training_args.do_train:
        if "train" not in ds:
            raise ValueError("--do_train requires a train dataset")
        if data_args.max_train_samples is not None:
            ds["train"] = ds["train"].shuffle(seed=training_args.seed).select(range(data_args.max_train_samples))
        
        ds["train"].set_transform(preprocess_images)

    if training_args.do_eval:
        if "validation" not in ds:
            raise ValueError("--do_eval requires a validation dataset")
        if data_args.max_eval_samples is not None:
            ds["validation"] = (
                ds["validation"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
            )
        
        ds["validation"].set_transform(preprocess_images)

    
    total_train_batch_size = (
        training_args.train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
    )
    if training_args.base_learning_rate is not None:
        training_args.learning_rate = training_args.base_learning_rate * total_train_batch_size / 256

    
    trainer = Trainer(
        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,
    )

    
    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()

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

    
    kwargs = {
        "tasks": "masked-auto-encoding",
        "dataset": data_args.dataset_name,
        "tags": ["masked-auto-encoding"],
    }
    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:
        trainer.create_model_card(**kwargs)


def _mp_fn(index):
    
    main()


if __name__ == "__main__":
    main()