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 limitations under the license finetuning any transformers model supported by automodelforsemanticsegmentation for semantic segmentation import argparse import json import math import os import random from pathlib import path import datasets import evaluate import numpy as np import torch from accelerate import accelerator from accelerate logging import getlogger from accelerate utils import setseed from datasets import loaddataset from huggingfacehub import repository createrepo hfhubdownload from pil import image from torch utils data import dataloader from torchvision import transforms from torchvision transforms import functional from tqdm auto import tqdm import transformers from transformers import autoconfig autoimageprocessor automodelforsemanticsegmentation schedulertype defaultdatacollator getscheduler from transformers utils import checkminversion sendexampletelemetry from transformers utils versions import requireversion will error if the minimal version of transformers is not installed remove at your own risks checkminversion4 36 0 dev0 logger getloggername requireversiondatasets2 0 0 to fix pip install r examplespytorchsemanticsegmentationrequirements txt def padifsmallerimg size fill0 minsize minimg size if minsize size originalwidth originalheight img size padheight size originalheight if originalheight size else 0 padwidth size originalwidth if originalwidth size else 0 img functional padimg 0 0 padwidth padheight fillfill return img class compose def initself transforms self transforms transforms def callself image target for t in self transforms image target timage target return image target class identity def initself pass def callself image target return image target class resize def initself size self size size def callself image target image functional resizeimage self size target functional resizetarget self size interpolationtransforms interpolationmode nearest return image target class randomresize def initself minsize maxsizenone self minsize minsize if maxsize is none maxsize minsize self maxsize maxsize def callself image target size random randintself minsize self maxsize image functional resizeimage size target functional resizetarget size interpolationtransforms interpolationmode nearest return image target class randomcrop def initself size self size size def callself image target image padifsmallerimage self size target padifsmallertarget self size fill255 cropparams transforms randomcrop getparamsimage self size self size image functional cropimage cropparams target functional croptarget cropparams return image target class randomhorizontalflip def initself flipprob self flipprob flipprob def callself image target if random random self flipprob image functional hflipimage target functional hfliptarget return image target class piltotensor def callself image target image functional piltotensorimage target torch astensornp arraytarget dtypetorch int64 return image target class convertimagedtype def initself dtype self dtype dtype def callself image target image functional convertimagedtypeimage self dtype return image target class normalize def initself mean std self mean mean self std std def callself image target image functional normalizeimage meanself mean stdself std return image target class reducelabels def callself image target if not isinstancetarget np ndarray target np arraytarget astypenp uint8 avoid using underflow conversion targettarget 0 255 target target 1 targettarget 254 255 target image fromarraytarget return image target def parseargs parser argparse argumentparserdescriptionfinetune a transformers model on a text classification task parser addargument modelnameorpath typestr helppath to a pretrained model or model identifier from huggingface comodels defaultnvidiamitb0 parser addargument datasetname typestr helpname of the dataset on the hub defaultsegmentssidewalksemantic parser addargument reducelabels actionstoretrue helpwhether or not to reduce all labels by 1 and replace background by 255 parser addargument trainvalsplit typefloat default0 15 helpfraction of the dataset to be used for validation parser addargument cachedir typestr helppath to a folder in which the model and dataset will be cached parser addargument useauthtoken actionstoretrue helpwhether to use an authentication token to access the model repository parser addargument perdevicetrainbatchsize typeint default8 helpbatch size per device for the training dataloader parser addargument perdeviceevalbatchsize typeint default8 helpbatch size per device for the evaluation dataloader parser addargument learningrate typefloat default5e5 helpinitial learning rate after the potential warmup period to use parser addargument adambeta1 typefloat default0 9 helpbeta1 for adamw optimizer parser addargument adambeta2 typefloat default0 999 helpbeta2 for adamw optimizer parser addargument adamepsilon typefloat default1e8 helpepsilon for adamw optimizer parser addargumentnumtrainepochs typeint default3 helptotal number of training epochs to perform parser addargument maxtrainsteps typeint defaultnone helptotal number of training steps to perform if provided overrides numtrainepochs parser addargument gradientaccumulationsteps typeint default1 helpnumber of updates steps to accumulate before performing a backwardupdate pass parser addargument lrschedulertype typeschedulertype defaultpolynomial helpthe scheduler type to use choiceslinear cosine cosinewithrestarts polynomial constant constantwithwarmup parser addargument numwarmupsteps typeint default0 helpnumber of steps for the warmup in the lr scheduler parser addargumentoutputdir typestr defaultnone helpwhere to store the final model parser addargumentseed typeint defaultnone helpa seed for reproducible training parser addargumentpushtohub actionstoretrue helpwhether or not to push the model to the hub parser addargument hubmodelid typestr helpthe name of the repository to keep in sync with the local outputdir parser addargumenthubtoken typestr helpthe token to use to push to the model hub parser addargument trustremotecode typebool defaultfalse 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 parser addargument checkpointingsteps typestr defaultnone helpwhether the various states should be saved at the end of every n steps or epoch for each epoch parser addargument resumefromcheckpoint typestr defaultnone helpif the training should continue from a checkpoint folder parser addargument withtracking requiredfalse actionstoretrue helpwhether to enable experiment trackers for logging parser addargument reportto typestr defaultall help the integration to report the results and logs to supported platforms are tensorboard wandb cometml and clearml use all default to report to all integrations only applicable when withtracking is passed args parser parseargs sanity checks if args pushtohub or args withtracking if args outputdir is none raise valueerror need an outputdir to create a repo when pushtohub or withtracking is specified if args outputdir is not none os makedirsargs outputdir existoktrue return args def main args parseargs 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 sendexampletelemetryrunsemanticsegmentationnotrainer args initialize the accelerator we will let the accelerator handle device placement for us in this example if we re using tracking we also need to initialize it here and it will by default pick up all supported trackers in the environment acceleratorlogkwargs if args withtracking acceleratorlogkwargslogwith args reportto acceleratorlogkwargsprojectdir args outputdir accelerator acceleratorgradientaccumulationstepsargs gradientaccumulationsteps acceleratorlogkwargs logger infoaccelerator state mainprocessonlyfalse if accelerator islocalmainprocess datasets utils logging setverbositywarning transformers utils logging setverbosityinfo else datasets utils logging setverbosityerror transformers utils logging setverbosityerror if passed along set the training seed now we set devicespecific to true as we want different data augmentation per device if args seed is not none setseedargs seed devicespecifictrue handle the repository creation if accelerator ismainprocess if args pushtohub retrieve of infer reponame reponame args hubmodelid if reponame is none reponame pathargs outputdir absolute name create repo and retrieve repoid repoid createreporeponame existoktrue tokenargs hubtoken repoid clone repo locally repo repositoryargs outputdir clonefromrepoid tokenargs hubtoken with openos path joinargs outputdir gitignore w as gitignore if step not in gitignore gitignore writestepn if epoch not in gitignore gitignore writeepochn elif args outputdir is not none os makedirsargs outputdir existoktrue accelerator waitforeveryone load dataset in distributed training the loaddataset function guarantees that only one local process can concurrently download the dataset todo support datasets from local folders dataset loaddatasetargs datasetname cachedirargs cachedir rename column names to standardized names only image and label need to be present if pixelvalues in datasettrain columnnames dataset dataset renamecolumnspixelvalues image if annotation in datasettrain columnnames dataset dataset renamecolumnsannotation label if we don t have a validation split split off a percentage of train as validation args trainvalsplit none if validation in dataset keys else args trainvalsplit if isinstanceargs trainvalsplit float and args trainvalsplit 0 0 split datasettrain traintestsplitargs 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 if args datasetname sceneparse150 repoid huggingfacelabelfiles filename ade20kid2label json else repoid args datasetname filename id2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items label2id v k for k v in id2label items load pretrained model and image processor config autoconfig frompretrained args modelnameorpath id2labelid2label label2idlabel2id trustremotecodeargs trustremotecode imageprocessor autoimageprocessor frompretrained args modelnameorpath trustremotecodeargs trustremotecode model automodelforsemanticsegmentation frompretrained args modelnameorpath configconfig trustremotecodeargs trustremotecode preprocessing the datasets define torchvision transforms to be applied to each image target not that straightforward in torchvision https github compytorchvisionissues9 currently based on official torchvision references https github compytorchvisionblobmainreferencessegmentationtransforms py if shortestedge in imageprocessor size we instead set the target size as shortestedge shortestedge to here to ensure all images are batchable size imageprocessor sizeshortestedge imageprocessor sizeshortestedge else size imageprocessor sizeheight imageprocessor sizewidth traintransforms compose reducelabels if args reducelabels else identity randomcropsizesize randomhorizontalflipflipprob0 5 piltotensor convertimagedtypetorch float normalizemeanimageprocessor imagemean stdimageprocessor imagestd define torchvision transform to be applied to each image jitter colorjitterbrightness0 25 contrast0 25 saturation0 25 hue0 1 valtransforms compose reducelabels if args reducelabels else identity resizesizesize piltotensor convertimagedtypetorch float normalizemeanimageprocessor imagemean stdimageprocessor imagestd def preprocesstrainexamplebatch pixelvalues labels for image target in zipexamplebatchimage examplebatchlabel image target traintransformsimage convertrgb target pixelvalues appendimage labels appendtarget encoding encodingpixelvalues torch stackpixelvalues encodinglabels torch stacklabels return encoding def preprocessvalexamplebatch pixelvalues labels for image target in zipexamplebatchimage examplebatchlabel image target valtransformsimage convertrgb target pixelvalues appendimage labels appendtarget encoding encodingpixelvalues torch stackpixelvalues encodinglabels torch stacklabels return encoding with accelerator mainprocessfirst traindataset datasettrain withtransformpreprocesstrain evaldataset datasetvalidation withtransformpreprocessval traindataloader dataloader traindataset shuffletrue collatefndefaultdatacollator batchsizeargs perdevicetrainbatchsize evaldataloader dataloader evaldataset collatefndefaultdatacollator batchsizeargs perdeviceevalbatchsize optimizer optimizer torch optim adamw listmodel parameters lrargs learningrate betasargs adambeta1 args adambeta2 epsargs adamepsilon figure out how many steps we should save the accelerator states checkpointingsteps args checkpointingsteps if checkpointingsteps is not none and checkpointingsteps isdigit checkpointingsteps intcheckpointingsteps scheduler and math around the number of training steps overrodemaxtrainsteps false numupdatestepsperepoch math ceillentraindataloader args gradientaccumulationsteps if args maxtrainsteps is none args maxtrainsteps args numtrainepochs numupdatestepsperepoch overrodemaxtrainsteps true lrscheduler getscheduler nameargs lrschedulertype optimizeroptimizer numwarmupstepsargs numwarmupsteps args gradientaccumulationsteps numtrainingstepsargs maxtrainsteps args gradientaccumulationsteps prepare everything with our accelerator model optimizer traindataloader evaldataloader lrscheduler accelerator prepare model optimizer traindataloader evaldataloader lrscheduler we need to recalculate our total training steps as the size of the training dataloader may have changed numupdatestepsperepoch math ceillentraindataloader args gradientaccumulationsteps if overrodemaxtrainsteps args maxtrainsteps args numtrainepochs numupdatestepsperepoch afterwards we recalculate our number of training epochs args numtrainepochs math ceilargs maxtrainsteps numupdatestepsperepoch instantiate metric metric evaluate loadmeaniou we need to initialize the trackers we use and also store our configuration the trackers initializes automatically on the main process if args withtracking experimentconfig varsargs tensorboard cannot log enums need the raw value experimentconfiglrschedulertype experimentconfiglrschedulertype value accelerator inittrackerssemanticsegmentationnotrainer experimentconfig train totalbatchsize args perdevicetrainbatchsize accelerator numprocesses args gradientaccumulationsteps logger info running training logger infof num examples lentraindataset logger infof num epochs args numtrainepochs logger infof instantaneous batch size per device args perdevicetrainbatchsize logger infof total train batch size w parallel distributed accumulation totalbatchsize logger infof gradient accumulation steps args gradientaccumulationsteps logger infof total optimization steps args maxtrainsteps only show the progress bar once on each machine progressbar tqdmrangeargs maxtrainsteps disablenot accelerator islocalmainprocess completedsteps 0 startingepoch 0 potentially load in the weights and states from a previous save if args resumefromcheckpoint if args resumefromcheckpoint is not none or args resumefromcheckpoint checkpointpath args resumefromcheckpoint path os path basenameargs resumefromcheckpoint else get the most recent checkpoint dirs f name for f in os scandiros getcwd if f isdir dirs sortkeyos path getctime path dirs1 sorts folders by date modified most recent checkpoint is the last checkpointpath path path os path basenamecheckpointpath accelerator printfresumed from checkpoint checkpointpath accelerator loadstatecheckpointpath extract epochi or stepi trainingdifference os path splitextpath0 if epoch in trainingdifference startingepoch inttrainingdifference replaceepoch 1 resumestep none completedsteps startingepoch numupdatestepsperepoch else need to multiply gradientaccumulationsteps to reflect real steps resumestep inttrainingdifference replacestep args gradientaccumulationsteps startingepoch resumestep lentraindataloader completedsteps resumestep args gradientaccumulationsteps resumestep startingepoch lentraindataloader update the progressbar if load from checkpoint progressbar updatecompletedsteps for epoch in rangestartingepoch args numtrainepochs model train if args withtracking totalloss 0 if args resumefromcheckpoint and epoch startingepoch and resumestep is not none we skip the first n batches in the dataloader when resuming from a checkpoint activedataloader accelerator skipfirstbatchestraindataloader resumestep else activedataloader traindataloader for step batch in enumerateactivedataloader with accelerator accumulatemodel outputs modelbatch loss outputs loss we keep track of the loss at each epoch if args withtracking totalloss loss detach float accelerator backwardloss optimizer step lrscheduler step optimizer zerograd checks if the accelerator has performed an optimization step behind the scenes if accelerator syncgradients progressbar update1 completedsteps 1 if isinstancecheckpointingsteps int if completedsteps checkpointingsteps 0 outputdir fstepcompletedsteps if args outputdir is not none outputdir os path joinargs outputdir outputdir accelerator savestateoutputdir if args pushtohub and epoch args numtrainepochs 1 accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if accelerator ismainprocess imageprocessor savepretrainedargs outputdir repo pushtohub commitmessageftraining in progress completedsteps steps blockingfalse autolfsprunetrue if completedsteps args maxtrainsteps break logger info running evaluation model eval for step batch in enumeratetqdmevaldataloader disablenot accelerator islocalmainprocess with torch nograd outputs modelbatch upsampledlogits torch nn functional interpolate outputs logits sizebatchlabels shape2 modebilinear aligncornersfalse predictions upsampledlogits argmaxdim1 predictions references accelerator gatherformetricspredictions batchlabels metric addbatch predictionspredictions referencesreferences evalmetrics metric compute numlabelslenid2label ignoreindex255 reducelabelsfalse we ve already reduced the labels before logger infofepoch epoch evalmetrics if args withtracking accelerator log meaniou evalmetricsmeaniou meanaccuracy evalmetricsmeanaccuracy overallaccuracy evalmetricsoverallaccuracy trainloss totalloss item lentraindataloader epoch epoch step completedsteps stepcompletedsteps if args pushtohub and epoch args numtrainepochs 1 accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if accelerator ismainprocess imageprocessor savepretrainedargs outputdir repo pushtohub commitmessageftraining in progress epoch epoch blockingfalse autolfsprunetrue if args checkpointingsteps epoch outputdir fepochepoch if args outputdir is not none outputdir os path joinargs outputdir outputdir accelerator savestateoutputdir if args withtracking accelerator endtraining if args outputdir is not none accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if accelerator ismainprocess imageprocessor savepretrainedargs outputdir if args pushtohub repo pushtohubcommitmessageend of training autolfsprunetrue allresults fevalk v tolist if isinstancev np ndarray else v for k v in evalmetrics items with openos path joinargs outputdir allresults json w as f json dumpallresults f if name main main 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 limitations under the license finetuning any transformers model supported by automodelforsemanticsegmentation for semantic segmentation will error if the minimal version of transformers is not installed remove at your own risks avoid using underflow conversion sanity checks 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 initialize the accelerator we will let the accelerator handle device placement for us in this example if we re using tracking we also need to initialize it here and it will by default pick up all supported trackers in the environment if passed along set the training seed now we set device_specific to true as we want different data augmentation per device handle the repository creation retrieve of infer repo_name create repo and retrieve repo_id clone repo locally load dataset in distributed training the load_dataset function guarantees that only one local process can concurrently download the dataset todo support datasets from local folders rename column names to standardized names only image and label need to be present 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 pretrained model and image processor preprocessing the datasets define torchvision transforms to be applied to each image target not that straightforward in torchvision https github com pytorch vision issues 9 currently based on official torchvision references https github com pytorch vision blob main references segmentation transforms py we instead set the target size as shortest_edge shortest_edge to here to ensure all images are batchable define torchvision transform to be applied to each image jitter colorjitter brightness 0 25 contrast 0 25 saturation 0 25 hue 0 1 optimizer figure out how many steps we should save the accelerator states scheduler and math around the number of training steps prepare everything with our accelerator we need to recalculate our total training steps as the size of the training dataloader may have changed afterwards we recalculate our number of training epochs instantiate metric we need to initialize the trackers we use and also store our configuration the trackers initializes automatically on the main process tensorboard cannot log enums need the raw value train only show the progress bar once on each machine potentially load in the weights and states from a previous save get the most recent checkpoint sorts folders by date modified most recent checkpoint is the last extract epoch_ i or step_ i need to multiply gradient_accumulation_steps to reflect real steps update the progress_bar if load from checkpoint we skip the first n batches in the dataloader when resuming from a checkpoint we keep track of the loss at each epoch checks if the accelerator has performed an optimization step behind the scenes we ve already reduced the labels before
import argparse
import json
import math
import os
import random
from pathlib import Path

