Code

`app` ¶

SegFormer evaluation pipeline for semantic segmentation tasks.

This module orchestrates the evaluation of SegFormer models with various quantization methods on the Scene Parse 150 dataset. It handles model loading, quantization, dataset preparation, evaluation, and result logging using Weights & Biases.

The pipeline supports multiple quantization levels and efficient processing through dataset sharding. Results are tracked and visualized for performance analysis across different model configurations.

Usage

python app.py

Environment variables

WANDB_API_KEY: API key for Weights & Biases logging WANDB_PROJECT: Name of the W&B project WANDB_ENTITY: Name of the W&B entity (team or user)

Functions¶

`main()` ¶

Main execution function for the SegFormer evaluation pipeline.

This function orchestrates the entire evaluation process, including model loading, dataset preparation, evaluation, and logging results.

Source code in src/app.py

def main():
    """
    Main execution function for the SegFormer evaluation pipeline.

    This function orchestrates the entire evaluation process, including
    model loading, dataset preparation, evaluation, and logging results.
    """	

    base_model = load_base_model(model_name, model_save_path, compute_dtype, device)
    image_processor = load_image_processor(model_name, tokenizer_save_path)
    dataset = load_dataset_custom(dataset_save_path, dataset_name)
    wandb.login(relogin=True, force=True, key=environ['WANDB_API_KEY'])

    # Quantize models
    models = quantize_models(base_model, model_name, model_save_path, device)
    print_model_sizes(models)

    # Evaluation loop
    ds_subset = 'validation'
    for model_name_i, model_i in models.items():
        model_i.eval()
        for k in range(0, ds_num_shards):
            if k == 0 or k % ds_shards_mod == 0:
                print(f"Evaluating {model_name_i}, shard {k}/{ds_num_shards}")
            wandb_run = create_wandb_run(
                environ['WANDB_PROJECT'], environ['WANDB_ENTITY'],
                model_name_i, dataset_name
            )
            create_wandb_run_meta(
                wandb_run, model_name_i, dataset_name, device, wandb_tag_mode,
                model_name_i, model_i, ds_num_shards, ds_shards_mod
            )
            dataset_shard = dataset[ds_subset].shard(
                num_shards=ds_num_shards, index=k
            )
            results = evaluate_model(
                model_i, dataset_shard, image_processor, device,
                metric_name, model_i.config.id2label
            )
            log_wandb_results(results, model_i)
            wandb.finish()

`config` ¶

Configuration settings for the SegFormer evaluation pipeline.

`utils.data_processing` ¶

Data processing module for SegFormer evaluation pipeline.

This module contains functions for loading, preprocessing, and handling dataset operations for semantic segmentation tasks using SegFormer models.

Functions:

Name	Description
`load_dataset`	Load or download and save the dataset.
`get_processed_inputs`	Process and prepare inputs for model inference.
`convert_to_RGB`	Convert dataset images and annotations to RGB and grayscale respectively.

The module uses the Hugging Face datasets library and image processing tools to prepare data for SegFormer model evaluation.

Functions¶

`convert_to_RGB(dataset)` ¶

Convert dataset images and annotations to RGB and grayscale respectively.

Parameters:

Name	Type	Description	Default
`dataset`	`Dict`	Dataset containing ‘image’ and ‘annotation’ keys.	required

Returns:

Type	Description
`Dict[List, List]`	Dict[List, List]: Processed dataset with converted images and annotations.

Source code in src/utils/data_processing.py

def convert_to_RGB(
    dataset: Dict[List, List]
) -> Dict[List, List]:
    """
    Convert dataset images and annotations to RGB and grayscale respectively.

    Args:
        dataset (Dict): Dataset containing 'image' and 'annotation' keys.

    Returns:
        Dict[List, List]: Processed dataset with converted images and annotations.
    """

    images = [img.convert("RGB") for img in dataset['image']]
    annotations = [img.convert("L") for img in dataset['annotation']]
    return {'image': images, 'annotation': annotations}

