| 51 | |
| 52 | |
| 53 | class CLIP(nn.Module): |
| 54 | def __init__(self,image_encode, text_encode, use_allgather): |
| 55 | super().__init__() |
| 56 | self.use_allgather = use_allgather |
| 57 | self.visual =image_encode |
| 58 | self.encode_text = text_encode |
| 59 | self.logit_scale = nn.Parameter(torch.ones([1])) |
| 60 | # self.logit_scale = nn.Parameter(torch.ones([])) |
| 61 | nn.init.constant_(self.logit_scale, np.log(1/0.07)) |
| 62 | #nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5) |
| 63 | |
| 64 | def text_parameters(self): |
| 65 | param = [self.logit_scale] |
| 66 | if self.encode_text.text_encode_type == 'Transformer': |
| 67 | param.append(self.encode_text.positional_embedding) |
| 68 | elif self.encode_text.text_encode_type == 'Bert': |
| 69 | # print('Bert', self.encode_text.text_transformer.cls.predictions, flush=True) |
| 70 | # param.extend([self.encode_text.text_transformer.cls.predictions.decoder.weight, |
| 71 | # self.encode_text.text_transformer.cls.predictions.bias]) |
| 72 | param.extend([self.encode_text.text_transformer.cls.predictions.bias]) |
| 73 | return param |
| 74 | |
| 75 | def text_modules(self): |
| 76 | if self.encode_text.text_encode_type == 'Transformer': |
| 77 | return [self.encode_text.transformer, self.encode_text.text_projection, self.encode_text.token_embedding, self.encode_text.ln_final] |
| 78 | elif self.encode_text.text_encode_type == 'Bert': |
| 79 | # print('Bert', self.encode_text.text_transformer, flush=True) |
| 80 | return [self.encode_text.text_transformer.bert, self.encode_text.text_projection, |
| 81 | self.encode_text.text_transformer.cls.predictions.transform] |
| 82 | # self.encode_text.text_transformer.cls.predictions.decoder, # decoder: bias |
| 83 | else: |
| 84 | import ipdb |
| 85 | ipdb.set_trace() |
| 86 | return [self.encode_text.text_transformer, self.encode_text.text_projection] |
| 87 | |
| 88 | def visual_parameters(self): |
| 89 | return [] |
| 90 | |
| 91 | def visual_modules(self): |
| 92 | return [self.visual] |
| 93 | |
| 94 | @property |
| 95 | def dtype(self): |
| 96 | try: |
| 97 | return self.visual.conv1.weight.dtype |
| 98 | except: |
| 99 | try: |
| 100 | return self.visual.head.weight.dtype |
| 101 | except: |
| 102 | try: |
| 103 | return self.visual.stem[0].weight.dtype |
| 104 | except: |
| 105 | return self.encode_text.text_projection.weight.dtype |
| 106 | |
| 107 | def encode_image(self, image): |
| 108 | return self.visual(image.type(self.dtype)) |
| 109 | |
| 110 | def sample_captions(self, texts): |
no outgoing calls
no test coverage detected