adding the best model to hugging face

app.py

@@ -51,7 +51,7 @@ def get_model():
     
     # recuperate the model
     model = ViTForImageClassification.from_pretrained(
-        'data/checkpoints/model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500',
+        'data/checkpoints/model_lhGqMDq/checkpoint-440',
         num_labels = len(characs['ids']),
         id2label = {name: key for key, name in characs['ids'].items()},
         label2id = characs['ids']
@@ -84,6 +84,12 @@ if file is not None:
     
     left.markdown("""---""")
     
+    # add a side for the scaler and the head number
+    scale = st.sidebar.slider("Attention scale", min_value=30, max_value =200)
+
+    head = int(st.sidebar.selectbox("Attention head", options=list(range(1, 13))))
+
+    
     if left.button("SUBMIT"):
         
         # Let us convert the image format to 'RGB'
@@ -116,7 +122,7 @@ if file is not None:
             attention = outputs.attentions[-1][0]
             
             # Let us recuperate the attention image
-            attention_image = get_attention(image, attention, size = (224, 224), patch_size = (14, 14))
+            attention_image = get_attention(image, attention, size = (224, 224), patch_size = (14, 14), scale = scale, head = head)
 
             # Let us transform the attention image to a opencv image
             attention_image = cv2.cvtColor(attention_image.astype('float32'), cv2.COLOR_RGB2BGR)

data/checkpoints/model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500/trainer_state.json

@@ -1,64 +0,0 @@
-{
-  "best_metric": 0.6927365064620972,
-  "best_model_checkpoint": "data/checkpoints/model_2yW4AcqNIb6zLKNIb6zLK\\checkpoint-1500",
-  "epoch": 1.710376282782212,
-  "global_step": 1500,
-  "is_hyper_param_search": false,
-  "is_local_process_zero": true,
-  "is_world_process_zero": true,
-  "log_history": [
-    {
-      "epoch": 0.57,
-      "learning_rate": 0.00012064414686134504,
-      "loss": 0.6945,
-      "step": 500
-    },
-    {
-      "epoch": 0.57,
-      "eval_accuracy": 0.5081081081081081,
-      "eval_f1": 0.38095238095238093,
-      "eval_loss": 0.6931825280189514,
-      "eval_runtime": 6.1462,
-      "eval_samples_per_second": 30.1,
-      "eval_steps_per_second": 3.905,
-      "step": 500
-    },
-    {
-      "epoch": 1.14,
-      "learning_rate": 0.00010514828346782865,
-      "loss": 0.6937,
-      "step": 1000
-    },
-    {
-      "epoch": 1.14,
-      "eval_accuracy": 0.4702702702702703,
-      "eval_f1": 0.0,
-      "eval_loss": 0.6942673325538635,
-      "eval_runtime": 11.0225,
-      "eval_samples_per_second": 16.784,
-      "eval_steps_per_second": 2.177,
-      "step": 1000
-    },
-    {
-      "epoch": 1.71,
-      "learning_rate": 8.962136623985633e-05,
-      "loss": 0.6936,
-      "step": 1500
-    },
-    {
-      "epoch": 1.71,
-      "eval_accuracy": 0.5297297297297298,
-      "eval_f1": 0.6925795053003534,
-      "eval_loss": 0.6927365064620972,
-      "eval_runtime": 6.7463,
-      "eval_samples_per_second": 27.423,
-      "eval_steps_per_second": 3.558,
-      "step": 1500
-    }
-  ],
-  "max_steps": 4385,
-  "num_train_epochs": 5,
-  "total_flos": 2.323984768541614e+17,
-  "trial_name": null,
-  "trial_params": null
-}

data/checkpoints/{model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500 → model_lhGqMDq/checkpoint-440}/optimizer.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2db8a17d7b7b6fb8c5b9ca808dfbd68977ee3eb53721a36287e69b2b0ca9a600
+oid sha256:34ae20bd389738b6680714ac81db29b7818f40d430f3ee5f99989e50498d9dbd
 size 686518917

data/checkpoints/{model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500 → model_lhGqMDq/checkpoint-440}/pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5f62ac6440dfa6ddd0737da7788da4569a54b1d0d8cc027ce7f623e6db55b05b
+oid sha256:2e40727744f63eb801348e4e0730accb61a5fdf014bb8086765db19a592fe248
 size 343268717

data/checkpoints/{model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500 → model_lhGqMDq/checkpoint-440}/rng_state.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f6f8978ea188121d9fecb3f115bfdfa2c5cad2a9b0c1d0a104dddd3e07af89f6
+oid sha256:26c8e955c34ff14bfa8f687fd165d49decdd3e034d94027e5187c7e7e7496c1a
 size 14575

data/checkpoints/{model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500 → model_lhGqMDq/checkpoint-440}/scaler.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7cefe4c84c26da8b0778281baf520dd6d4b41a18ea28fd317c86c1f2b76d30fb
+oid sha256:76c820f9778523d807e0a80d015f245f8f1a75ff5c8ee1aa94e258d65b1066f5
 size 557

data/checkpoints/{model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500 → model_lhGqMDq/checkpoint-440}/scheduler.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ea1d89fe668b9776e02fd071e17549ee7882e574b4efb0629a6d930572aab462
+oid sha256:d56c15a0722c57b78cdab7fc0e6a1424469ef928d58365fbf9e1bfe0832b43e7
 size 627

data/checkpoints/model_lhGqMDq/checkpoint-440/trainer_state.json

@@ -0,0 +1,80 @@
+{
+  "best_metric": 0.32448598742485046,
+  "best_model_checkpoint": "data/checkpoints/model_lhGqMDq\\checkpoint-440",
+  "epoch": 4.0,
+  "global_step": 440,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 1.0,
+      "learning_rate": 6.923196230748668e-05,
+      "loss": 0.6551,
+      "step": 110
+    },
+    {
+      "epoch": 1.0,
+      "eval_accuracy": 0.6702702702702703,
+      "eval_f1": 0.6772486772486773,
+      "eval_loss": 0.6143904328346252,
+      "eval_runtime": 7.9911,
+      "eval_samples_per_second": 23.151,
+      "eval_steps_per_second": 3.003,
+      "step": 110
+    },
+    {
+      "epoch": 2.0,
+      "learning_rate": 4.615464153832446e-05,
+      "loss": 0.5106,
+      "step": 220
+    },
+    {
+      "epoch": 2.0,
+      "eval_accuracy": 0.7675675675675676,
+      "eval_f1": 0.7860696517412936,
+      "eval_loss": 0.4895593523979187,
+      "eval_runtime": 7.6354,
+      "eval_samples_per_second": 24.229,
+      "eval_steps_per_second": 3.143,
+      "step": 220
+    },
+    {
+      "epoch": 3.0,
+      "learning_rate": 2.307732076916223e-05,
+      "loss": 0.4299,
+      "step": 330
+    },
+    {
+      "epoch": 3.0,
+      "eval_accuracy": 0.8108108108108109,
+      "eval_f1": 0.8044692737430168,
+      "eval_loss": 0.435648649930954,
+      "eval_runtime": 8.1072,
+      "eval_samples_per_second": 22.819,
+      "eval_steps_per_second": 2.96,
+      "step": 330
+    },
+    {
+      "epoch": 4.0,
+      "learning_rate": 0.0,
+      "loss": 0.2903,
+      "step": 440
+    },
+    {
+      "epoch": 4.0,
+      "eval_accuracy": 0.8594594594594595,
+      "eval_f1": 0.8673469387755102,
+      "eval_loss": 0.32448598742485046,
+      "eval_runtime": 7.1509,
+      "eval_samples_per_second": 25.871,
+      "eval_steps_per_second": 3.356,
+      "step": 440
+    }
+  ],
+  "max_steps": 440,
+  "num_train_epochs": 4,
+  "total_flos": 5.433738311775191e+17,
+  "trial_name": null,
+  "trial_params": null
+}