`get_processed_inputs(dataset, image_processor, device, bias_dtype=None)` ¶

Process and prepare inputs for model inference.

Parameters:

Name	Type	Description	Default
`dataset`	`Dataset`	The dataset to process.	required
`image_processor`	`SegformerImageProcessor`	The image processor to use.	required
`device`	`device`	The device to load tensors to.	required
`bias_dtype`	`dtype`	Dtype for bias, if any.	`None`

Returns:

Type	Description
`Tuple[Tensor, Tensor]`	Tuple[torch.Tensor, torch.Tensor]: Processed pixel values and labels.

Source code in src/utils/data_processing.py

def get_processed_inputs(
    dataset: Dataset,
    image_processor: SegformerImageProcessor,
    device: device,
    bias_dtype: dtype = None
) -> Tuple[Tensor, Tensor]:
    """
    Process and prepare inputs for model inference.

    Args:
        dataset (Dataset): The dataset to process.
        image_processor (SegformerImageProcessor): The image processor to use.
        device (torch.device): The device to load tensors to.
        bias_dtype (torch.dtype, optional): Dtype for bias, if any.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Processed pixel values and labels.
    """

    set_verbosity_error()
    dataset = dataset.map(convert_to_RGB, batched=True)
    dataset = image_processor.preprocess(
        images=dataset['image'],
        segmentation_maps=dataset['annotation'],
        return_tensors="pt",
    )
    pixel_values = dataset['pixel_values'].to(device)
    labels = dataset['labels'].to(device)
    if bias_dtype in [half, float16]:
        pixel_values = pixel_values.half()
    return pixel_values, labels

`load_dataset_custom(dataset_save_path, dataset_name)` ¶

Load or download and save the dataset.

Parameters:

Name	Type	Description	Default
`dataset_save_path`	`str`	Path to save/load the dataset.	required
`dataset_name`	`str`	Name of the dataset to load.	required

Returns:

Name	Type	Description
`Dataset`		The loaded dataset.

Source code in src/utils/data_processing.py

def load_dataset_custom(
    dataset_save_path: str,
    dataset_name: str
):
    """
    Load or download and save the dataset.

    Args:
        dataset_save_path (str): Path to save/load the dataset.
        dataset_name (str): Name of the dataset to load.

    Returns:
        Dataset: The loaded dataset.
    """

    try:
        return load_from_disk(dataset_save_path)
    except Exception:
        dataset = load_dataset(dataset_name, trust_remote_code=True)
        dataset.save_to_disk(dataset_save_path)
        return dataset

`utils.evaluator` ¶

Evaluator module for SegFormer model inference and metric computation.

This module provides functions for model inference and evaluation on semantic segmentation tasks using SegFormer models.

Functions:

Name	Description
`infer_model`	Perform model inference and return loss and logits.
`evaluate_model`	Evaluate the model on a dataset shard and compute metrics.

The module uses PyTorch for model inference and the ‘evaluate’ library for computing semantic segmentation metrics.

Functions¶

`evaluate_model(model, dataset_shard, image_processor, device, metric_name, id2label)` ¶

Evaluate the model on a dataset shard and compute metrics.

Parameters:

Name	Type	Description	Default
`model`	`SegformerForSemanticSegmentation`	The model to evaluate.	required
`dataset_shard`	`Dataset`	A shard of the dataset to evaluate on.	required
`image_processor`	`SegformerImageProcessor`	The image processor to use.	required
`device`	`device`	The device to run evaluation on.	required
`metric_name`	`str`	Name of the metric to use.	required
`id2label`	`dict`	Mapping of label IDs to label names.	required

Returns:

Type	Description
`Dict[str, float]`	Dict[str, float]: Computed evaluation metrics where keys are metric names and values are the corresponding scores.

Source code in src/utils/evaluator.py

