import os import albumentations as A import numpy as np import torch from datasets import load_dataset from loguru import logger from sklearn import metrics from transformers import ( AutoConfig, AutoImageProcessor, AutoModelForImageClassification, EarlyStoppingCallback, Trainer, TrainingArguments, ) from autotrain import utils from autotrain.params import ImageBinaryClassificationParams, ImageMultiClassClassificationParams BINARY_CLASSIFICATION_EVAL_METRICS = ( "eval_loss", "eval_accuracy", "eval_f1", "eval_auc", "eval_precision", "eval_recall", ) MULTI_CLASS_CLASSIFICATION_EVAL_METRICS = ( "eval_loss", "eval_accuracy", "eval_f1_macro", "eval_f1_micro", "eval_f1_weighted", "eval_precision_macro", "eval_precision_micro", "eval_precision_weighted", "eval_recall_macro", "eval_recall_micro", "eval_recall_weighted", ) MODEL_CARD = """ --- tags: - autotrain - image-classification widget: - src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg example_title: Tiger - src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/teapot.jpg example_title: Teapot - src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/palace.jpg example_title: Palace datasets: - {dataset} co2_eq_emissions: emissions: {co2} --- # Model Trained Using AutoTrain - Problem type: Image Classification - CO2 Emissions (in grams): {co2:.4f} ## Validation Metricsg {validation_metrics} """ class Dataset: def __init__(self, data, transforms): self.data = data self.transforms = transforms def __len__(self): return len(self.data) def __getitem__(self, item): image = self.data[item]["image"] target = int(self.data[item]["label"]) image = self.transforms(image=np.array(image.convert("RGB")))["image"] image = np.transpose(image, (2, 0, 1)).astype(np.float32) return { "pixel_values": torch.tensor(image, dtype=torch.float), "labels": torch.tensor(target, dtype=torch.long), } def _binary_classification_metrics(pred): raw_predictions, labels = pred predictions = np.argmax(raw_predictions, axis=1) result = { "f1": metrics.f1_score(labels, predictions), "precision": metrics.precision_score(labels, predictions), "recall": metrics.recall_score(labels, predictions), "auc": metrics.roc_auc_score(labels, raw_predictions[:, 1]), "accuracy": metrics.accuracy_score(labels, predictions), } return result def _multi_class_classification_metrics(pred): raw_predictions, labels = pred predictions = np.argmax(raw_predictions, axis=1) results = { "f1_macro": metrics.f1_score(labels, predictions, average="macro"), "f1_micro": metrics.f1_score(labels, predictions, average="micro"), "f1_weighted": metrics.f1_score(labels, predictions, average="weighted"), "precision_macro": metrics.precision_score(labels, predictions, average="macro"), "precision_micro": metrics.precision_score(labels, predictions, average="micro"), "precision_weighted": metrics.precision_score(labels, predictions, average="weighted"), "recall_macro": metrics.recall_score(labels, predictions, average="macro"), "recall_micro": metrics.recall_score(labels, predictions, average="micro"), "recall_weighted": metrics.recall_score(labels, predictions, average="weighted"), "accuracy": metrics.accuracy_score(labels, predictions), } return results def process_data(train_data, valid_data, 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"]) try: height, width = size except TypeError: height = size width = size train_transforms = A.Compose( [ A.RandomResizedCrop(height=height, width=width), A.RandomRotate90(), A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.Normalize(mean=image_processor.image_mean, std=image_processor.image_std), ] ) val_transforms = A.Compose( [ A.Resize(height=height, width=width), A.Normalize(mean=image_processor.image_mean, std=image_processor.image_std), ] ) train_data = Dataset(train_data, train_transforms) valid_data = Dataset(valid_data, val_transforms) return train_data, valid_data @utils.job_watcher def train(co2_tracker, payload, huggingface_token, model_path): # create model repo model_repo = utils.create_repo( project_name=payload["proj_name"], autotrain_user=payload["username"], huggingface_token=huggingface_token, model_path=model_path, ) data_path = f"{payload['username']}/autotrain-data-{payload['proj_name']}" data = load_dataset(data_path, use_auth_token=huggingface_token) logger.info(f"Loaded data from {data_path}") job_config = payload["config"]["params"][0] job_config["model_name"] = payload["config"]["hub_model"] train_data = data["train"] valid_data = data["validation"] labels = train_data.features["label"].names label2id, id2label = {}, {} for i, label in enumerate(labels): label2id[label] = str(i) id2label[str(i)] = label num_classes = len(labels) model_name = job_config["model_name"] device = job_config.get("device", "cuda") # remove model_name from job config del job_config["model_name"] if num_classes == 2: job_config["task"] = "image_binary_classification" job_config = ImageBinaryClassificationParams(**job_config) elif num_classes > 2: job_config["task"] = "image_multi_class_classification" job_config = ImageMultiClassClassificationParams(**job_config) else: raise ValueError("Invalid number of classes") model_config = AutoConfig.from_pretrained( model_name, num_labels=num_classes, use_auth_token=huggingface_token, ) model_config._num_labels = len(label2id) model_config.label2id = label2id model_config.id2label = id2label logger.info(model_config) try: model = AutoModelForImageClassification.from_pretrained( model_name, config=model_config, use_auth_token=huggingface_token, ignore_mismatched_sizes=True, ) except OSError: model = AutoModelForImageClassification.from_pretrained( model_name, config=model_config, use_auth_token=huggingface_token, from_tf=True, ignore_mismatched_sizes=True, ) image_processor = AutoImageProcessor.from_pretrained(model_name, use_auth_token=huggingface_token) train_dataset, valid_dataset = process_data(train_data, valid_data, image_processor) # trainer specific logging_steps = int(0.2 * len(valid_dataset) / job_config.train_batch_size) if logging_steps == 0: logging_steps = 1 fp16 = True if device == "cpu": fp16 = False training_args = dict( output_dir=model_path, per_device_train_batch_size=job_config.train_batch_size, per_device_eval_batch_size=job_config.train_batch_size, learning_rate=job_config.learning_rate, num_train_epochs=job_config.num_train_epochs, fp16=fp16, load_best_model_at_end=True, evaluation_strategy="epoch", logging_steps=logging_steps, save_total_limit=1, save_strategy="epoch", disable_tqdm=not bool(os.environ.get("ENABLE_TQDM", 0)), gradient_accumulation_steps=job_config.gradient_accumulation_steps, report_to="none", auto_find_batch_size=True, lr_scheduler_type=job_config.scheduler, optim=job_config.optimizer, warmup_ratio=job_config.warmup_ratio, weight_decay=job_config.weight_decay, max_grad_norm=job_config.max_grad_norm, ) early_stop = EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.01) callbacks_to_use = [early_stop] args = TrainingArguments(**training_args) trainer_args = dict( args=args, model=model, callbacks=callbacks_to_use, compute_metrics=_binary_classification_metrics if num_classes == 2 else _multi_class_classification_metrics, ) trainer = Trainer( **trainer_args, train_dataset=train_dataset, eval_dataset=valid_dataset, ) trainer.train() logger.info("Finished training") logger.info(trainer.state.best_metric) eval_scores = trainer.evaluate() # create and save model card co2_consumed = co2_tracker.stop() co2_consumed = co2_consumed * 1000 if co2_consumed is not None else 0 valid_metrics = BINARY_CLASSIFICATION_EVAL_METRICS if num_classes == 2 else MULTI_CLASS_CLASSIFICATION_EVAL_METRICS eval_scores = [f"{k[len('eval_'):]}: {v}" for k, v in eval_scores.items() if k in valid_metrics] eval_scores = "\n\n".join(eval_scores) model_card = MODEL_CARD.format( language=payload["config"]["language"], dataset=data_path, co2=co2_consumed, validation_metrics=eval_scores, ) logger.info(model_card) utils.save_model_card(model_card, model_path) # save model, image_processor and config model = utils.update_model_config(trainer.model, job_config) utils.save_tokenizer(image_processor, model_path) utils.save_model(model, model_path) utils.remove_checkpoints(model_path=model_path) # push model to hub logger.info("Pushing model to Hub") model_repo.git_pull() model_repo.git_add() model_repo.git_commit(commit_message="Commit From AutoTrain") model_repo.git_push()