import datasets
import evaluate
import numpy as np
import torch
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from huggingface_hub import Repository, create_repo, hf_hub_download
from PIL import Image
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.transforms import functional
from tqdm.auto import tqdm

import transformers
from transformers import (
    AutoConfig,
    AutoImageProcessor,
    AutoModelForSemanticSegmentation,
    SchedulerType,
    default_data_collator,
    get_scheduler,
)
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version



check_min_version("4.36.0.dev0")

logger = get_logger(__name__)

require_version("datasets>=2.0.0", "To fix: pip install -r examples/pytorch/semantic-segmentation/requirements.txt")


def pad_if_smaller(img, size, fill=0):
    min_size = min(img.size)
    if min_size < size:
        original_width, original_height = img.size
        pad_height = size - original_height if original_height < size else 0
        pad_width = size - original_width if original_width < size else 0
        img = functional.pad(img, (0, 0, pad_width, pad_height), fill=fill)
    return img


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class Identity:
    def __init__(self):
        pass

    def __call__(self, image, target):
        return image, target


class Resize:
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = functional.resize(image, self.size)
        target = functional.resize(target, self.size, interpolation=transforms.InterpolationMode.NEAREST)
        return image, target


class RandomResize:
    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        if max_size is None:
            max_size = min_size
        self.max_size = max_size

    def __call__(self, image, target):
        size = random.randint(self.min_size, self.max_size)
        image = functional.resize(image, size)
        target = functional.resize(target, size, interpolation=transforms.InterpolationMode.NEAREST)
        return image, target


