chendl/compositional_test
1
1# coding=utf-82# Copyright 2022 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" Finetuning any 🤗 Transformers model supported by AutoModelForSemanticSegmentation for semantic segmentation."""16 17import argparse18import json19import math20import os21import random22from pathlib import Path23 24import datasets25import evaluate26import numpy as np27import torch28from accelerate import Accelerator29from accelerate.logging import get_logger30from accelerate.utils import set_seed31from datasets import load_dataset32from huggingface_hub import Repository, create_repo, hf_hub_download33from PIL import Image34from torch.utils.data import DataLoader35from torchvision import transforms36from torchvision.transforms import functional37from tqdm.auto import tqdm38 39import transformers40from transformers import (41 AutoConfig,42 AutoImageProcessor,43 AutoModelForSemanticSegmentation,44 SchedulerType,45 default_data_collator,46 get_scheduler,47)48from transformers.utils import check_min_version, get_full_repo_name, send_example_telemetry49from transformers.utils.versions import require_version50 51 52# Will error if the minimal version of Transformers is not installed. Remove at your own risks.53check_min_version("4.28.0")54 55logger = get_logger(__name__)56 57require_version("datasets>=2.0.0", "To fix: pip install -r examples/pytorch/semantic-segmentation/requirements.txt")58 59 60def pad_if_smaller(img, size, fill=0):61 min_size = min(img.size)62 if min_size < size:63 original_width, original_height = img.size64 pad_height = size - original_height if original_height < size else 065 pad_width = size - original_width if original_width < size else 066 img = functional.pad(img, (0, 0, pad_width, pad_height), fill=fill)67 return img68 69 70class Compose:71 def __init__(self, transforms):72 self.transforms = transforms73 74 def __call__(self, image, target):75 for t in self.transforms:76 image, target = t(image, target)77 return image, target78 79 80class Identity:81 def __init__(self):82 pass83 84 def __call__(self, image, target):85 return image, target86 87 88class Resize:89 def __init__(self, size):90 self.size = size91 92 def __call__(self, image, target):93 image = functional.resize(image, self.size)94 target = functional.resize(target, self.size, interpolation=transforms.InterpolationMode.NEAREST)95 return image, target96 97 98class RandomResize:99 def __init__(self, min_size, max_size=None):100 self.min_size = min_size101 if max_size is None:102 max_size = min_size103 self.max_size = max_size104 105 def __call__(self, image, target):106 size = random.randint(self.min_size, self.max_size)107 image = functional.resize(image, size)108 target = functional.resize(target, size, interpolation=transforms.InterpolationMode.NEAREST)109 return image, target110 111 112class RandomCrop:113 def __init__(self, size):114 self.size = size115 116 def __call__(self, image, target):117 image = pad_if_smaller(image, self.size)118 target = pad_if_smaller(target, self.size, fill=255)119 crop_params = transforms.RandomCrop.get_params(image, (self.size, self.size))120 image = functional.crop(image, *crop_params)121 target = functional.crop(target, *crop_params)122 return image, target123 124 125class RandomHorizontalFlip:126 def __init__(self, flip_prob):127 self.flip_prob = flip_prob128 129 def __call__(self, image, target):130 if random.random() < self.flip_prob:131 image = functional.hflip(image)132 target = functional.hflip(target)133 return image, target134 135 136class PILToTensor:137 def __call__(self, image, target):138 image = functional.pil_to_tensor(image)139 target = torch.as_tensor(np.array(target), dtype=torch.int64)140 return image, target141 142 143class ConvertImageDtype:144 def __init__(self, dtype):145 self.dtype = dtype146 147 def __call__(self, image, target):148 image = functional.convert_image_dtype(image, self.dtype)149 return image, target150 151 152class Normalize:153 def __init__(self, mean, std):154 self.mean = mean155 self.std = std156 157 def __call__(self, image, target):158 image = functional.normalize(image, mean=self.mean, std=self.std)159 return image, target160 161 162class ReduceLabels:163 def __call__(self, image, target):164 if not isinstance(target, np.ndarray):165 target = np.array(target).astype(np.uint8)166 # avoid using underflow conversion167 target[target == 0] = 255168 target = target - 1169 target[target == 254] = 255170 171 target = Image.fromarray(target)172 return image, target173 174 175def parse_args():176 parser = argparse.ArgumentParser(description="Finetune a transformers model on a text classification task")177 parser.add_argument(178 "--model_name_or_path",179 type=str,180 help="Path to a pretrained model or model identifier from huggingface.co/models.",181 default="nvidia/mit-b0",182 )183 parser.add_argument(184 "--dataset_name",185 type=str,186 help="Name of the dataset on the hub.",187 default="segments/sidewalk-semantic",188 )189 parser.add_argument(190 "--reduce_labels",191 action="store_true",192 help="Whether or not to reduce all labels by 1 and replace background by 255.",193 )194 parser.add_argument(195 "--train_val_split",196 type=float,197 default=0.15,198 help="Fraction of the dataset to be used for validation.",199 )200 parser.add_argument(201 "--cache_dir",202 type=str,203 help="Path to a folder in which the model and dataset will be cached.",204 )205 parser.add_argument(206 "--use_auth_token",207 action="store_true",208 help="Whether to use an authentication token to access the model repository.",209 )210 parser.add_argument(211 "--per_device_train_batch_size",212 type=int,213 default=8,214 help="Batch size (per device) for the training dataloader.",215 )216 parser.add_argument(217 "--per_device_eval_batch_size",218 type=int,219 default=8,220 help="Batch size (per device) for the evaluation dataloader.",221 )222 parser.add_argument(223 "--learning_rate",224 type=float,225 default=5e-5,226 help="Initial learning rate (after the potential warmup period) to use.",227 )228 parser.add_argument(229 "--adam_beta1",230 type=float,231 default=0.9,232 help="Beta1 for AdamW optimizer",233 )234 parser.add_argument(235 "--adam_beta2",236 type=float,237 default=0.999,238 help="Beta2 for AdamW optimizer",239 )240 parser.add_argument(241 "--adam_epsilon",242 type=float,243 default=1e-8,244 help="Epsilon for AdamW optimizer",245 )246 parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")247 parser.add_argument(248 "--max_train_steps",249 type=int,250 default=None,251 help="Total number of training steps to perform. If provided, overrides num_train_epochs.",252 )253 parser.add_argument(254 "--gradient_accumulation_steps",255 type=int,256 default=1,257 help="Number of updates steps to accumulate before performing a backward/update pass.",258 )259 parser.add_argument(260 "--lr_scheduler_type",261 type=SchedulerType,262 default="polynomial",263 help="The scheduler type to use.",264 choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],265 )266 parser.add_argument(267 "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."268 )269 parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")270 parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")271 parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")272 parser.add_argument(273 "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."274 )275 parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")276 parser.add_argument(277 "--checkpointing_steps",278 type=str,279 default=None,280 help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",281 )282 parser.add_argument(283 "--resume_from_checkpoint",284 type=str,285 default=None,286 help="If the training should continue from a checkpoint folder.",287 )288 parser.add_argument(289 "--with_tracking",290 required=False,291 action="store_true",292 help="Whether to enable experiment trackers for logging.",293 )294 parser.add_argument(295 "--report_to",296 type=str,297 default="all",298 help=(299 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'300 ' `"wandb"`, `"comet_ml"` and `"clearml"`. Use `"all"` (default) to report to all integrations.'301 "Only applicable when `--with_tracking` is passed."302 ),303 )304 args = parser.parse_args()305 306 # Sanity checks307 if args.push_to_hub or args.with_tracking:308 if args.output_dir is None:309 raise ValueError(310 "Need an `output_dir` to create a repo when `--push_to_hub` or `with_tracking` is specified."311 )312 313 if args.output_dir is not None:314 os.makedirs(args.output_dir, exist_ok=True)315 316 return args317 318 319def main():320 args = parse_args()321 322 # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The323 # information sent is the one passed as arguments along with your Python/PyTorch versions.324 send_example_telemetry("run_semantic_segmentation_no_trainer", args)325 326 # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.327 # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers328 # in the environment329 accelerator_log_kwargs = {}330 331 if args.with_tracking:332 accelerator_log_kwargs["log_with"] = args.report_to333 accelerator_log_kwargs["logging_dir"] = args.output_dir334 335 accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, **accelerator_log_kwargs)336 337 logger.info(accelerator.state, main_process_only=False)338 if accelerator.is_local_main_process:339 datasets.utils.logging.set_verbosity_warning()340 transformers.utils.logging.set_verbosity_info()341 else:342 datasets.utils.logging.set_verbosity_error()343 transformers.utils.logging.set_verbosity_error()344 345 # If passed along, set the training seed now.346 # We set device_specific to True as we want different data augmentation per device.347 if args.seed is not None:348 set_seed(args.seed, device_specific=True)349 350 # Handle the repository creation351 if accelerator.is_main_process:352 if args.push_to_hub:353 if args.hub_model_id is None:354 repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)355 else:356 repo_name = args.hub_model_id357 create_repo(repo_name, exist_ok=True, token=args.hub_token)358 repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)359 360 with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:361 if "step_*" not in gitignore:362 gitignore.write("step_*\n")363 if "epoch_*" not in gitignore:364 gitignore.write("epoch_*\n")365 elif args.output_dir is not None:366 os.makedirs(args.output_dir, exist_ok=True)367 accelerator.wait_for_everyone()368 369 # Load dataset370 # In distributed training, the load_dataset function guarantees that only one local process can concurrently371 # download the dataset.372 # TODO support datasets from local folders373 dataset = load_dataset(args.dataset_name, cache_dir=args.cache_dir)374 375 # Rename column names to standardized names (only "image" and "label" need to be present)376 if "pixel_values" in dataset["train"].column_names:377 dataset = dataset.rename_columns({"pixel_values": "image"})378 if "annotation" in dataset["train"].column_names:379 dataset = dataset.rename_columns({"annotation": "label"})380 381 # If we don't have a validation split, split off a percentage of train as validation.382 args.train_val_split = None if "validation" in dataset.keys() else args.train_val_split383 if isinstance(args.train_val_split, float) and args.train_val_split > 0.0:384 split = dataset["train"].train_test_split(args.train_val_split)385 dataset["train"] = split["train"]386 dataset["validation"] = split["test"]387 388 # Prepare label mappings.389 # We'll include these in the model's config to get human readable labels in the Inference API.390 if args.dataset_name == "scene_parse_150":391 repo_id = "huggingface/label-files"392 filename = "ade20k-id2label.json"393 else:394 repo_id = args.dataset_name395 filename = "id2label.json"396 id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))397 id2label = {int(k): v for k, v in id2label.items()}398 label2id = {v: k for k, v in id2label.items()}399 400 # Load pretrained model and image processor401 config = AutoConfig.from_pretrained(args.model_name_or_path, id2label=id2label, label2id=label2id)402 image_processor = AutoImageProcessor.from_pretrained(args.model_name_or_path)403 model = AutoModelForSemanticSegmentation.from_pretrained(args.model_name_or_path, config=config)404 405 # Preprocessing the datasets406 # Define torchvision transforms to be applied to each image + target.407 # Not that straightforward in torchvision: https://github.com/pytorch/vision/issues/9408 # Currently based on official torchvision references: https://github.com/pytorch/vision/blob/main/references/segmentation/transforms.py409 if "shortest_edge" in image_processor.size:410 # We instead set the target size as (shortest_edge, shortest_edge) to here to ensure all images are batchable.411 size = (image_processor.size["shortest_edge"], image_processor.size["shortest_edge"])412 else:413 size = (image_processor.size["height"], image_processor.size["width"])414 train_transforms = Compose(415 [416 ReduceLabels() if args.reduce_labels else Identity(),417 RandomCrop(size=size),418 RandomHorizontalFlip(flip_prob=0.5),419 PILToTensor(),420 ConvertImageDtype(torch.float),421 Normalize(mean=image_processor.image_mean, std=image_processor.image_std),422 ]423 )424 # Define torchvision transform to be applied to each image.425 # jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)426 val_transforms = Compose(427 [428 ReduceLabels() if args.reduce_labels else Identity(),429 Resize(size=size),430 PILToTensor(),431 ConvertImageDtype(torch.float),432 Normalize(mean=image_processor.image_mean, std=image_processor.image_std),433 ]434 )435 436 def preprocess_train(example_batch):437 pixel_values = []438 labels = []439 for image, target in zip(example_batch["image"], example_batch["label"]):440 image, target = train_transforms(image.convert("RGB"), target)441 pixel_values.append(image)442 labels.append(target)443 444 encoding = {}445 encoding["pixel_values"] = torch.stack(pixel_values)446 encoding["labels"] = torch.stack(labels)447 448 return encoding449 450 def preprocess_val(example_batch):451 pixel_values = []452 labels = []453 for image, target in zip(example_batch["image"], example_batch["label"]):454 image, target = val_transforms(image.convert("RGB"), target)455 pixel_values.append(image)456 labels.append(target)457 458 encoding = {}459 encoding["pixel_values"] = torch.stack(pixel_values)460 encoding["labels"] = torch.stack(labels)461 462 return encoding463 464 with accelerator.main_process_first():465 train_dataset = dataset["train"].with_transform(preprocess_train)466 eval_dataset = dataset["validation"].with_transform(preprocess_val)467 468 train_dataloader = DataLoader(469 train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size470 )471 eval_dataloader = DataLoader(472 eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size473 )474 475 # Optimizer476 optimizer = torch.optim.AdamW(477 list(model.parameters()),478 lr=args.learning_rate,479 betas=[args.adam_beta1, args.adam_beta2],480 eps=args.adam_epsilon,481 )482 483 # Figure out how many steps we should save the Accelerator states484 checkpointing_steps = args.checkpointing_steps485 if checkpointing_steps is not None and checkpointing_steps.isdigit():486 checkpointing_steps = int(checkpointing_steps)487 488 # Scheduler and math around the number of training steps.489 overrode_max_train_steps = False490 num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)491 if args.max_train_steps is None:492 args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch493 overrode_max_train_steps = True494 495 lr_scheduler = get_scheduler(496 name=args.lr_scheduler_type,497 optimizer=optimizer,498 num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,499 num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,500 )501 502 # Prepare everything with our `accelerator`.503 model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(504 model, optimizer, train_dataloader, eval_dataloader, lr_scheduler505 )506 507 # We need to recalculate our total training steps as the size of the training dataloader may have changed.508 num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)509 if overrode_max_train_steps:510 args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch511 # Afterwards we recalculate our number of training epochs512 args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)513 514 # Instantiate metric515 metric = evaluate.load("mean_iou")516 517 # We need to initialize the trackers we use, and also store our configuration.518 # The trackers initializes automatically on the main process.519 if args.with_tracking:520 experiment_config = vars(args)521 # TensorBoard cannot log Enums, need the raw value522 experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value523 accelerator.init_trackers("semantic_segmentation_no_trainer", experiment_config)524 525 # Train!526 total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps527 528 logger.info("***** Running training *****")529 logger.info(f" Num examples = {len(train_dataset)}")530 logger.info(f" Num Epochs = {args.num_train_epochs}")531 logger.info(f" Instantaneous batch size per device = {args.per_device_train_batch_size}")532 logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")533 logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")534 logger.info(f" Total optimization steps = {args.max_train_steps}")535 # Only show the progress bar once on each machine.536 progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)537 completed_steps = 0538 starting_epoch = 0539 540 # Potentially load in the weights and states from a previous save541 if args.resume_from_checkpoint:542 if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":543 accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")544 accelerator.load_state(args.resume_from_checkpoint)545 path = os.path.basename(args.resume_from_checkpoint)546 else:547 # Get the most recent checkpoint548 dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]549 dirs.sort(key=os.path.getctime)550 path = dirs[-1] # Sorts folders by date modified, most recent checkpoint is the last551 # Extract `epoch_{i}` or `step_{i}`552 training_difference = os.path.splitext(path)[0]553 554 if "epoch" in training_difference:555 starting_epoch = int(training_difference.replace("epoch_", "")) + 1556 resume_step = None557 else:558 resume_step = int(training_difference.replace("step_", ""))559 starting_epoch = resume_step // len(train_dataloader)560 resume_step -= starting_epoch * len(train_dataloader)561 562 for epoch in range(starting_epoch, args.num_train_epochs):563 if args.with_tracking:564 total_loss = 0565 model.train()566 for step, batch in enumerate(train_dataloader):567 # We need to skip steps until we reach the resumed step568 if args.resume_from_checkpoint and epoch == starting_epoch:569 if resume_step is not None and step < resume_step:570 completed_steps += 1571 continue572 573 with accelerator.accumulate(model):574 outputs = model(**batch)575 loss = outputs.loss576 # We keep track of the loss at each epoch577 if args.with_tracking:578 total_loss += loss.detach().float()579 accelerator.backward(loss)580 optimizer.step()581 lr_scheduler.step()582 optimizer.zero_grad()583 584 # Checks if the accelerator has performed an optimization step behind the scenes585 if accelerator.sync_gradients:586 progress_bar.update(1)587 completed_steps += 1588 589 if isinstance(checkpointing_steps, int):590 if completed_steps % checkpointing_steps == 0:591 output_dir = f"step_{completed_steps }"592 if args.output_dir is not None:593 output_dir = os.path.join(args.output_dir, output_dir)594 accelerator.save_state(output_dir)595 596 if args.push_to_hub and epoch < args.num_train_epochs - 1:597 accelerator.wait_for_everyone()598 unwrapped_model = accelerator.unwrap_model(model)599 unwrapped_model.save_pretrained(600 args.output_dir,601 is_main_process=accelerator.is_main_process,602 save_function=accelerator.save,603 )604 if accelerator.is_main_process:605 image_processor.save_pretrained(args.output_dir)606 repo.push_to_hub(607 commit_message=f"Training in progress {completed_steps} steps",608 blocking=False,609 auto_lfs_prune=True,610 )611 612 if completed_steps >= args.max_train_steps:613 break614 615 logger.info("***** Running evaluation *****")616 model.eval()617 for step, batch in enumerate(tqdm(eval_dataloader, disable=not accelerator.is_local_main_process)):618 with torch.no_grad():619 outputs = model(**batch)620 621 upsampled_logits = torch.nn.functional.interpolate(622 outputs.logits, size=batch["labels"].shape[-2:], mode="bilinear", align_corners=False623 )624 predictions = upsampled_logits.argmax(dim=1)625 626 predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))627 628 metric.add_batch(629 predictions=predictions,630 references=references,631 )632 633 eval_metrics = metric.compute(634 num_labels=len(id2label),635 ignore_index=255,636 reduce_labels=False, # we've already reduced the labels before637 )638 logger.info(f"epoch {epoch}: {eval_metrics}")639 640 if args.with_tracking:641 accelerator.log(642 {643 "mean_iou": eval_metrics["mean_iou"],644 "mean_accuracy": eval_metrics["mean_accuracy"],645 "overall_accuracy": eval_metrics["overall_accuracy"],646 "train_loss": total_loss.item() / len(train_dataloader),647 "epoch": epoch,648 "step": completed_steps,649 },650 step=completed_steps,651 )652 653 if args.push_to_hub and epoch < args.num_train_epochs - 1:654 accelerator.wait_for_everyone()655 unwrapped_model = accelerator.unwrap_model(model)656 unwrapped_model.save_pretrained(657 args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save658 )659 if accelerator.is_main_process:660 image_processor.save_pretrained(args.output_dir)661 repo.push_to_hub(662 commit_message=f"Training in progress epoch {epoch}", blocking=False, auto_lfs_prune=True663 )664 665 if args.checkpointing_steps == "epoch":666 output_dir = f"epoch_{epoch}"667 if args.output_dir is not None:668 output_dir = os.path.join(args.output_dir, output_dir)669 accelerator.save_state(output_dir)670 671 if args.with_tracking:672 accelerator.end_training()673 674 if args.output_dir is not None:675 accelerator.wait_for_everyone()676 unwrapped_model = accelerator.unwrap_model(model)677 unwrapped_model.save_pretrained(678 args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save679 )680 if accelerator.is_main_process:681 image_processor.save_pretrained(args.output_dir)682 if args.push_to_hub:683 repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)684 685 all_results = {f"eval_{k}": v for k, v in eval_metrics.items()}686 with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:687 json.dump(all_results, f)688 689 690if __name__ == "__main__":691 main()692