usrbinenv python codingutf8 2021 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 finetuning a transformers model for image classification 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 examplespytorchimageclassificationrequirements txt modelconfigclasses listmodelforimageclassificationmapping keys modeltypes tupleconf modeltype for conf in modelconfigclasses def pilloaderpath str with openpath rb as f im image openf return im convertrgb dataclass class datatrainingarguments datasetname optionalstr field defaultnone metadata help name of a dataset from the hub could be your own possibly private dataset hosted on the hub datasetconfigname optionalstr field defaultnone metadatahelp the configuration name of the dataset to use via the datasets library 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 if self datasetname is none and self traindir is none and self validationdir is none raise valueerror you must specify either a dataset name from the hub or a train andor validation directory dataclass class modelarguments modelnameorpath str field defaultgooglevitbasepatch16224in21k metadatahelp path to pretrained model or model identifier from huggingface comodels modeltype optionalstr field defaultnone metadatahelp if training from scratch pass a model type from the list joinmodeltypes configname optionalstr field defaultnone metadatahelp pretrained config name or path if not the same as modelname 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 trustremotecode bool field defaultfalse metadata help whether or not to allow for custom models defined on the hub in their own modeling files this option should only be set to true for repositories you trust and in which you have read the code as it will execute code present on the hub on your local machine ignoremismatchedsizes bool field defaultfalse metadatahelp will enable to load a pretrained model whose head dimensions are different def collatefnexamples pixelvalues torch stackexamplepixelvalues for example in examples labels torch tensorexamplelabels for example in examples return pixelvalues pixelvalues labels labels 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 trainingarguments 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 sendexampletelemetryrunimageclassification 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 set seed before initializing model setseedtrainingargs seed initialize our dataset and prepare it for the imageclassification task if dataargs datasetname is not none dataset loaddataset dataargs datasetname dataargs datasetconfigname cachedirmodelargs cachedir taskimageclassification tokenmodelargs token else datafiles if dataargs traindir is not none datafilestrain os path joindataargs traindir if dataargs validationdir is not none datafilesvalidation os path joindataargs validationdir dataset loaddataset imagefolder datafilesdatafiles cachedirmodelargs cachedir taskimageclassification if we don t have a validation split split off a percentage of train as validation dataargs trainvalsplit none if validation in dataset keys else dataargs trainvalsplit if isinstancedataargs trainvalsplit float and dataargs trainvalsplit 0 0 split datasettrain traintestsplitdataargs trainvalsplit datasettrain splittrain datasetvalidation splittest prepare label mappings we ll include these in the model s config to get human readable labels in the inference api labels datasettrain featureslabels names label2id id2label for i label in enumeratelabels label2idlabel stri id2labelstri label load the accuracy metric from the datasets package metric evaluate loadaccuracy define our computemetrics function it takes an evalprediction object a namedtuple with a predictions and labelids field and has to return a dictionary string to float def computemetricsp define torchvision transforms to be applied to each image apply traintransforms across a batch examplebatchpixelvalues traintransformspilimg convertrgb for pilimg in examplebatchimage return examplebatch def valtransformsexamplebatch set the training transforms set the validation transforms initalize our trainer training evaluation write model card and optionally push to hub usr bin env python coding utf 8 2021 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 fine tuning a transformers model for image classification 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 tokenizer we are going to fine tune from 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 set seed before initializing model initialize our dataset and prepare it for the image classification task if we don t have a validation split split off a percentage of train as validation prepare label mappings we ll include these in the model s config to get human readable labels in the inference api load the accuracy metric from the datasets package define our compute_metrics function it takes an evalprediction object a namedtuple with a predictions and label_ids field and has to return a dictionary string to float computes accuracy on a batch of predictions define torchvision transforms to be applied to each image apply _train_transforms across a batch apply _val_transforms across a batch set the training transforms set the validation transforms initalize our trainer training evaluation write model card and optionally push to hub
import logging
import os
import sys
import warnings
from dataclasses import dataclass, field
from typing import Optional

import evaluate
import numpy as np
import torch
from datasets import load_dataset
from PIL import Image
from torchvision.transforms import (
    CenterCrop,
    Compose,
    Lambda,
    Normalize,
    RandomHorizontalFlip,
    RandomResizedCrop,
    Resize,
    ToTensor,
)