class RandomCrop:
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = pad_if_smaller(image, self.size)
        target = pad_if_smaller(target, self.size, fill=255)
        crop_params = transforms.RandomCrop.get_params(image, (self.size, self.size))
        image = functional.crop(image, *crop_params)
        target = functional.crop(target, *crop_params)
        return image, target


class RandomHorizontalFlip:
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = functional.hflip(image)
            target = functional.hflip(target)
        return image, target


class PILToTensor:
    def __call__(self, image, target):
        image = functional.pil_to_tensor(image)
        target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target


class ConvertImageDtype:
    def __init__(self, dtype):
        self.dtype = dtype

    def __call__(self, image, target):
        image = functional.convert_image_dtype(image, self.dtype)
        return image, target


class Normalize:
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = functional.normalize(image, mean=self.mean, std=self.std)
        return image, target


class ReduceLabels:
    def __call__(self, image, target):
        if not isinstance(target, np.ndarray):
            target = np.array(target).astype(np.uint8)
        
        target[target == 0] = 255
        target = target - 1
        target[target == 254] = 255

        target = Image.fromarray(target)
        return image, target


def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a transformers model on a text classification task")
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to a pretrained model or model identifier from huggingface.co/models.",
        default="nvidia/mit-b0",
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        help="Name of the dataset on the hub.",
        default="segments/sidewalk-semantic",
    )
    parser.add_argument(
        "--reduce_labels",
        action="store_true",
        help="Whether or not to reduce all labels by 1 and replace background by 255.",
    )
    parser.add_argument(
        "--train_val_split",
        type=float,
        default=0.15,
        help="Fraction of the dataset to be used for validation.",
    )
    parser.add_argument(
        "--cache_dir",
        type=str,
        help="Path to a folder in which the model and dataset will be cached.",
    )
    parser.add_argument(
        "--use_auth_token",
        action="store_true",
        help="Whether to use an authentication token to access the model repository.",
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the evaluation dataloader.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-5,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--adam_beta1",
        type=float,
        default=0.9,
        help="Beta1 for AdamW optimizer",
    )
    parser.add_argument(
        "--adam_beta2",
        type=float,
        default=0.999,
        help="Beta2 for AdamW optimizer",
    )
    parser.add_argument(
        "--adam_epsilon",
        type=float,
        default=1e-8,
        help="Epsilon for AdamW optimizer",
    )
    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="polynomial",
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )
    parser.add_argument(
        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument(
        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
    )
    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
    parser.add_argument(
        "--trust_remote_code",
        type=bool,
        default=False,
        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."
        ),
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=str,
        default=None,
        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help="If the training should continue from a checkpoint folder.",
    )
    parser.add_argument(
        "--with_tracking",
        required=False,
        action="store_true",
        help="Whether to enable experiment trackers for logging.",
    )
    parser.add_argument(
        "--report_to",
        type=str,
        default="all",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
            ' `"wandb"`, `"comet_ml"` and `"clearml"`. Use `"all"` (default) to report to all integrations. '
            "Only applicable when `--with_tracking` is passed."
        ),
    )
    args = parser.parse_args()

    
    if args.push_to_hub or args.with_tracking:
        if args.output_dir is None:
            raise ValueError(
                "Need an `output_dir` to create a repo when `--push_to_hub` or `with_tracking` is specified."
            )

    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)

    return args


