examples/tensorflow/image-classification/run_image_classification.py/main
def main():
    # 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.
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TFTrainingArguments))
    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()

    if not (training_args.do_train or training_args.do_eval or training_args.do_predict):
        exit("Must specify at least one of --do_train, --do_eval or --do_predict!")

    # 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/TensorFlow versions.
    send_example_telemetry("run_image_classification", model_args, data_args, framework="tensorflow")

    # Checkpoints. Find the checkpoint the use when loading the model.
    checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        checkpoint = get_last_checkpoint(training_args.output_dir)
        if 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 checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)

    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
    # Log on each process the small summary:
    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}")

    # region Dataset and labels
    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Initialize our dataset and prepare it for the 'image-classification' task.
    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",
            use_auth_token=True if model_args.use_auth_token else None,
        )
    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",
        )
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.

    # Prepare label mappings.
    # We'll include these in the model's config to get human readable labels in the Inference API.
    labels = dataset["train"].features["labels"].names
    label2id, id2label = {}, {}
    for i, label in enumerate(labels):
        label2id[label] = str(i)
        id2label[str(i)] = label

    # Load model image processor and configuration
    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,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    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,
        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 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"]

    # Define our data preprocessing function. It takes an image file path as input and returns
    # Write a note describing the resizing behaviour.
    if "shortest_edge" in image_processor.size:
        # We instead set the target size as (shortest_edge, shortest_edge) to here to ensure all images are batchable.
        image_size = (image_processor.size["shortest_edge"], image_processor.size["shortest_edge"])
    else:
        image_size = (image_processor.size["height"], image_processor.size["width"])

    def _train_transforms(image):
        img_size = image_size
        image = tf.keras.utils.img_to_array(image)
        image = random_resized_crop(image, size=img_size)
        image = tf.image.random_flip_left_right(image)
        image /= 255.0
        image = (image - image_processor.image_mean) / image_processor.image_std
        image = tf.transpose(image, perm=[2, 0, 1])
        return image

    def _val_transforms(image):
        image = tf.keras.utils.img_to_array(image)
        image = tf.image.resize(image, size=image_size)
        # image = np.array(image) # FIXME - use tf.image function
        image = center_crop(image, size=image_size)
        image /= 255.0
        image = (image - image_processor.image_mean) / image_processor.image_std
        image = tf.transpose(image, perm=[2, 0, 1])
        return image

    def train_transforms(example_batch):
        """Apply _train_transforms across a 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):
        """Apply _val_transforms across a batch."""
        example_batch["pixel_values"] = [_val_transforms(pil_img.convert("RGB")) for pil_img in example_batch["image"]]
        return example_batch

    train_dataset = None
    if training_args.do_train:
        if "train" not in dataset:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = dataset["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.shuffle(seed=training_args.seed).select(range(data_args.max_train_samples))
        train_dataset = train_dataset.map(
            train_transforms,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    eval_dataset = None
    if training_args.do_eval:
        if "validation" not in dataset:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = dataset["validation"]
        if data_args.max_eval_samples is not None:
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
        # Set the validation transforms
        eval_dataset = eval_dataset.map(
            val_transforms,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    predict_dataset = None
    if training_args.do_predict:
        if "test" not in dataset:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = dataset["test"]
        if data_args.max_predict_samples is not None:
            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
        # Set the test transforms
        predict_dataset = predict_dataset.map(
            val_transforms,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    collate_fn = DefaultDataCollator(return_tensors="np")

    # Load the accuracy metric from the datasets package
    metric = evaluate.load("accuracy")

    # 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.
    def compute_metrics(p):
        """Computes accuracy on a batch of predictions"""
        logits, label_ids = p
        predictions = np.argmax(logits, axis=-1)
        metrics = metric.compute(predictions=predictions, references=label_ids)
        return metrics

    with training_args.strategy.scope():
        if checkpoint is None:
            model_path = model_args.model_name_or_path
        else:
            model_path = checkpoint

        model = TFAutoModelForImageClassification.from_pretrained(
            model_path,
            config=config,
            from_pt=bool(".bin" in model_path),
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
            ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
        )
        num_replicas = training_args.strategy.num_replicas_in_sync
        total_train_batch_size = training_args.per_device_train_batch_size * num_replicas
        total_eval_batch_size = training_args.per_device_eval_batch_size * num_replicas

        dataset_options = tf.data.Options()
        dataset_options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF

        if training_args.do_train:
            num_train_steps = int(len(train_dataset) * training_args.num_train_epochs)
            if training_args.warmup_steps > 0:
                num_warmpup_steps = int(training_args.warmup_steps)
            elif training_args.warmup_ratio > 0:
                num_warmpup_steps = int(training_args.warmup_ratio * num_train_steps)
            else:
                num_warmpup_steps = 0

            optimizer, _ = create_optimizer(
                init_lr=training_args.learning_rate,
                num_train_steps=num_train_steps,
                num_warmup_steps=num_warmpup_steps,
                adam_beta1=training_args.adam_beta1,
                adam_beta2=training_args.adam_beta2,
                adam_epsilon=training_args.adam_epsilon,
                weight_decay_rate=training_args.weight_decay,
                adam_global_clipnorm=training_args.max_grad_norm,
            )
            # model.prepare_tf_dataset() wraps a Hugging Face dataset in a tf.data.Dataset which is ready to use in
            # training. This is the recommended way to use a Hugging Face dataset when training with Keras. You can also
            # use the lower-level dataset.to_tf_dataset() method, but you will have to specify things like column names
            # yourself if you use this method, whereas they are automatically inferred from the model input names when
            # using model.prepare_tf_dataset()
            # For more info see the docs:
            # https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.TFPreTrainedModel.prepare_tf_dataset
            # https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.to_tf_dataset
            train_dataset = model.prepare_tf_dataset(
                train_dataset,
                shuffle=True,
                batch_size=total_train_batch_size,
                collate_fn=collate_fn,
            ).with_options(dataset_options)
        else:
            optimizer = None

        if training_args.do_eval:
            eval_dataset = model.prepare_tf_dataset(
                eval_dataset,
                shuffle=False,
                batch_size=total_eval_batch_size,
                collate_fn=collate_fn,
            ).with_options(dataset_options)

        if training_args.do_predict:
            predict_dataset = model.prepare_tf_dataset(
                predict_dataset,
                shuffle=False,
                batch_size=total_eval_batch_size,
                collate_fn=collate_fn,
            ).with_options(dataset_options)

        model.compile(optimizer=optimizer, jit_compile=training_args.xla, metrics=["accuracy"])

        push_to_hub_model_id = training_args.push_to_hub_model_id
        if not push_to_hub_model_id:
            model_name = model_args.model_name_or_path.split("/")[-1]
            push_to_hub_model_id = f"{model_name}-finetuned-image-classification"

        model_card_kwargs = {
            "finetuned_from": model_args.model_name_or_path,
            "tasks": "image-classification",
            "dataset": data_args.dataset_name,
            "tags": ["image-classification", "tensorflow", "vision"],
        }

        callbacks = []
        if eval_dataset is not None:
            callbacks.append(KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=eval_dataset))
        if training_args.push_to_hub:
            callbacks.append(
                PushToHubCallback(
                    output_dir=training_args.output_dir,
                    hub_model_id=push_to_hub_model_id,
                    hub_token=training_args.push_to_hub_token,
                    tokenizer=image_processor,
                    **model_card_kwargs,
                )
            )

        if training_args.do_train:
            model.fit(
                train_dataset,
                validation_data=eval_dataset,
                epochs=int(training_args.num_train_epochs),
                callbacks=callbacks,
            )

        if training_args.do_eval:
            n_eval_batches = len(eval_dataset)
            eval_predictions = model.predict(eval_dataset, steps=n_eval_batches)
            eval_labels = dataset["validation"]["labels"][: n_eval_batches * total_eval_batch_size]
            eval_metrics = compute_metrics((eval_predictions.logits, eval_labels))
            logging.info("Eval metrics:")
            for metric_name, value in eval_metrics.items():
                logging.info(f"{metric_name}: {value:.3f}")

        if training_args.output_dir is not None:
            with open(os.path.join(training_args.output_dir, "all_results.json"), "w") as f:
                f.write(json.dumps(eval_metrics))

        if training_args.do_predict:
            n_predict_batches = len(predict_dataset)
            test_predictions = model.predict(predict_dataset, steps=n_predict_batches)
            test_labels = dataset["validation"]["labels"][: n_predict_batches * total_eval_batch_size]
            test_metrics = compute_metrics((test_predictions.logits, test_labels))
            logging.info("Test metrics:")
            for metric_name, value in test_metrics.items():
                logging.info(f"{metric_name}: {value:.3f}")

        if training_args.output_dir is not None and not training_args.push_to_hub:
            # If we're not pushing to hub, at least save a local copy when we're done
            model.save_pretrained(training_args.output_dir)