import transformers
from transformers import (
    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    AutoModelForImageClassification,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
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-classification/requirements.txt")

MODEL_CONFIG_CLASSES = list(MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


def pil_loader(path: str):
    with open(path, "rb") as f:
        im = Image.open(f)
        return im.convert("RGB")


@dataclass
class DataTrainingArguments:
    

    dataset_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Name of a dataset from the hub (could be your own, possibly private dataset hosted on the hub)."
        },
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    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):
        if self.dataset_name is None and (self.train_dir is None and self.validation_dir is None):
            raise ValueError(
                "You must specify either a dataset name from the hub or a train and/or validation directory."
            )


@dataclass
class ModelArguments:
    

    model_name_or_path: str = field(
        default="google/vit-base-patch16-224-in21k",
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"},
    )
    model_type: Optional[str] = field(
        default=None,
        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    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."
        },
    )
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option"
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
                "execute code present on the Hub on your local machine."
            )
        },
    )
    ignore_mismatched_sizes: bool = field(
        default=False,
        metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."},
    )


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


def main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    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_image_classification", 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."
            )

    
    set_seed(training_args.seed)

    
    if data_args.dataset_name is not None:
        dataset = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            task="image-classification",
            token=model_args.token,
        )
    else:
        data_files = {}
        if data_args.train_dir is not None:
            data_files["train"] = os.path.join(data_args.train_dir, "**")
        if data_args.validation_dir is not None:
            data_files["validation"] = os.path.join(data_args.validation_dir, "**")
        dataset = load_dataset(
            "imagefolder",
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            task="image-classification",
        )

    
    data_args.train_val_split = None if "validation" in dataset.keys() else data_args.train_val_split
    if isinstance(data_args.train_val_split, float) and data_args.train_val_split > 0.0:
        split = dataset["train"].train_test_split(data_args.train_val_split)
        dataset["train"] = split["train"]
        dataset["validation"] = split["test"]

    
    
    labels = dataset["train"].features["labels"].names
    label2id, id2label = {}, {}
    for i, label in enumerate(labels):
        label2id[label] = str(i)
        id2label[str(i)] = label

    
    metric = evaluate.load("accuracy")

    
    
    def compute_metrics(p):
        
        return metric.compute(predictions=np.argmax(p.predictions, axis=1), references=p.label_ids)

    config = AutoConfig.from_pretrained(
        model_args.config_name or model_args.model_name_or_path,
        num_labels=len(labels),
        label2id=label2id,
        id2label=id2label,
        finetuning_task="image-classification",
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    model = AutoModelForImageClassification.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,
        trust_remote_code=model_args.trust_remote_code,
        ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
    )
    image_processor = AutoImageProcessor.from_pretrained(
        model_args.image_processor_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )

    
    if "shortest_edge" in image_processor.size:
        size = image_processor.size["shortest_edge"]
    else:
        size = (image_processor.size["height"], image_processor.size["width"])
    normalize = (
        Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
        if hasattr(image_processor, "image_mean") and hasattr(image_processor, "image_std")
        else Lambda(lambda x: x)
    )
    _train_transforms = Compose(
        [
            RandomResizedCrop(size),
            RandomHorizontalFlip(),
            ToTensor(),
            normalize,
        ]
    )
    _val_transforms = Compose(
        [
            Resize(size),
            CenterCrop(size),
            ToTensor(),
            normalize,
        ]
    )

    def train_transforms(example_batch):
        
        example_batch["pixel_values"] = [
            _train_transforms(pil_img.convert("RGB")) for pil_img in example_batch["image"]
        ]
        return example_batch

    def val_transforms(example_batch):
        
        example_batch["pixel_values"] = [_val_transforms(pil_img.convert("RGB")) for pil_img in example_batch["image"]]
        return example_batch

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

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

    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"] if training_args.do_train else None,
        eval_dataset=dataset["validation"] if training_args.do_eval else None,
        compute_metrics=compute_metrics,
        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 = {
        "finetuned_from": model_args.model_name_or_path,
        "tasks": "image-classification",
        "dataset": data_args.dataset_name,
        "tags": ["image-classification", "vision"],
    }
    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:
        trainer.create_model_card(**kwargs)


if __name__ == "__main__":
    main()