def evaluate_model(
    model, # : 'SegformerForSemanticSegmentation',
    dataset_shard, # : 'Dataset',
    image_processor, # : 'SegformerImageProcessor',
    device: torch.device,
    metric_name: str,
    id2label: Dict[int, str]
) -> Dict[str, float]:
    """
    Evaluate the model on a dataset shard and compute metrics.

    Args:
        model (SegformerForSemanticSegmentation): The model to evaluate.
        dataset_shard (Dataset): A shard of the dataset to evaluate on.
        image_processor (SegformerImageProcessor): The image processor to use.
        device (torch.device): The device to run evaluation on.
        metric_name (str): Name of the metric to use.
        id2label (dict): Mapping of label IDs to label names.

    Returns:
        Dict[str, float]: Computed evaluation metrics where keys are metric names and values are the corresponding scores.
    """

    pixel_values, labels = get_processed_inputs(dataset_shard, image_processor, device, model.config.torch_dtype)
    loss, logits = infer_model(model, pixel_values, labels)
    predictions = torch.nn.functional.interpolate(
        logits,
        size=labels.shape[-2:],
        mode="bilinear", align_corners=False
    ).argmax(dim=1)
    metric = load(metric_name)
    results = metric.compute(
        predictions=predictions.cpu().numpy(),
        references=labels.cpu().numpy(),
        num_labels=len(id2label),
        ignore_index=model.config.semantic_loss_ignore_index
    )
    return results

`infer_model(model, pixel_values, labels)` ¶

Perform model inference and return loss and logits.

Parameters:

Name	Type	Description	Default
`model`	`SegformerForSemanticSegmentation`	The model to use for inference.	required
`pixel_values`	`Tensor`	Input pixel values.	required
`labels`	`Tensor`	Ground truth labels.	required

Returns:

Type	Description
`Tuple[float, Tensor]`	Tuple[float, torch.Tensor]: A tuple containing the model loss as a float and the logits as a torch.Tensor.

Source code in src/utils/evaluator.py

def infer_model(
    model, # : 'SegformerForSemanticSegmentation',
    pixel_values: torch.Tensor,
    labels: torch.Tensor
) -> Tuple[float, torch.Tensor]:
    """
    Perform model inference and return loss and logits.

    Args:
        model (SegformerForSemanticSegmentation): The model to use for inference.
        pixel_values (torch.Tensor): Input pixel values.
        labels (torch.Tensor): Ground truth labels.

    Returns:
        Tuple[float, torch.Tensor]: A tuple containing the model loss as a float and the logits as a torch.Tensor.
    """

    with torch.no_grad():
        outputs = model(pixel_values=pixel_values, labels=labels)
    return outputs.loss, outputs.logits

`utils.general_utils` ¶

`utils.model_loader` ¶

Model loading module for SegFormer evaluation pipeline.

This module provides functions for loading and initializing SegFormer models and image processors, with support for local and remote loading.

Functions:

Name	Description
`load_base_model`	Load or download and save the base SegFormer model.
`load_image_processor`	Load or download and save the image processor.

The module uses the transformers library for model and processor management.

Functions¶

`load_base_model(model_name, model_save_path, compute_dtype, device)` ¶

Load or download and save the base SegFormer model.

Parameters:

Name	Type	Description	Default
`model_name`	`str`	Name of the model to load.	required
`model_save_path`	`str`	Path to save/load the model.	required
`compute_dtype`	`dtype`	Computation dtype for the model.	required
`device`	`device`	Device to load the model to.	required

Returns:

Name	Type	Description
`SegformerForSemanticSegmentation`	`SegformerForSemanticSegmentation`	The loaded model.

Source code in src/utils/model_loader.py