data/checkpoints/{model_2yW4AcqNIb6zLKNIb6zLK/checkpoint-1500 → model_lhGqMDq/checkpoint-440}/training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f595feb7a87c355b9cc6060e25a6fe727190de55b96c20541d829a9c3b0fe7af
-size 3643
+oid sha256:74817fa301a4e63d886a80850533c40044c6316a596cd85b00f752249a979780
+size 3579

fake_face_detection/data/lion_cheetah_collator.py

@@ -0,0 +1,33 @@
+
+import torch
+import numpy as np
+
+def lion_cheetah_collator(batch): 
+    """The data collator for training vision transformer models on the lion cheetah dataset
+
+    Args:
+        batch (list): A dictionary containing the pixel values and the labels
+
+    Returns:
+        dict: The final dictionary
+    """
+    
+    new_batch = {
+        'pixel_values': [],
+        'labels': []
+    }
+    
+    for x in batch:
+        
+        pixel_values = torch.from_numpy(x['pixel_values'][0]) if isinstance(x['pixel_values'][0], np.ndarray) \
+            else x['pixel_values'][0]
+        
+        new_batch['pixel_values'].append(pixel_values)
+        
+        new_batch['labels'].append(torch.tensor(x['labels']))
+    
+    new_batch['pixel_values'] = torch.stack(new_batch['pixel_values'])
+    
+    new_batch['labels'] = torch.stack(new_batch['labels'])
+    
+    return new_batch

