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