def load_base_model(
    model_name: str,
    model_save_path: str,
    compute_dtype: dtype, 
    device: device
) -> SegformerForSemanticSegmentation:
    """
    Load or download and save the base SegFormer model.

    Args:
        model_name (str): Name of the model to load.
        model_save_path (str): Path to save/load the model.
        compute_dtype (torch.dtype): Computation dtype for the model.
        device (torch.device): Device to load the model to.

    Returns:
        SegformerForSemanticSegmentation: The loaded model.
    """

    try:
        print("loading from disk")
        model = SegformerForSemanticSegmentation.from_pretrained(model_save_path)
    except Exception as e:
        print(
            f"Error loading local model: {e}."
            f"Loading from source and saving to disk."
        )
        if not exists(model_save_path):
            makedirs(model_save_path, exist_ok=True)
        model = SegformerForSemanticSegmentation.from_pretrained(
            model_name,
            torch_dtype=compute_dtype,
        )
        model.save_pretrained(model_save_path)
    return model.to(device)

`load_image_processor(model_name, tokenizer_save_path)` ¶

Load or download and save the image processor.

Parameters:

Name	Type	Description	Default
`model_name`	`str`	Name of the model to load the processor for.	required
`tokenizer_save_path`	`str`	Path to save/load the processor.	required

Returns:

Name	Type	Description
`SegformerImageProcessor`	`SegformerImageProcessor`	The loaded image processor.

Source code in src/utils/model_loader.py

def load_image_processor(
    model_name: str,
    tokenizer_save_path: str
) -> SegformerImageProcessor:
    """
    Load or download and save the image processor.

    Args:
        model_name (str): Name of the model to load the processor for.
        tokenizer_save_path (str): Path to save/load the processor.

    Returns:
        SegformerImageProcessor: The loaded image processor.
    """

    try:
        return SegformerImageProcessor.from_pretrained(tokenizer_save_path)
    except Exception as e:
        print(
            f"Error loading local ImageProcessor: {e}."
            f"Loading from source and saving to disk."
        )
        image_processor = SegformerImageProcessor.from_pretrained(model_name)
        image_processor.save_pretrained(tokenizer_save_path)
        return image_processor

`utils.quantization` ¶

Quantization module for SegFormer models.

This module provides functions for quantizing SegFormer models using various quantization methods supported by the Quanto library.

Functions:

Name	Description
`quantize_models`	Quantize a base SegFormer model using multiple quantization levels.

The module uses Quanto for quantization and supports float8, int8, int4, and int2 quantization.

Functions¶

`quantize_models(base_model, model_name, model_save_path, torch_device)` ¶

Quantize the base model using various quantization methods.

Parameters:

Name	Type	Description	Default
`base_model`	`SegformerForSemanticSegmentation`	The base model to quantize.	required
`model_name`	`str`	Name of the model.	required
`model_save_path`	`str`	Path to save quantized models.	required
`torch_device`	`device`	Device to load models to.	required

Returns:

Type	Description
`Dict[str, SegformerForSemanticSegmentation]`	Dict[str, SegformerForSemanticSegmentation]: Dictionary of quantized models.

Source code in src/utils/quantization.py

def quantize_models(
    base_model: SegformerForSemanticSegmentation,
    model_name: str,
    model_save_path: str,
    torch_device: device
) -> Dict[str, SegformerForSemanticSegmentation]:
    """
    Quantize the base model using various quantization methods.

    Args:
        base_model (SegformerForSemanticSegmentation): The base model to quantize.
        model_name (str): Name of the model.
        model_save_path (str): Path to save quantized models.
        torch_device (torch.device): Device to load models to.

    Returns:
        Dict[str, SegformerForSemanticSegmentation]: Dictionary of quantized models.
    """

    models = {}
    config_quanto = {}
    bits_quanto_w = ['float8', 'int8', 'int4', 'int2']
    for nbits in bits_quanto_w:
        config_quanto[nbits] = QuantoConfig(weights=nbits)

    for nbits in bits_quanto_w:
        model_htype = f"{model_name}-quanto-{nbits}"
        model_save_path_quanto = f"{model_save_path}/{model_name}-{model_htype}"
        try:
            print(f"loading local {model_htype}")
            models[model_htype] = SegformerForSemanticSegmentation.from_pretrained(model_save_path_quanto)
        except Exception:
            try:
                print(f"loading local {model_name}")
                models[model_htype] = SegformerForSemanticSegmentation.from_pretrained(
                    model_save_path,
                    local_files_only=True,
                    torch_dtype=base_model.config.torch_dtype,
                    quantization_config=config_quanto[nbits],
                )
            except Exception:
                print(f"loading online {model_name}")
                models[model_htype] = SegformerForSemanticSegmentation.from_pretrained(
                    model_name,
                    torch_dtype=base_model.config.torch_dtype,
                    quantization_config=config_quanto[nbits],
                )
        quanto.freeze(models[model_htype])
        models[model_htype] = models[model_htype].to(torch_device)
    return models

