| 247 | |
| 248 | |
| 249 | class CLIP(nn.Module): |
| 250 | def __init__(self, |
| 251 | embed_dim: int, |
| 252 | # vision |
| 253 | image_resolution: int, |
| 254 | vision_layers: Union[Tuple[int, int, int, int], int], |
| 255 | vision_width: int, |
| 256 | vision_patch_size: int, |
| 257 | # text |
| 258 | context_length: int, |
| 259 | vocab_size: int, |
| 260 | transformer_width: int, |
| 261 | transformer_heads: int, |
| 262 | transformer_layers: int |
| 263 | ): |
| 264 | super().__init__() |
| 265 | |
| 266 | self.context_length = context_length |
| 267 | |
| 268 | if isinstance(vision_layers, (tuple, list)): |
| 269 | vision_heads = vision_width * 32 // 64 |
| 270 | self.visual = ModifiedResNet( |
| 271 | layers=vision_layers, |
| 272 | output_dim=embed_dim, |
| 273 | heads=vision_heads, |
| 274 | input_resolution=image_resolution, |
| 275 | width=vision_width |
| 276 | ) |
| 277 | else: |
| 278 | vision_heads = vision_width // 64 |
| 279 | self.visual = VisualTransformer( |
| 280 | input_resolution=image_resolution, |
| 281 | patch_size=vision_patch_size, |
| 282 | width=vision_width, |
| 283 | layers=vision_layers, |
| 284 | heads=vision_heads, |
| 285 | output_dim=embed_dim |
| 286 | ) |
| 287 | |
| 288 | self.transformer = Transformer( |
| 289 | width=transformer_width, |
| 290 | layers=transformer_layers, |
| 291 | heads=transformer_heads, |
| 292 | attn_mask=self.build_attention_mask() |
| 293 | ) |
| 294 | |
| 295 | self.vocab_size = vocab_size |
| 296 | self.token_embedding = nn.Embedding(vocab_size, transformer_width) |
| 297 | self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width)) |
| 298 | self.ln_final = LayerNorm(transformer_width) |
| 299 | |
| 300 | self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim)) |
| 301 | self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) |
| 302 | |
| 303 | self.initialize_parameters() |
| 304 | |
| 305 | def initialize_parameters(self): |
| 306 | nn.init.normal_(self.token_embedding.weight, std=0.02) |