def main():
    args = parse_args()

    
    
    send_example_telemetry("run_semantic_segmentation_no_trainer", args)

    
    
    
    accelerator_log_kwargs = {}

    if args.with_tracking:
        accelerator_log_kwargs["log_with"] = args.report_to
        accelerator_log_kwargs["project_dir"] = args.output_dir

    accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, **accelerator_log_kwargs)

    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    
    
    if args.seed is not None:
        set_seed(args.seed, device_specific=True)

    
    if accelerator.is_main_process:
        if args.push_to_hub:
            
            repo_name = args.hub_model_id
            if repo_name is None:
                repo_name = Path(args.output_dir).absolute().name
            
            repo_id = create_repo(repo_name, exist_ok=True, token=args.hub_token).repo_id
            
            repo = Repository(args.output_dir, clone_from=repo_id, token=args.hub_token)

            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
                if "step_*" not in gitignore:
                    gitignore.write("step_*\n")
                if "epoch_*" not in gitignore:
                    gitignore.write("epoch_*\n")
        elif args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    
    
    
    
    dataset = load_dataset(args.dataset_name, cache_dir=args.cache_dir)

    
    if "pixel_values" in dataset["train"].column_names:
        dataset = dataset.rename_columns({"pixel_values": "image"})
    if "annotation" in dataset["train"].column_names:
        dataset = dataset.rename_columns({"annotation": "label"})

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

    
    
    if args.dataset_name == "scene_parse_150":
        repo_id = "huggingface/label-files"
        filename = "ade20k-id2label.json"
    else:
        repo_id = args.dataset_name
        filename = "id2label.json"
    id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
    id2label = {int(k): v for k, v in id2label.items()}
    label2id = {v: k for k, v in id2label.items()}

    
    config = AutoConfig.from_pretrained(
        args.model_name_or_path, id2label=id2label, label2id=label2id, trust_remote_code=args.trust_remote_code
    )
    image_processor = AutoImageProcessor.from_pretrained(
        args.model_name_or_path, trust_remote_code=args.trust_remote_code
    )
    model = AutoModelForSemanticSegmentation.from_pretrained(
        args.model_name_or_path, config=config, trust_remote_code=args.trust_remote_code
    )

    
    
    
    
    if "shortest_edge" in image_processor.size:
        
        size = (image_processor.size["shortest_edge"], image_processor.size["shortest_edge"])
    else:
        size = (image_processor.size["height"], image_processor.size["width"])
    train_transforms = Compose(
        [
            ReduceLabels() if args.reduce_labels else Identity(),
            RandomCrop(size=size),
            RandomHorizontalFlip(flip_prob=0.5),
            PILToTensor(),
            ConvertImageDtype(torch.float),
            Normalize(mean=image_processor.image_mean, std=image_processor.image_std),
        ]
    )
    
    
    val_transforms = Compose(
        [
            ReduceLabels() if args.reduce_labels else Identity(),
            Resize(size=size),
            PILToTensor(),
            ConvertImageDtype(torch.float),
            Normalize(mean=image_processor.image_mean, std=image_processor.image_std),
        ]
    )

    def preprocess_train(example_batch):
        pixel_values = []
        labels = []
        for image, target in zip(example_batch["image"], example_batch["label"]):
            image, target = train_transforms(image.convert("RGB"), target)
            pixel_values.append(image)
            labels.append(target)

        encoding = {}
        encoding["pixel_values"] = torch.stack(pixel_values)
        encoding["labels"] = torch.stack(labels)

        return encoding

    def preprocess_val(example_batch):
        pixel_values = []
        labels = []
        for image, target in zip(example_batch["image"], example_batch["label"]):
            image, target = val_transforms(image.convert("RGB"), target)
            pixel_values.append(image)
            labels.append(target)

        encoding = {}
        encoding["pixel_values"] = torch.stack(pixel_values)
        encoding["labels"] = torch.stack(labels)

        return encoding

    with accelerator.main_process_first():
        train_dataset = dataset["train"].with_transform(preprocess_train)
        eval_dataset = dataset["validation"].with_transform(preprocess_val)

    train_dataloader = DataLoader(
        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
    )
    eval_dataloader = DataLoader(
        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
    )

    
    optimizer = torch.optim.AdamW(
        list(model.parameters()),
        lr=args.learning_rate,
        betas=[args.adam_beta1, args.adam_beta2],
        eps=args.adam_epsilon,
    )

    
    checkpointing_steps = args.checkpointing_steps
    if checkpointing_steps is not None and checkpointing_steps.isdigit():
        checkpointing_steps = int(checkpointing_steps)

    
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
    )

    
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    
    metric = evaluate.load("mean_iou")

    
    
    if args.with_tracking:
        experiment_config = vars(args)
        
        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
        accelerator.init_trackers("semantic_segmentation_no_trainer", experiment_config)

    
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    
    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    starting_epoch = 0

    
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
            checkpoint_path = args.resume_from_checkpoint
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            
            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
            dirs.sort(key=os.path.getctime)
            path = dirs[-1]  
            checkpoint_path = path
            path = os.path.basename(checkpoint_path)

        accelerator.print(f"Resumed from checkpoint: {checkpoint_path}")
        accelerator.load_state(checkpoint_path)
        
        training_difference = os.path.splitext(path)[0]

        if "epoch" in training_difference:
            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
            resume_step = None
            completed_steps = starting_epoch * num_update_steps_per_epoch
        else:
            
            resume_step = int(training_difference.replace("step_", "")) * args.gradient_accumulation_steps
            starting_epoch = resume_step // len(train_dataloader)
            completed_steps = resume_step // args.gradient_accumulation_steps
            resume_step -= starting_epoch * len(train_dataloader)

    
    progress_bar.update(completed_steps)

    for epoch in range(starting_epoch, args.num_train_epochs):
        model.train()
        if args.with_tracking:
            total_loss = 0
        if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
            
            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
        else:
            active_dataloader = train_dataloader
        for step, batch in enumerate(active_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                
                if args.with_tracking:
                    total_loss += loss.detach().float()
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1

            if isinstance(checkpointing_steps, int):
                if completed_steps % checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps}"
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)

                    if args.push_to_hub and epoch < args.num_train_epochs - 1:
                        accelerator.wait_for_everyone()
                        unwrapped_model = accelerator.unwrap_model(model)
                        unwrapped_model.save_pretrained(
                            args.output_dir,
                            is_main_process=accelerator.is_main_process,
                            save_function=accelerator.save,
                        )
                        if accelerator.is_main_process:
                            image_processor.save_pretrained(args.output_dir)
                            repo.push_to_hub(
                                commit_message=f"Training in progress {completed_steps} steps",
                                blocking=False,
                                auto_lfs_prune=True,
                            )

            if completed_steps >= args.max_train_steps:
                break

        logger.info("***** Running evaluation *****")
        model.eval()
        for step, batch in enumerate(tqdm(eval_dataloader, disable=not accelerator.is_local_main_process)):
            with torch.no_grad():
                outputs = model(**batch)

            upsampled_logits = torch.nn.functional.interpolate(
                outputs.logits, size=batch["labels"].shape[-2:], mode="bilinear", align_corners=False
            )
            predictions = upsampled_logits.argmax(dim=1)

            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))

            metric.add_batch(
                predictions=predictions,
                references=references,
            )

        eval_metrics = metric.compute(
            num_labels=len(id2label),
            ignore_index=255,
            reduce_labels=False,  
        )
        logger.info(f"epoch {epoch}: {eval_metrics}")

        if args.with_tracking:
            accelerator.log(
                {
                    "mean_iou": eval_metrics["mean_iou"],
                    "mean_accuracy": eval_metrics["mean_accuracy"],
                    "overall_accuracy": eval_metrics["overall_accuracy"],
                    "train_loss": total_loss.item() / len(train_dataloader),
                    "epoch": epoch,
                    "step": completed_steps,
                },
                step=completed_steps,
            )

        if args.push_to_hub and epoch < args.num_train_epochs - 1:
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(
                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
            )
            if accelerator.is_main_process:
                image_processor.save_pretrained(args.output_dir)
                repo.push_to_hub(
                    commit_message=f"Training in progress epoch {epoch}", blocking=False, auto_lfs_prune=True
                )

        if args.checkpointing_steps == "epoch":
            output_dir = f"epoch_{epoch}"
            if args.output_dir is not None:
                output_dir = os.path.join(args.output_dir, output_dir)
            accelerator.save_state(output_dir)

    if args.with_tracking:
        accelerator.end_training()

    if args.output_dir is not None:
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        unwrapped_model.save_pretrained(
            args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
        )
        if accelerator.is_main_process:
            image_processor.save_pretrained(args.output_dir)
            if args.push_to_hub:
                repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)

            all_results = {
                f"eval_{k}": v.tolist() if isinstance(v, np.ndarray) else v for k, v in eval_metrics.items()
            }
            with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
                json.dump(all_results, f)


if __name__ == "__main__":
    main()