`utils.wandb_utils` ¶

Weights & Biases utility module for SegFormer evaluation pipeline.

Provides functions for managing W&B runs, including initialization, metadata setting, and result logging for the SegFormer evaluation pipeline.

Functions:

Name	Description
`create_wandb_run`	Initialize a new W&B run.
`create_wandb_run_meta`	Set metadata for a W&B run.
`log_wandb_results`	Log evaluation results to W&B.

Note: Requires WANDB_API_KEY, WANDB_PROJECT, and WANDB_ENTITY environment variables.

Functions¶

`create_wandb_run(project, entity, name, group)` ¶

Initialize and create a new Weights & Biases run.

Parameters:

Name	Type	Description	Default
`project`	`str`	Name of the W&B project.	required
`entity`	`str`	Name of the W&B entity (team or user).	required
`name`	`str`	Name of the run.	required
`group`	`str`	Group name for the run.	required

Returns:

Type	Description
`Run`	wandb.Run: The created W&B run object.

Source code in src/utils/wandb_utils.py

def create_wandb_run(
    project: str,
    entity: str,
    name: str,
    group: str
) -> Run:
    """
    Initialize and create a new Weights & Biases run.

    Args:
        project (str): Name of the W&B project.
        entity (str): Name of the W&B entity (team or user).
        name (str): Name of the run.
        group (str): Group name for the run.

    Returns:
        wandb.Run: The created W&B run object.
    """

    wandb_run = init(
        project=project,
        entity=entity,
        name=name,
        group=group
    )
    assert wandb_run is run
    return wandb_run

`create_wandb_run_meta(wandb_run, model_name, dataset_name, torch_device, wandb_tag_mode, quant_used, model_used, ds_num_shards, ds_shards_mod)` ¶

Set metadata for the Weights & Biases run.

Parameters:

Name	Type	Description	Default
`wandb_run`	`Run`	The W&B run object.	required
`model_name`	`str`	Name of the model.	required
`dataset_name`	`str`	Name of the dataset.	required
`torch_device`	`device`	Device used for computation.	required
`wandb_tag_mode`	`str`	Tag for the run mode.	required
`quant_used`	`str`	Quantization method used, if any.	required
`model_used`	`SegformerForSemanticSegmentation`	The model being evaluated.	required
`ds_num_shards`	`int`	Number of dataset shards.	required
`ds_shards_mod`	`float`	Modulo for dataset shard logging.	required

Returns:

Type	Description
`Run`	wandb.Run: The updated W&B run object.

Source code in src/utils/wandb_utils.py

