| 174 | |
| 175 | |
| 176 | class CLIP(nn.Module): |
| 177 | output_dict: torch.jit.Final[bool] |
| 178 | |
| 179 | def __init__( |
| 180 | self, |
| 181 | embed_dim: int, |
| 182 | vision_cfg: CLIPVisionCfg, |
| 183 | text_cfg: CLIPTextCfg, |
| 184 | quick_gelu: bool = False, |
| 185 | cast_dtype: Optional[torch.dtype] = None, |
| 186 | output_dict: bool = False, |
| 187 | ): |
| 188 | super().__init__() |
| 189 | self.output_dict = output_dict |
| 190 | self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype) |
| 191 | |
| 192 | text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype) |
| 193 | self.transformer = text.transformer |
| 194 | self.vocab_size = text.vocab_size |
| 195 | self.token_embedding = text.token_embedding |
| 196 | self.positional_embedding = text.positional_embedding |
| 197 | self.ln_final = text.ln_final |
| 198 | self.text_projection = text.text_projection |
| 199 | self.register_buffer('attn_mask', text.attn_mask, persistent=False) |
| 200 | |
| 201 | self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) |
| 202 | |
| 203 | def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False): |
| 204 | # lock image tower as per LiT - https://arxiv.org/abs/2111.07991 |
| 205 | self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats) |
| 206 | |
| 207 | def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True): |
| 208 | locked_layers = [] |
| 209 | locked_layers.append(self.token_embedding) |
| 210 | self.positional_embedding.requires_grad = False |
| 211 | if unlocked_layers > 0: |
| 212 | locked_layers.append(self.transformer.resblocks[:-unlocked_layers]) |
| 213 | else: |
| 214 | locked_layers.append(self.transformer) |
| 215 | locked_layers.append(self.ln_final) |
| 216 | self.text_projection.requires_grad = False |
| 217 | |
| 218 | # freeze layers |
| 219 | for module in locked_layers: |
| 220 | for n, p in module.named_parameters(): |
| 221 | p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False |
| 222 | |
| 223 | @torch.jit.ignore |
| 224 | def set_grad_checkpointing(self, enable=True): |
| 225 | self.visual.set_grad_checkpointing(enable) |
| 226 | self.transformer.grad_checkpointing = enable |
| 227 | |
| 228 | def encode_image(self, image, normalize: bool = False): |
| 229 | features = self.visual(image) |
| 230 | return F.normalize(features, dim=-1) if normalize else features |
| 231 | |
| 232 | def encode_text(self, text, normalize: bool = False): |
| 233 | cast_dtype = self.transformer.get_cast_dtype() |
no outgoing calls
no test coverage detected