fake_face_detection/data/lion_cheetah_dataset.py

@@ -0,0 +1,102 @@
+
+from fake_face_detection.utils.compute_weights import compute_weights
+from torch.utils.data import Dataset
+from PIL import Image
+from glob import glob
+import numpy as np
+import torch
+import os
+
+class LionCheetahDataset(Dataset):
+
+    def __init__(self, lion_path: str, cheetah_path: str, id_map: dict, transformer, **transformer_kwargs):
+        
+        # let us recuperate the transformer
+        self.transformer = transformer
+        
+        # let us recuperate the transformer kwargs
+        self.transformer_kwargs = transformer_kwargs
+        
+        # let us load the images 
+        lion_images = glob(os.path.join(lion_path, "*"))
+        
+        cheetah_images = glob(os.path.join(cheetah_path, "*"))
+        
+        # recuperate rgb images
+        self.lion_images = []
+        
+        self.cheetah_images = []
+        
+        for lion in lion_images:
+            
+            try:
+                
+                with Image.open(lion) as img:
+                    
+                    # let us add a transformation on the images
+                    if self.transformer:
+                        
+                        image = self.transformer(img, **self.transformer_kwargs)
+                
+                self.lion_images.append(lion)
+            
+            except Exception as e:
+                
+                pass
+            
+        for cheetah in cheetah_images:
+            
+            try:
+                
+                with Image.open(cheetah) as img:
+                    
+                    # let us add a transformation on the images
+                    if self.transformer:
+                        
+                        image = self.transformer(img, **self.transformer_kwargs)
+                
+                self.cheetah_images.append(cheetah)
+            
+            except Exception as e:
+                
+                pass
+        
+        self.images = self.lion_images + self.cheetah_images
+        
+        # let us recuperate the labels
+        self.lion_labels = [int(id_map['lion'])] * len(self.lion_images)
+        
+        self.cheetah_labels = [int(id_map['cheetah'])] * len(self.cheetah_images)
+        
+        self.labels = self.lion_labels + self.cheetah_labels
+        
+        # let us recuperate the weights
+        self.weights = torch.from_numpy(compute_weights(self.labels))
+        
+        # let us recuperate the length
+        self.length = len(self.labels)
+        
+    def __getitem__(self, index):
+        
+        # let us recuperate an image
+        image = self.images[index]
+        
+        with Image.open(image) as img:
+            
+            # let us recuperate a label
+            label = self.labels[index]
+            
+            # let us add a transformation on the images
+            if self.transformer:
+                
+                image = self.transformer(img, **self.transformer_kwargs)
+                
+        # let us add the label inside the obtained dictionary
+        image['labels'] = label
+        
+        return image    
+        
+    def __len__(self):
+        
+        return self.length
+        

fake_face_detection/metrics/make_predictions.py

@@ -16,7 +16,7 @@ import numpy as np
 import torch
 import os
 
-def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, size: tuple, patch_size: tuple):
+def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, size: tuple, patch_size: tuple, scale: int = 50, head: int = 1):
     
     # recuperate the image as a numpy array
     if isinstance(image, str):
