given the code below, consider how the code could be improved scale to thousands or even millions of images or models. Would it make sense to qualntize and dequantize the model? Use distributed training?

Code:
import pandas as pd
from PIL import Image
from pathlib import Path
import torch
import torch.nn as nn
import os
import json
from torch.utils.data import Dataset, DataLoader
import statistics


class ImageCaptionDataset(Dataset):
    """
     A PyTorch Dataset for image-caption pairs, which applies specified transformations to the images, and tokenizes the
     captions.

     Args:
         df (pandas.DataFrame):
             The DataFrame containing image-caption data.
         transform (callable or None):
             The transformation function applied to images.
         tokenizer (callable or None):
             The tokenizer function applied to captions.

     """
    def __init__(self, df, transform=None, tokenizer=None):
        self.transform = transform
        self.tokenizer = tokenizer
        self.df = df

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        image_path = self.df.iloc[idx]['image_path']
        caption = self.df.iloc[idx]['caption']
        image = Image.open(image_path)
        if self.transform is not None:
            image = self.transform(image)
        if self.tokenizer is not None:
            caption = self.tokenizer(caption, truncate=True).squeeze()
        return image, caption


def load_data(csv_path, images_dir):
    """
        Load image-caption data from a CSV file and verify the existence of image files.

    Args:
        csv_path (str):
            The file path to the CSV containing image URLs and captions.

        images_dir (str):
            The directory path where image files are stored.

    Returns:
        pandas.DataFrame:
            A DataFrame containing the validated image-caption pairs.
    """
    df = pd.read_csv(csv_path)
    df['image_path'] = df['url'].apply(lambda x: os.path.join(images_dir, Path(x).parts[3] + '.png'))
    missing_images = df[~df['image_path'].apply(os.path.exists)]
    if not missing_images.empty:
        print("Warning: The following images are missing and will be skipped:")
        print(missing_images['url'].tolist())
        df = df[df['image_path'].apply(os.path.exists)]
    return df


def compute_similarity_statistics(similarities):
    """
    Compute statistics for similarity scores.

    Args:
        similarities (list): List of similarity scores.

    Returns:
        dict: Statistics including average, min, and max similarity.
    """
    if not similarities:
        return {
            'average_similarity': 0,
            'min_similarity': 0,
            'max_similarity': 0
        }

    avg_sim = statistics.mean(similarities)
    min_sim = min(similarities)
    max_sim = max(similarities)

    return {
        'average_similarity': avg_sim,
        'min_similarity': min_sim,
        'max_similarity': max_sim
    }


def save_embeddings(image_features, text_features, start_idx, save_dir='embeddings'):
    """
    Save image and text feature embeddings to disk.
    Args:
        image_features (torch.Tensor):
            A tensor containing the image embeddings. Typically produced by `model.encode_image()`.

        text_features (torch.Tensor):
            A tensor containing the text embeddings. Typically produced by `model.encode_text()`.

        start_idx (int):
            The starting index of the current batch.
        save_dir (str, optional):
            The directory where the embedding files will be saved.
    """
    os.makedirs(save_dir, exist_ok=True)
    torch.save(image_features.cpu(), os.path.join(save_dir, f'image_features_{start_idx}.pt'))
    torch.save(text_features.cpu(), os.path.join(save_dir, f'text_features_{start_idx}.pt'))


def save_checkpoint(state, filename='checkpoint.json'):
    """
    Save the current state to a checkpoint file.

    Args:
        state (dict): The state to save (e.g., similarities, last_idx).
        filename (str): The filename for the checkpoint.
    """
    try:
        with open(filename, 'w') as f:
            json.dump(state, f)
        print(f"Checkpoint saved at batch {state['last_idx']}.")
    except Exception as e:
        print(f"Error saving checkpoint: {e}")


def load_checkpoint(filename='checkpoint.json'):
    """
    Load the state from a checkpoint file if it exists.

    Args:
        filename (str): The filename of the checkpoint.

    Returns:
        dict or None: The loaded state dictionary or None if loading fails.
    """
    if os.path.exists(filename):
        try:
            with open(filename, 'r') as f:
                state = json.load(f)
            print(f"Checkpoint loaded. Resuming from batch {state['last_idx']}.")
            return state
        except Exception as e:
            print(f"Error loading checkpoint: {e}")
            return None
    return None


class CLIPModelWrapper(nn.Module):
    """
    A wrapper for the CLIP (Contrastive Language-Image Pretraining) model to facilitate
    parallel processing and streamline the computation of image-text similarities.
    It is designed to work seamlessly with PyTorch's `DataParallel` for multi-GPU utilization.

    Args:
        model (nn.Module):
            The wrapped CLIP model instance.
    """
    def __init__(self, model):
        super(CLIPModelWrapper, self).__init__()
        self.model = model

    def forward(self, images, texts):
        image_features = self.model.encode_image(images)
        text_features = self.model.encode_text(texts)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        similarity = (image_features * text_features).sum(dim=-1)
        return similarity, image_features, text_features