def create_wandb_run_meta(
    wandb_run: Run,
    model_name: str,
    dataset_name: str,
    torch_device: device,
    wandb_tag_mode: str,
    quant_used: str,
    model_used: SegformerForSemanticSegmentation,
    ds_num_shards: int,
    ds_shards_mod: float
) -> Run:
    """
    Set metadata for the Weights & Biases run.

    Args:
        wandb_run (wandb.Run): The W&B run object.
        model_name (str): Name of the model.
        dataset_name (str): Name of the dataset.
        torch_device (torch.device): Device used for computation.
        wandb_tag_mode (str): Tag for the run mode.
        quant_used (str): Quantization method used, if any.
        model_used (SegformerForSemanticSegmentation): The model being evaluated.
        ds_num_shards (int): Number of dataset shards.
        ds_shards_mod (float): Modulo for dataset shard logging.

    Returns:
        wandb.Run: The updated W&B run object.
    """

    wandb_tags = [model_name, dataset_name, torch_device.type, wandb_tag_mode]
    if quant_used:
        wandb_tags += [quant_used]
    else:
        wandb_tags += [str(model_used.config.torch_dtype)]
    wandb_run.tags = wandb_tags
    wandb_run.notes = f"{datetime.now().isoformat()}, " \
        f"model size {model_used.get_memory_footprint()*1.0e-6:.2f} MB, " \
        f"{ds_num_shards=}, {ds_shards_mod=}"
    return wandb_run

`log_wandb_results(results, model)` ¶

Log evaluation results to Weights & Biases.

Parameters:

Name	Type	Description	Default
`results`	`dict`	Evaluation results to log.	required
`model`	`SegformerForSemanticSegmentation`	The model being evaluated.	required

Returns:

Type	Description
`None`	None

Source code in src/utils/wandb_utils.py

def log_wandb_results(
    results: Dict,
    model: SegformerForSemanticSegmentation
) -> None:
    """
    Log evaluation results to Weights & Biases.

    Args:
        results (dict): Evaluation results to log.
        model (SegformerForSemanticSegmentation): The model being evaluated.

    Returns:
        None
    """

    log({
        'mean_iou': results['mean_iou'],
        'mean_accuracy': results['mean_accuracy'],
        'overall_accuracy': results['overall_accuracy'],
        'memory_footprint_MB': float(f"{model.get_memory_footprint()*1.0e-6:.2f}")
    })

Code

app ¶

Functions¶

main() ¶

config ¶

utils.data_processing ¶

Functions¶

convert_to_RGB(dataset) ¶

get_processed_inputs(dataset, image_processor, device, bias_dtype=None) ¶

load_dataset_custom(dataset_save_path, dataset_name) ¶

utils.evaluator ¶

Functions¶

evaluate_model(model, dataset_shard, image_processor, device, metric_name, id2label) ¶

infer_model(model, pixel_values, labels) ¶

utils.general_utils ¶

utils.model_loader ¶

Functions¶

load_base_model(model_name, model_save_path, compute_dtype, device) ¶

load_image_processor(model_name, tokenizer_save_path) ¶

utils.quantization ¶

Functions¶

quantize_models(base_model, model_name, model_save_path, torch_device) ¶

utils.wandb_utils ¶

Functions¶

create_wandb_run(project, entity, name, group) ¶

create_wandb_run_meta(wandb_run, model_name, dataset_name, torch_device, wandb_tag_mode, quant_used, model_used, ds_num_shards, ds_shards_mod) ¶

log_wandb_results(results, model) ¶

`app` ¶

`main()` ¶

`config` ¶

`utils.data_processing` ¶

`convert_to_RGB(dataset)` ¶

`get_processed_inputs(dataset, image_processor, device, bias_dtype=None)` ¶

`load_dataset_custom(dataset_save_path, dataset_name)` ¶

`utils.evaluator` ¶

`evaluate_model(model, dataset_shard, image_processor, device, metric_name, id2label)` ¶

`infer_model(model, pixel_values, labels)` ¶

`utils.general_utils` ¶

`utils.model_loader` ¶

`load_base_model(model_name, model_save_path, compute_dtype, device)` ¶

`load_image_processor(model_name, tokenizer_save_path)` ¶

`utils.quantization` ¶

`quantize_models(base_model, model_name, model_save_path, torch_device)` ¶

`utils.wandb_utils` ¶

`create_wandb_run(project, entity, name, group)` ¶

`create_wandb_run_meta(wandb_run, model_name, dataset_name, torch_device, wandb_tag_mode, quant_used, model_used, ds_num_shards, ds_shards_mod)` ¶

`log_wandb_results(results, model)` ¶