Update clipGPT.py

clipGPT.py

@@ -7,6 +7,7 @@ import torch
 import torch.nn as nn
 from torch.utils.data import Dataset, DataLoader
 from torch.nn import functional as F
+from typing import Tuple
 
 import pandas as pd
 import numpy as np
@@ -16,19 +17,30 @@ import nltk
 nltk.download('punkt')
 
 
+class Adapter(nn.Module):
+    def forward(self, x):
+        return self.model(x)
+
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(Adapter, self).__init__()
+        layers = []
+        for i in range(len(sizes) -1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+        self.model = nn.Sequential(*layers)
+
+
+
 class ClipGPT2Model(nn.Module):
     def __init__(self, img_feature_length, img_feature_size = 512):
         super(ClipGPT2Model, self).__init__()
-        torch.cuda.empty_cache()
-        gc.collect()
         self.img_feature_length = img_feature_length
-
         self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
         self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
         self.clip_project = Adapter((img_feature_size,
                                        (self.gpt_embedding_size * img_feature_length) // 2,
                                        self.gpt_embedding_size * img_feature_length))
-        torch.cuda.empty_cache()
     def get_dummy_token(self,
                         batch_size: int,
                         device: torch.device) -> torch.Tensor:
@@ -39,8 +51,6 @@ class ClipGPT2Model(nn.Module):
                 feature: torch.Tensor,
                 mask = None,
                 labels = None):
-        torch.cuda.empty_cache()
-        gc.collect()
 
         embedding_text = self.gpt.transformer.wte(tokens)
         feature_projections = self.clip_project(feature).view(-1, self.img_feature_length, self.gpt_embedding_size)