import torch
import torch.nn as nn
import clip
import time
from memory_profiler import memory_usage
from torch.utils.data import Dataset, DataLoader
import statistics
import sys
import traceback
from utils import (ImageCaptionDataset, load_data, compute_similarity_statistics, save_embeddings, save_checkpoint,
                   load_checkpoint, CLIPModelWrapper)


def compute_similarities(dataloader, model, device, checkpoint_file=None, checkpoint_interval=100):
    """
    Compute similarity between image and caption embeddings using CLIP with optional checkpointing.

    Args:
        dataloader (DataLoader): DataLoader for the dataset.
        model (torch.nn.Module): The CLIP model.
        device (torch.device): Device to perform computations on.
        checkpoint_file (str, optional): Filename for saving/loading checkpoints. Defaults to None.
        checkpoint_interval (int, optional): Number of batches between saving checkpoints. Defaults to 100.

    Returns:
        tuple: (similarities list, processing_metrics dict)
    """
    similarities = []
    model.eval()

    # Metrics initialization
    total_batches = 0
    failed_batches = 0
    batch_times = []
    start_idx = 0

    # Load existing similarities if checkpoint_file is provided and exists
    if checkpoint_file is not None:
        state = load_checkpoint(checkpoint_file)
        if state and 'similarities' in state and 'last_idx' in state:
            similarities = state['similarities']
            start_idx = state['last_idx']
        elif state is not None:
            print("Checkpoint file is corrupted or incomplete. Starting from scratch.")

    with torch.no_grad():
        for batch_idx, (images, texts) in enumerate(dataloader):
            # Skip batches if resuming from a checkpoint
            if batch_idx < start_idx:
                continue

            batch_start_time = time.time()
            try:
                images = images.to(device, non_blocking=True)
                texts = texts.to(device, non_blocking=True)

                similarity, image_features, text_features = model(images, texts)
                similarity = similarity.cpu().numpy()
                similarities.extend(similarity.tolist())

                batch_end_time = time.time()
                batch_processing_time = batch_end_time - batch_start_time
                batch_times.append(batch_processing_time)
                total_batches += 1
            except Exception as e:
                failed_batches += 1
                print(f"Error processing batch {batch_idx}: {e}")
                traceback.print_exc(file=sys.stdout)
                continue

            # Checkpointing at specified intervals
            if checkpoint_file is not None and (batch_idx + 1) % checkpoint_interval == 0:
                checkpoint_state = {
                    'last_idx': batch_idx + 1,
                    'similarities': similarities
                }
                save_checkpoint(checkpoint_state, filename=checkpoint_file)

            # Save embeddings
            save_embeddings(image_features, text_features, batch_idx * dataloader.batch_size)

    # Compile metrics
    processing_metrics = {
        'total_images_processed': len(similarities),
        'total_batches_processed': total_batches,
        'failed_batches': failed_batches,
        'average_batch_time_seconds': statistics.mean(batch_times) if batch_times else 0,
        'min_batch_time_seconds': min(batch_times) if batch_times else 0,
        'max_batch_time_seconds': max(batch_times) if batch_times else 0,
    }

    return similarities, processing_metrics