@@ -33,7 +33,7 @@ def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, siz
     attention = attention[:, -1, 1:]
 
     # calculate the mean attention
-    attention = attention.mean(axis = 0)
+    attention = attention[head - 1]
 
     # let us reshape transform the image to a numpy array
 
@@ -48,9 +48,9 @@ def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, siz
     attention = attention.reshape(size[0], size[1], 1)
     
     # recuperate the result
-    attention_image = img / 255 * attention.numpy()
+    attention_image = img / 255 * attention.numpy() * scale
     
-    return attention_image
+    return np.clip(attention_image, 0, 1)
 
 
 def make_predictions(test_dataset: FakeFaceDetectionDataset,
@@ -60,7 +60,28 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
                      batch_size: int = 3,
                      size: tuple = (224, 224), 
                      patch_size: tuple = (14, 14),
-                     figsize: tuple = (24, 24)):
+                     figsize: tuple = (24, 24),
+                     attention_scale: int = 50,
+                     show: bool = True, 
+                     head: int = 1):
+    """Make predictions with a vision transformer model
+
+    Args:
+        test_dataset (FakeFaceDetectionDataset): The test dataset
+        model (_type_): The model
+        log_dir (str, optional): The log directory. Defaults to "fake_face_logs".
+        tag (str, optional): The tag. Defaults to "Attentions".
+        batch_size (int, optional): The batch size. Defaults to 3.
+        size (tuple, optional): The size of the attention image. Defaults to (224, 224).
+        patch_size (tuple, optional): The path size. Defaults to (14, 14).
+        figsize (tuple, optional): The figure size. Defaults to (24, 24).
+        attention_scale (int, optional): The attention scale. Defaults to 50.
+        show (bool, optional): A boolean value indicating if we want to recuperate the figure. Defaults to True.
+        head (int, optional): The head number. Defaults to 1.
+
+    Returns:
+        Union[Tuple[pd.DataFrame, dict], Tuple[pd.DataFame, dict, figure]]: The return prediction and the metrics
+    """
     
     with torch.no_grad():
         
@@ -86,22 +107,22 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
         for data in test_dataloader:
             
             # recuperate the pixel values
-            pixel_values = data['pixel_values'][0].cuda()
+            pixel_values = data['pixel_values'][0]
             
             # recuperate the labels
-            labels_ = data['labels'].cuda()
+            labels_ = data['labels']
             
             # # recuperate the outputs
             outputs = model(pixel_values, labels = labels_, output_attentions = True)
             
             # recuperate the predictions
-            predictions['predictions'].append(torch.softmax(outputs.logits.detach().cpu(), axis = -1).numpy())
+            predictions['predictions'].append(torch.softmax(outputs.logits.detach(), axis = -1).numpy())
             
             # recuperate the attentions of the last encoder layer
-            predictions['attentions'].append(outputs.attentions[-1].detach().cpu())
+            predictions['attentions'].append(outputs.attentions[-1].detach())
             
             # add the loss
-            loss += outputs.loss.detach().cpu().item()
+            loss += outputs.loss.detach().item()
         
         predictions['predictions'] = np.concatenate(predictions['predictions'], axis = 0)
         
@@ -140,8 +161,11 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
         del predictions['predictions']
         del predictions['attentions']
         
-        # let us recuperate the metrics and the predictions
-        return pd.DataFrame(predictions), metrics
+        # show the figure if necessary
+        if show: return pd.DataFrame(predictions), metrics, fig
+        else:
+            # let us recuperate the metrics and the predictions
+            return pd.DataFrame(predictions), metrics
         
         
         

fake_face_detection/optimization/fake_face_bayesian_optimization.py

@@ -53,15 +53,15 @@ class SimpleBayesianOptimizationForFakeReal:
             
                 pickler = pickle.Unpickler(f)
                 
-                self.checkpoint = pickler.load()
+                checkpoint = pickler.load()
                 
-                self.data = self.checkpoint['data']
+                self.data = checkpoint['data']
                 
-                self.scores = self.checkpoint['scores']
+                self.scores = checkpoint['scores']
                 
-                self.model = self.checkpoint['model']
+                self.model = checkpoint['model']
                 
