Merge pull request #53 from jmisilo/develop

Develop

readme.md

@@ -10,17 +10,17 @@ The Model uses prefixes as in the [ClipCap](https://arxiv.org/abs/2111.09734) pa
 
 The Model was trained with a frozen CLIP, a fully trained Mapping Module (6x Transformer Encoder Layers) and with partially frozen GPT-2 (the first and last 14 layers were trained).
 
-The training process was carried out using the [Kaggle](https://www.kaggle.com/) P100 GPU. Training time is about 2 x 11h (106 epochs) with a linearly changing learning rate (from 0 to 0.0001908) and batch size 64. Originally, the Model was supposed to be trained longer - which results in a non-standard LR. *I also tried a longer training session (150 epochs), but overtraining was noticeable.*
+The training process was carried out using the [Kaggle](https://www.kaggle.com/) P100 GPU. Training time - about 3 x 11h (150 epochs) with a linear learning rate warmup (max LR `3e-3`) and batch size 64. 
 
-### Example results
-
-![Example1](./examples/23012796.jpg)
+#### Loss and Learning Rate during training
 
-![Example2](./examples/3787801.jpg)
+![LOSSxLR](./examples/loss_lr.jpg)
 
-![Example3](./examples/7757242158.jpg)
+### Example results
 
-As I said, the goal was to test the Model's ability to recognize the situation. In the next phase of the experiments, I will try to improve the Model process and parameters to achieve better captions with the same dataset.
+![Example1](./examples/23012796.jpg)
+![Example2](./examples/36979.jpg)
+![Example3](./examples/89407459.jpg)
 
 ### Usage
 
@@ -36,7 +36,10 @@ Create environment and install requirements:
 
 ```bash
 python -m venv venv
+# for windows
 .\venv\Scripts\activate
+# for linux/mac
+source venv/bin/activate
 
 pip install -r requirements.txt
 ```

src/evaluate.py

@@ -44,6 +44,14 @@
     help='Path to the results folder'
 )
 
+parser.add_argument(
+    '-T', 
+    '--temperature',
+    type=float,
+    default=1.0,
+    help='Temperature for sampling'
+)
+
 args = parser.parse_args()
 
 ckp_path = os.path.join(config.weights_dir, args.checkpoint_name)
@@ -93,4 +101,4 @@
     if not os.path.exists(save_path):
         os.mkdir(save_path)
 
-    evaluate_dataset(model, test_dataset, args.img_path, save_path)
+    evaluate_dataset(model, test_dataset, args.img_path, save_path, args.temperature)

src/model/loops.py

@@ -130,7 +130,7 @@ def test_step(model, dataset, img_path, num_examples=4):
 
     return Image.open(buf)
 
-def evaluate_dataset(model, dataset, img_path, save_path):
+def evaluate_dataset(model, dataset, img_path, save_path, temperature=1.0):
     '''
         Evaluate model on dataset.
     
@@ -147,7 +147,7 @@ def evaluate_dataset(model, dataset, img_path, save_path):
         img = Image.open(os.path.join(img_path, img_name))
 
         with torch.no_grad():
-            caption, _ = model(img)
+            caption, _ = model(img, temperature)
 
         plt.imshow(img)
         plt.title(caption)

src/model/model.py

@@ -146,7 +146,7 @@ def freeze_layers(self):
         for p in [*list(self.ie.parameters()), *list(self.td.parameters())[14:-14]]: # freeze everything, except 1st and last transformer layer in Decoder
             p.requires_grad = False
 
-    def forward(self, img):
+    def forward(self, img, temperature=1.0):
         '''
             Caption generation for a single image.
 
@@ -157,7 +157,10 @@ def forward(self, img):
                 caption: generated caption [str]
                 tokens: generated tokens [torch.Tensor]
         '''
-        # only one image at a time
+
+        if temperature <= 0.0:
+            temperature = 1.0
+            print('Temperature must be positive. Setting it to 1.0')
 
         with torch.no_grad():
             img_embedded = self.ie(img)
@@ -188,6 +191,8 @@ def forward(self, img):
                 emb += pos_emb
                 pred = self.td(emb)
 
+                pred = torch.softmax(pred / temperature, dim=-1)
+
                 _, pred = torch.max(pred, dim=1)
 
                 last_token = pred[-1].item()

src/predict.py

@@ -43,6 +43,14 @@
     help='Path to the results folder'
 )
 
+parser.add_argument(
+    '-T', 
+    '--temperature',
+    type=float,
+    default=1.0,
+    help='Temperature for sampling'
+)
+
 args = parser.parse_args()
 
 # set seed
@@ -87,7 +95,7 @@
     model.eval()
 
     with torch.no_grad():
-        caption, _ = model(img)
+        caption, _ = model(img, args.temperature)
 
     plt.imshow(img)
     plt.title(caption)

src/training.py

@@ -87,7 +87,11 @@
     scaler = torch.cuda.amp.GradScaler()
 
     ckp_path = os.path.join(config.weights_dir, args.checkpoint_name)
-    start_epoch = load_ckp(ckp_path, model, optimizer, scheduler, scaler, device) if os.path.isfile(ckp_path) else 0
+    start_epoch, total_train_loss, total_valid_loss = (
+        load_ckp(ckp_path, model, optimizer, scheduler, scaler, device) 
+        if os.path.isfile(ckp_path) else 
+        (0, [], [])
+    )
 
     # build train model process with experiment tracking from wandb
     wandb.init(project='clipXgpt2 captioner', config=config.__dict__)
@@ -107,19 +111,22 @@
             'examples': wandb.Image(test_results)
         })
 
+        total_train_loss.append(train_loss)
+        total_valid_loss.append(valid_loss)
+
         if not os.path.exists(config.weights_dir):
             os.makedirs(config.weights_dir)
 
         if (epoch + 1) % 10 == 0: 
             torch.save(
                 {
                     'epoch': epoch,
-                    'train_loss': train_loss,
-                    'valid_loss': valid_loss,
-                    'model1_state_dict': model.state_dict(),
+                    'model_state_dict': model.state_dict(),
                     'optimizer_state_dict': optimizer.state_dict(),
                     'scheduler_state_dict': scheduler.state_dict(),
                     'scaler_state_dict': scaler.state_dict(),
+                    'tloss': total_train_loss,
+                    'vloss': total_valid_loss
                 }, 
                 os.path.join(config.weights_dir, f'epoch_{epoch}.pt')
             )

src/utils/config.py

@@ -15,8 +15,8 @@ class Config:
     num_workers: int = 2
     train_size: int = 0.84
     val_size: int = 0.13
-    epochs: int = 200
-    lr: int = 6e-3
+    epochs: int = 150
+    lr: int = 3e-3
     k: float = 0.33
     batch_size_exp: int = 6
     ep_len: int = 4

src/utils/load_ckp.py

@@ -12,7 +12,7 @@ def load_ckp(checkpoint_fpath, model, optimizer=None, scheduler=None, scaler=Non
 
     checkpoint = torch.load(checkpoint_fpath, map_location=device)
 
-    model.load_state_dict(checkpoint['model1_state_dict'])
+    model.load_state_dict(checkpoint['model_state_dict'])
     if optimizer is not None:
         optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
 
@@ -22,11 +22,11 @@ def load_ckp(checkpoint_fpath, model, optimizer=None, scheduler=None, scaler=Non
     if scaler is not None:
         scaler.load_state_dict(checkpoint['scaler_state_dict'])
 
-    return checkpoint['epoch']
+    return checkpoint['epoch'], checkpoint['tloss'], checkpoint['vloss']
 
 def download_weights(checkpoint_fpath):
     '''
         Downloads weights from Google Drive.
     '''
 
-    gdown.download('https://drive.google.com/uc?id=1lEufQVOETFEIhPdFDYaez31uroq_5Lby', checkpoint_fpath, quiet=False)
+    gdown.download('https://drive.google.com/uc?id=10ieSMMJzE9EeiPIF3CMzeT4timiQTjHV', checkpoint_fpath, quiet=False)