def filter_high_quality_pairs(df, threshold=0.3, output_file='filtered_results.csv'):
    """
    Filter out low-quality image-caption pairs and save results to CSV.

    Args:
        df (pd.DataFrame): DataFrame containing image-caption pairs with similarity scores.
        threshold (float, optional): Similarity threshold. Pairs with similarity >= threshold are kept. Defaults to 0.3.
        output_file (str, optional): Filename for the filtered CSV. Defaults to 'filtered_results.csv'.
    """
    # Filter out pairs with similarity below the threshold
    high_quality_df = df[df['similarity'] >= threshold]

    # Save the filtered DataFrame back to CSV
    high_quality_df.to_csv(output_file, index=False)
    print(f"Filtered results saved to '{output_file}' with {len(high_quality_df)} high-quality pairs.")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Compute image-text similarity using CLIP.")
    parser.add_argument('--csv', type=str, required=True, help="Path to the CSV file.")
    parser.add_argument('--images_dir', type=str, required=True, help="Directory containing images.")
    parser.add_argument('--batch-size', type=int, default=8,
                        help='Batch size to use')
    parser.add_argument('--output_csv', type=str, default='output.csv', help="Output CSV file with similarity scores.")
    parser.add_argument('--metrics_file', type=str, default='metrics.txt', help="File to store performance metrics.")
    parser.add_argument('--device', type=str, default=None,
                        help='Device to use')
    parser.add_argument('--clip-model', type=str, default='ViT-B/32',
                        help='CLIP model to use')
    parser.add_argument('--threshold', type=float, default=0.3,
                        help='Similarity threshold for filtering high-quality pairs')
    parser.add_argument('--num_workers', type=int, default=4,
                        help='Number of worker threads for data loading')
    parser.add_argument('--checkpoint_interval', type=int, default=100,
                        help='Number of batches between checkpoints')
    parser.add_argument('--checkpoint_file', type=str, default='checkpoint.json',
                        help='Filename for saving checkpoints')

    args = parser.parse_args()

    start_time = time.time()

    # Determine the device
    if args.device is None:
        if torch.cuda.is_available():
            device = torch.device('cuda')
        else:
            device = torch.device('cpu')
    else:
        device = torch.device(args.device)
    print(f"Using device: {device}")

    # Load CLIP model
    model, preprocess = clip.load(args.clip_model, device=device)
    # Wrap the model
    model = CLIPModelWrapper(model)

    # Decide whether to use DataParallel
    use_dataparallel = False
    num_available_gpus = torch.cuda.device_count()

    if device.type == 'cuda':
        # If device is 'cuda' without specific index, consider using DataParallel
        if args.device.lower() == 'cuda' and num_available_gpus > 1:
            use_dataparallel = True
        # If device is 'cuda:X' where X is a specific GPU index, do not use DataParallel
        elif args.device.startswith('cuda:') and num_available_gpus > 1:
            # Extract the GPU index
            try:
                specified_gpu = int(args.device.split(':')[1])
                if specified_gpu < num_available_gpus:
                    use_dataparallel = False
                else:
                    print(f"Specified GPU index {specified_gpu} is out of range. Using CPU instead.")
                    device = torch.device('cpu')
                    model = CLIPModelWrapper(model)  # Re-wrap if device changes
            except ValueError:
                print(f"Invalid device format: {args.device}. Using single GPU or CPU.")
                use_dataparallel = False
    else:
        # If not using CUDA, do not use DataParallel
        use_dataparallel = False

    if use_dataparallel:
        print(f"Using DataParallel with {num_available_gpus} GPUs.")
        model = nn.DataParallel(model)
    else:
        if device.type == 'cuda':
            print(f"Using single GPU: {device}.")
        else:
            print("Using CPU.")
    model.to(device)

    # Load data
    df = load_data(args.csv, args.images_dir)
    dataset = ImageCaptionDataset(df, transform=preprocess, tokenizer=clip.tokenize)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers,
                            pin_memory=True)

    # Compute similarity
    similarities, processing_metrics = compute_similarities(
        dataloader,
        model,
        device,
        checkpoint_interval=args.checkpoint_interval,
        checkpoint_file=args.checkpoint_file
    )

    df['similarity'] = similarities

    # Save updated CSV
    df.to_csv(args.output_csv, index=False)
    print(f"Similarity scores saved to {args.output_csv}")

    filter_high_quality_pairs(df, threshold=args.threshold)

    # Compute similarity statistics
    similarity_stats = compute_similarity_statistics(similarities)

    end_time = time.time()
    total_time = end_time - start_time

    # Compile all metrics
    all_metrics = {
        'Total Execution Time (seconds)': round(total_time, 2),
        'Total Images Processed': len(similarities),
        'Total Captions Processed': len(similarities),
        'Total Batches Processed': processing_metrics['total_batches_processed'],
        'Failed Batches': processing_metrics['failed_batches'],
        'Average Batch Processing Time (seconds)': round(processing_metrics['average_batch_time_seconds'], 4),
        'Min Batch Processing Time (seconds)': round(processing_metrics['min_batch_time_seconds'], 4),
        'Max Batch Processing Time (seconds)': round(processing_metrics['max_batch_time_seconds'], 4),
        'Average Similarity': round(similarity_stats['average_similarity'], 4),
        'Min Similarity': round(similarity_stats['min_similarity'], 4),
        'Max Similarity': round(similarity_stats['max_similarity'], 4)
    }
    # Add memory_profiler's peak memory
    try:
        peak_memory = max(memory_usage((lambda: None)))  # Dummy call to get current memory usage
        all_metrics['Peak Memory Usage During Execution (MiB)'] = round(peak_memory, 2)
    except Exception as e:
        print(f"Error retrieving peak memory usage: {e}")

    # Save metrics to the metrics file
    print(f"Saving performance metrics to '{args.metrics_file}'...")
    with open(args.metrics_file, 'w') as f:
        for key, value in all_metrics.items():
            f.write(f"{key}: {value}\n")
    print(f"Performance metrics saved to '{args.metrics_file}'.")


if __name__ == "__main__":
    main()