-                self.current_trial = self.checkpoint['trial']
+                self.current_trial = checkpoint['trial']
                 
                 print(f"Checkpoint loaded at trial {self.current_trial}")
         
@@ -113,17 +113,8 @@ class SimpleBayesianOptimizationForFakeReal:
             new_sample = generate_sample(self.data, self.model, self.search_spaces, n_tests, maximize = self.maximize)
             config = {key: new_sample[i] for i, key in enumerate(self.search_spaces)}
             
-            # initialize the random kwargs with a random values
-            random_kwargs = {key: value + ''.join(random.choice(letters) for i in range(7)) for key, value in self.random_kwargs.items()}
-            
-            # add random kwargs to the kwargs
-            self.kwargs.update(random_kwargs)
-            
-            # add config to kwargs
-            self.kwargs['config'] = config
-            
-            # calculate the first score
-            new_score = self.objective(**self.kwargs)
+            # recuperate a new score
+            new_score = self.get_score(config)
             
             # let us add the new sample, target and score to their lists
             self.data.append(new_sample)
@@ -148,6 +139,25 @@ class SimpleBayesianOptimizationForFakeReal:
                 }
                 
                 pickler.dump(checkpoint)
+    
+    def get_score(self, config: dict):
+        
+        # add random seed (since we have always the same problem of randomizing the seed)
+        random.seed(None)
+        
+        # initialize the random kwargs with a random values
+        random_kwargs = {key: value + ''.join(random.choice(letters) for i in range(7)) for key, value in self.random_kwargs.items()}
+        print(random_kwargs)
+        # add random kwargs to the kwargs
+        self.kwargs.update(random_kwargs)
+        
+        # add config to kwargs
+        self.kwargs['config'] = config
+        
+        # calculate the first score
+        new_score = self.objective(**self.kwargs)
+        
+        return new_score
         
     def get_results(self):
         """Recuperate the generated samples and the scores

fake_face_detection/trainers/lion_cheetah_search_train.py

@@ -0,0 +1,80 @@
+
+from fake_face_detection.metrics.compute_metrics import compute_metrics
+from fake_face_detection.data.lion_cheetah_collator import lion_cheetah_collator
+from transformers import Trainer, TrainingArguments, set_seed
+from torch.utils.tensorboard import SummaryWriter
+from torch import nn
+from typing import *
+import numpy as np
+import json
+import os
+
+def train(epochs: int, output_dir: str, config: dict, model: nn.Module, trainer, get_datasets: Callable, log_dir: str = "fake_face_logs", metric = 'accuracy', seed: int = 0):
+    
+    print("------------------------- Beginning of training")
+    
+    set_seed(seed)
+    
+    # initialize the model
+    model = model()
+    
+    # reformat the config integer type
+    for key, value in config.items():
+        
+        if isinstance(value, np.int32): config[key] = int(value)
+    
+    pretty = json.dumps(config, indent = 4)
+    
+    print(f"Current Config: \n {pretty}")
+    
+    print(f"Checkpoints in {output_dir}")
+    
+    # recuperate the dataset
+    train_dataset, test_dataset = get_datasets(config['h_flip_p'], config['v_flip_p'], config['gray_scale_p'], config['rotation'])
+    
+    # initialize the arguments of the training
+    training_args = TrainingArguments(output_dir,
+                                      per_device_train_batch_size=config['batch_size'],
+                                      evaluation_strategy='steps',
+                                      save_strategy='steps',
+                                      logging_strategy='steps',
+                                      num_train_epochs=epochs,
+                                      fp16=True,
+                                      save_total_limit=2,
+                                      remove_unused_columns=True,
+                                      push_to_hub=False,
+                                      logging_dir=os.path.join(log_dir, os.path.basename(output_dir)),
+                                      load_best_model_at_end=True,
+                                      learning_rate=config['lr'],
+                                      weight_decay=config['weight_decay']
+                                      )
+    
+    # train the model
+    trainer_ = trainer(
+        model = model,
+        args = training_args,
+        data_collator = lion_cheetah_collator,
+        compute_metrics = compute_metrics,
+        train_dataset = train_dataset,
+        eval_dataset = test_dataset
+    )
+    
+    # train the model
+    trainer_.train()
+    
+    # evaluate the model and recuperate metrics
+    metrics = trainer_.evaluate(test_dataset)
+    
+    # add metrics and config to the hyperparameter panel of tensorboard
+    with SummaryWriter(os.path.join(log_dir, 'lchparams')) as logger:
+        
+        logger.add_hparams(
+            config, metrics
+        )
+    
+    print(metrics)
+    
+    print("------------------------- End of training")
+    # recuperate the metric to evaluate
+    return metrics[f'eval_{metric}']
+        

fake_face_detection/trainers/search_train.py

@@ -1,7 +1,7 @@
 
 from fake_face_detection.metrics.compute_metrics import compute_metrics
 from fake_face_detection.data.collator import fake_face_collator
-from transformers import Trainer, TrainingArguments
+from transformers import Trainer, TrainingArguments, set_seed
 from torch.utils.tensorboard import SummaryWriter
 from torch import nn
 from typing import *
@@ -9,10 +9,15 @@ import numpy as np
 import json
 import os
 
-def train(epochs: int, output_dir: str, config: dict, model: nn.Module, trainer, get_datasets: Callable, log_dir: str = "fake_face_logs", metric = 'accuracy'):
+def train(epochs: int, output_dir: str, config: dict, model: nn.Module, trainer, get_datasets: Callable, log_dir: str = "fake_face_logs", metric = 'accuracy', seed: int = 0):
     
     print("------------------------- Beginning of training")
     
+    set_seed(seed)
+    
+    # initialize the model
+    model = model()
+    
     # reformat the config integer type
     for key, value in config.items():
         
@@ -22,24 +27,24 @@ def train(epochs: int, output_dir: str, config: dict, model: nn.Module, trainer,
     
     print(f"Current Config: \n {pretty}")
     
+    print(f"Checkpoints in {output_dir}")
+    
     # recuperate the dataset
     train_dataset, test_dataset = get_datasets(config['h_flip_p'], config['v_flip_p'], config['gray_scale_p'], config['rotation'])
     
     # initialize the arguments of the training
     training_args = TrainingArguments(output_dir,
                                       per_device_train_batch_size=config['batch_size'],
-                                      evaluation_strategy='steps',
-                                      save_strategy='steps',
-                                      logging_strategy='steps',
+                                      evaluation_strategy='epoch',
+                                      save_strategy='epoch',
+                                      logging_strategy='epoch',
                                       num_train_epochs=epochs,
                                       fp16=True,
                                       save_total_limit=2,
-                                      remove_unused_columns=True,
                                       push_to_hub=False,
                                       logging_dir=os.path.join(log_dir, os.path.basename(output_dir)),
                                       load_best_model_at_end=True,
-                                      learning_rate=config['lr'],
-                                      weight_decay=config['weight_decay']
+                                      learning_rate=config['lr']
                                       )
     
     # train the model

fake_face_detection/utils/visualize_images.py

@@ -15,7 +15,19 @@ def visualize_images(images_dict: Dict[str, Iterable[Union[JpegImageFile, torch.
                      log_directory: str = "fake_face_logs",
                      n_images: int = 40,
                      figsize = (15, 15), 
-                     seed: Union[int, None] = None):
+                     seed: Union[int, None] = None,
+                     show: bool = True
+                     ):
+    """Visualize some images from a dictionary
+
+    Args:
+        images_dict (Dict[str, Iterable[Union[JpegImageFile, torch.Tensor, np.ndarray]]]): The dictionary of the images with key indicating the tag
+        log_directory (str, optional): The tensorboard log directory. Defaults to "fake_face_logs".
+        n_images (int, optional): The number of images. Defaults to 40.
+        figsize (tuple, optional): The figure size. Defaults to (15, 15).
+        seed (Union[int, None], optional): The seed. Defaults to None.
+        show (bool): Indicate if we want to show the figure. Defaults to True.
+    """
     
     assert len(images_dict) > 0
     
@@ -77,3 +89,5 @@ def visualize_images(images_dict: Dict[str, Iterable[Union[JpegImageFile, torch.
         
         writer.add_figure(tag = tag, figure = fig)
     
+    if show: return fig
+