Uses the OpenCLIP transformer encoder for text
| 146 | |
| 147 | |
| 148 | class FrozenOpenCLIPEmbedder(nn.Module): |
| 149 | """ |
| 150 | Uses the OpenCLIP transformer encoder for text |
| 151 | """ |
| 152 | LAYERS = ["last", "penultimate"] |
| 153 | def __init__( |
| 154 | self, |
| 155 | arch="ViT-H-14", |
| 156 | version="laion2b_s32b_b79k", |
| 157 | max_length=77, |
| 158 | freeze=True, |
| 159 | layer="penultimate" |
| 160 | ): |
| 161 | super().__init__() |
| 162 | assert layer in self.LAYERS |
| 163 | model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version) |
| 164 | del model.visual |
| 165 | self.model = model |
| 166 | |
| 167 | self.max_length = max_length |
| 168 | if freeze: |
| 169 | self.freeze() |
| 170 | self.layer = layer |
| 171 | if self.layer == "last": |
| 172 | self.layer_idx = 0 |
| 173 | elif self.layer == "penultimate": |
| 174 | self.layer_idx = 1 |
| 175 | else: |
| 176 | raise NotImplementedError() |
| 177 | |
| 178 | self.uncond = None |
| 179 | |
| 180 | def freeze(self): |
| 181 | self.model = self.model.eval() |
| 182 | for param in self.parameters(): |
| 183 | param.requires_grad = False |
| 184 | |
| 185 | def forward(self, text): |
| 186 | dev = next(self.parameters()).device |
| 187 | tokens = open_clip.tokenize(text) |
| 188 | z = self.encode_with_transformer(tokens.to(dev)) |
| 189 | return z |
| 190 | |
| 191 | def encode_with_transformer(self, text): |
| 192 | x = self.model.token_embedding(text) # [batch_size, n_ctx, d_model] |
| 193 | x = x + self.model.positional_embedding |
| 194 | x = x.permute(1, 0, 2) # NLD -> LND |
| 195 | x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask) |
| 196 | x = x.permute(1, 0, 2) # LND -> NLD |
| 197 | x = self.model.ln_final(x) |
| 198 | return x |
| 199 | |
| 200 | def text_transformer_forward(self, x: torch.Tensor, attn_mask=None): |
| 201 | for i, r in enumerate(self.model.transformer.resblocks): |
| 202 | if i == len(self.model.transformer.resblocks) - self.layer_idx: |
| 203 | break |
| 204 | if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(): |
| 205 | x = checkpoint(r, x, attn_mask) |
nothing calls this directly
no outgoing calls
no test coverage detected