Class CLIP

prototype/model/slip.py:112–206

class CLIP(nn.Module):
    def __init__(self,image_encode, text_encode, use_allgather):
        super().__init__()
        self.use_allgather = use_allgather
        self.visual =image_encode
        self.text_encoder = text_encode
        self.logit_scale = nn.Parameter(torch.ones([1]))
        # self.logit_scale = nn.Parameter(torch.ones([]))
        nn.init.constant_(self.logit_scale, np.log(1/0.07))
        #nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)

    def text_parameters(self):
        param = [self.logit_scale]
        if self.text_encoder.text_encode_type == 'Transformer':
            param.append(self.text_encoder.positional_embedding)
        elif self.text_encoder.text_encode_type == 'Bert':
            # print('Bert', self.text_encoder.text_transformer.cls.predictions, flush=True)
            # param.extend([self.text_encoder.text_transformer.cls.predictions.decoder.weight,
            #               self.text_encoder.text_transformer.cls.predictions.bias])
            param.extend([self.text_encoder.text_transformer.cls.predictions.bias])
        return param

    def text_modules(self):
        if self.text_encoder.text_encode_type == 'Transformer':
            return [self.text_encoder.transformer, self.text_encoder.text_projection, self.text_encoder.token_embedding, self.text_encoder.ln_final]
        elif self.text_encoder.text_encode_type == 'Bert':
            # print('Bert', self.text_encoder.text_transformer, flush=True)
            return [self.text_encoder.text_transformer.bert, self.text_encoder.text_projection,
                    self.text_encoder.text_transformer.cls.predictions.transform]
                    # self.text_encoder.text_transformer.cls.predictions.decoder,  # decoder: bias
        else:
            import ipdb
            ipdb.set_trace()
            return [self.text_encoder.text_transformer, self.text_encoder.text_projection]

    def visual_parameters(self):
        return []

    def visual_modules(self):
        return [self.visual]

    @property
    def dtype(self):
        try:
            return self.visual.conv1.weight.dtype
        except:
            try:
                return self.visual.head.weight.dtype
            except:
                try:
                    return self.visual.stem[0].weight.dtype
                except:
                    return self.text_encoder.text_projection.weight.dtype

    def encode_image(self, image):
        return self.visual(image.type(self.dtype))

    def sample_captions(self, texts):