updated vision_transformer.py

vision_transformer.py

@@ -1,362 +1,8 @@
-import os
-import random
 import torch
 import torchvision
-import torch._dynamo
-import matplotlib.pyplot as plt
-from typing import List
 from torch import nn
-from torch.utils.data import DataLoader
-from torch.nn.init import trunc_normal_, xavier_normal_, zeros_, orthogonal_, kaiming_normal_
-from torchvision import datasets
-from torchvision.transforms import v2
-
-def display_random_images(dataset: torch.utils.data.dataset.Dataset, # or torchvision.datasets.ImageFolder?
-                          classes: List[str] = None,
-                          n: int = 10,
-                          display_shape: bool = True,
-                          rows: int = 5,
-                          cols: int = 5,
-                          seed: int = None):
-    
-   
-    """Displays a number of random images from a given dataset.
-
-    Args:
-        dataset (torch.utils.data.dataset.Dataset): Dataset to select random images from.
-        classes (List[str], optional): Names of the classes. Defaults to None.
-        n (int, optional): Number of images to display. Defaults to 10.
-        display_shape (bool, optional): Whether to display the shape of the image tensors. Defaults to True.
-        rows: number of rows of the subplot
-        cols: number of columns of the subplot
-        seed (int, optional): The seed to set before drawing random images. Defaults to None.
-    
-    Usage:
-    display_random_images(train_data, 
-                      n=16, 
-                      classes=class_names,
-                      rows=4,
-                      cols=4,
-                      display_shape=False,
-                      seed=None)
-    """
-
-    # Setup the range to select images
-    n = min(n, len(dataset))
-    # Adjust display if n too high
-    if n > rows*cols:
-        n = rows*cols
-        #display_shape = False
-        print(f"For display purposes, n shouldn't be larger than {rows*cols}, setting to {n} and removing shape display.")
-    
-    # Set random seed
-    if seed:
-        random.seed(seed)
-
-    # Get random sample indexes
-    random_samples_idx = random.sample(range(len(dataset)), k=n)
-
-    # Setup plot
-    plt.figure(figsize=(cols*4, rows*4))
-
-    #Loop through samples and display random samples 
-    for i, targ_sample in enumerate(random_samples_idx):
-        targ_image, targ_label = dataset[targ_sample][0], dataset[targ_sample][1]
-
-        # 7. Adjust image tensor shape for plotting: [color_channels, height, width] -> [color_channels, height, width]
-        targ_image_adjust = targ_image.permute(1, 2, 0)
-
-        # Plot adjusted samples
-        plt.subplot(rows, cols, i+1)        
-        plt.imshow(targ_image_adjust)
-        plt.axis("off")
-        if classes:
-            title = f"class: {classes[targ_label]}"
-            if display_shape:
-                title = title + f"\nshape: {targ_image_adjust.shape}"
-        plt.title(title)
-
-def create_dataloaders(
-    train_dir: str, 
-    test_dir: str, 
-    train_transform: v2.Compose, 
-    test_transform: v2.Compose,
-    batch_size: int, 
-    num_workers: int=os.cpu_count()
-):
-  """Creates training and testing DataLoaders.
-
-  Takes in a training directory and testing directory path and turns
-  them into PyTorch Datasets and then into PyTorch DataLoaders.
-
-  Args:
-    train_dir: Path to training directory.
-    test_dir: Path to testing directory.
-    train_transform: torchvision transforms to perform on training data.
-    test_transform: torchvision transforms to perform on test data.
-    batch_size: Number of samples per batch in each of the DataLoaders.
-    num_workers: An integer for number of workers per DataLoader.
-
-  Returns:
-    A tuple of (train_dataloader, test_dataloader, class_names).
-    Where class_names is a list of the target classes.
-    Example usage:
-      train_dataloader, test_dataloader, class_names = \
-        = create_dataloaders(train_dir=path/to/train_dir,
-                             test_dir=path/to/test_dir,
-                             transform=some_transform,
-                             batch_size=32,
-                             num_workers=4)
-  """
-  # Use ImageFolder to create dataset(s)
-  train_data = datasets.ImageFolder(train_dir, transform=train_transform)
-  test_data = datasets.ImageFolder(test_dir, transform=test_transform)
-
-  # Get class names
-  class_names = train_data.classes
-
-  # Turn images into data loaders
-  train_dataloader = DataLoader(
-      train_data,
-      batch_size=batch_size,
-      shuffle=True,
-      num_workers=num_workers,
-      pin_memory=True, #enables fast data transfers to CUDA-enabled GPU
-  )
-  test_dataloader = DataLoader(
-      test_data,
-      batch_size=batch_size,
-      shuffle=False,
-      num_workers=num_workers,
-      pin_memory=True, #enables fast data transfers to CUDA-enabled GPU
-  )
-
-  return train_dataloader, test_dataloader, class_names
-
-def create_dataloader_for_vit(
-        vit_model: str="bitbase16",
-        train_dir: str="./",
-        test_dir: str="./",
-        batch_size: int=64,
-        aug: bool=True,
-        display_imgs: bool=True,
-        num_workers: int=os.cpu_count()
-        ):
-    
-    """
-    Creates data loaders for the training and test datasets to be used to traing visiton transformers.
-
-    Args:
-        vit_model (str): The name of the ViT model to use. Default is "bitbase16".
-        train_dir (str): The path to the training dataset directory. Default is TRAIN_DIR.
-        test_dir (str): The path to the test dataset directory. Default is TEST_DIR.
-        batch_size (int): The batch size for the data loaders. Default is BATCH_SIZE.
-        aug (bool): Whether to apply data augmentation or not. Default is True.
-        display_imgs (bool): Whether to display sample images or not. Default is True.
-
-    Returns:
-        train_dataloader (torch.utils.data.DataLoader): The data loader for the training dataset.
-        test_dataloader (torch.utils.data.DataLoader): The data loader for the test dataset.
-        class_names (list): A list of class names.
-    """
-
-    IMG_SIZE = 224
-    IMG_SIZE_2 = 384
-
-    # Manual transforms for the training dataset
-    manual_transforms = v2.Compose([           
-        v2.RandomCrop((IMG_SIZE, IMG_SIZE)),    
-        v2.ToImage(),
-        v2.ToDtype(torch.float32, scale=True),    
-    ])
-
-    # ViT-Base/16 transforms
-    if vit_model == "vitbase16":
-
-        # Manual transforms for the training dataset
-        if aug:
-            manual_transforms_train_vitb = v2.Compose([    
-                v2.TrivialAugmentWide(),
-                v2.Resize((256, 256)),
-                v2.RandomCrop((IMG_SIZE, IMG_SIZE)),    
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225]) 
-            ])
-        else:
-            manual_transforms_train_vitb = v2.Compose([
-                v2.Resize((256, 256)),
-                v2.CenterCrop((IMG_SIZE, IMG_SIZE)),
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225])
-            ])
-
-        # Manual transforms for the test dataset
-        manual_transforms_test_vitb = v2.Compose([    
-            v2.Resize((256, 256)),
-            v2.CenterCrop((IMG_SIZE, IMG_SIZE)),    
-            v2.ToImage(),
-            v2.ToDtype(torch.float32, scale=True),
-            v2.Normalize(mean=[0.485, 0.456, 0.406],
-                        std=[0.229, 0.224, 0.225]) 
-        ])
-
-        # Create data loaders for ViT-Base
-        train_dataloader, test_dataloader, class_names = create_dataloaders(
-            train_dir=train_dir,
-            test_dir=test_dir,
-            train_transform=manual_transforms_train_vitb,
-            test_transform=manual_transforms_test_vitb,
-            batch_size=batch_size,
-            num_workers=num_workers
-            )
-    
-    if vit_model == "vitbase16_2":
-
-        # Manual transforms for the training dataset
-        if aug:
-            manual_transforms_train_vitb = v2.Compose([    
-                v2.TrivialAugmentWide(),
-                v2.Resize((IMG_SIZE_2, IMG_SIZE_2)),
-                v2.CenterCrop((IMG_SIZE_2, IMG_SIZE_2)),    
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225]) 
-            ])
-        else:
-            manual_transforms_train_vitb = v2.Compose([
-                v2.Resize((IMG_SIZE_2, IMG_SIZE_2)),
-                v2.CenterCrop((IMG_SIZE_2, IMG_SIZE_2)),
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225])
-            ])
-
-        # Manual transforms for the test dataset
-        manual_transforms_test_vitb = v2.Compose([    
-            v2.Resize((IMG_SIZE_2, IMG_SIZE_2)),
-            v2.CenterCrop((IMG_SIZE_2, IMG_SIZE_2)),    
-            v2.ToImage(),
-            v2.ToDtype(torch.float32, scale=True),
-            v2.Normalize(mean=[0.485, 0.456, 0.406],
-                        std=[0.229, 0.224, 0.225]) 
-        ])
-
-        # Create data loaders for ViT-Base
-        train_dataloader, test_dataloader, class_names = create_dataloaders(
-            train_dir=train_dir,
-            test_dir=test_dir,
-            train_transform=manual_transforms_train_vitb,
-            test_transform=manual_transforms_test_vitb,
-            batch_size=batch_size,
-            num_workers=num_workers
-            )
-
-    # ViT-Large/16 transforms
-    elif vit_model == "vitlarge16":
-
-        # Manual transforms for the training dataset
-        if aug:
-            manual_transforms_train_vitl = v2.Compose([    
-                v2.TrivialAugmentWide(),
-                v2.Resize((242, 242)),
-                v2.RandomCrop((IMG_SIZE, IMG_SIZE)),    
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225]) 
-            ])
-        else:
-            manual_transforms_train_vitl = v2.Compose([
-                v2.Resize((242, 242)),
-                v2.CenterCrop((IMG_SIZE, IMG_SIZE)),
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225])
-            ])
-
-        # Manual transforms for the test dataset
-        manual_transforms_test_vitl = v2.Compose([    
-            v2.Resize((242, 242)),
-            v2.CenterCrop((IMG_SIZE, IMG_SIZE)),    
-            v2.ToImage(),
-            v2.ToDtype(torch.float32, scale=True),
-            v2.Normalize(mean=[0.485, 0.456, 0.406],
-                        std=[0.229, 0.224, 0.225]) 
-        ])
-
-        # Create data loaders for ViT-Large/16
-        train_dataloader, test_dataloader, class_names = create_dataloaders(
-            train_dir=train_dir,
-            test_dir=test_dir,
-            train_transform=manual_transforms_train_vitl,
-            test_transform=manual_transforms_test_vitl,
-            batch_size=batch_size,
-            num_workers=num_workers
-        )
-
-    # ViT-Large/32 transforms
-    else:
-        # Manual transforms for the training dataset
-        if aug:
-            manual_transforms_train_vitl = v2.Compose([    
-                v2.TrivialAugmentWide(),
-                v2.Resize((256, 256)),
-                v2.RandomCrop((IMG_SIZE, IMG_SIZE)),    
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225]) 
-            ])
-        else:
-            manual_transforms_train_vitl = v2.Compose([
-                v2.Resize((256, 256)),
-                v2.CenterCrop((IMG_SIZE, IMG_SIZE)),
-                v2.ToImage(),
-                v2.ToDtype(torch.float32, scale=True),
-                v2.Normalize(mean=[0.485, 0.456, 0.406],
-                            std=[0.229, 0.224, 0.225])
-            ])
-
-        # Manual transforms for the test dataset
-        manual_transforms_test_vitl = v2.Compose([    
-            v2.Resize((256, 256)),
-            v2.CenterCrop((IMG_SIZE, IMG_SIZE)),    
-            v2.ToImage(),
-            v2.ToDtype(torch.float32, scale=True),
-            v2.Normalize(mean=[0.485, 0.456, 0.406],
-                        std=[0.229, 0.224, 0.225]) 
-        ])
-
-        # Create data loaders for ViT-Large/32
-        train_dataloader, test_dataloader, class_names = create_dataloaders(
-            train_dir=train_dir,
-            test_dir=test_dir,
-            train_transform=manual_transforms_train_vitl,
-            test_transform=manual_transforms_test_vitl,
-            batch_size=batch_size,
-            num_workers=num_workers
-        )
-
-    # Display images
-    if display_imgs:
-        train_data = datasets.ImageFolder(train_dir, transform=manual_transforms)
-        display_random_images(train_data,
-                        n=25,
-                        classes=class_names,
-                        rows=5,
-                        cols=5,
-                        display_shape=False,
-                        seed=None)
-    
-    return train_dataloader, test_dataloader, class_names
-
+from torch.nn.init import trunc_normal_
+#, xavier_normal_, zeros_, orthogonal_, kaiming_normal_
 
 # Create Pytorch's default ViT models
 def create_vit(