Class PLOTPP

plot-pp/trainers/plotpp.py:266–396

It is based on PLOT.

@TRAINER_REGISTRY.register()
class PLOTPP(TrainerX):
    """
    It is based on PLOT.
    """
    
    def check_cfg(self, cfg):
        assert cfg.TRAINER.PLOTPP.PREC in ["fp16", "fp32", "amp"]
    
    def build_model(self):
        cfg = self.cfg
        classnames = self.dm.dataset.classnames
        name_tp_update = cfg.TRAINER.PLOTPP.MODEL_UPD
        
        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)
        
        if cfg.TRAINER.PLOTPP.PREC == "fp32" or cfg.TRAINER.PLOTPP.PREC == "amp":
            # CLIP's default precision is fp16
            clip_model.float()
        
        print("Building custom CLIP")
        self.model = CustomCLIP(cfg, classnames, clip_model)
        
        print("Turning off gradients in both the image and the text encoder")
        for name, param in self.model.named_parameters():
                
                if "prompt_learner" not in name:
                    param.requires_grad_(False)
                else:
                    if name_tp_update == "vision" and name_tp_update not in name:
                        param.requires_grad_(False)
     # Double check
        enabled = set()
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                enabled.add(name)
                
        print(f"Parameters to be updated: {enabled}")
        
        # if cfg.MODEL.INIT_WEIGHTS:
        #     load_pretrained_weights(self.model.prompt_learner, cfg.MODEL.INIT_WEIGHTS)
        
        if cfg.TRAINER.PLOTPP.PRETRAIN_DIR:
            load_pretrained_weights(self.model, cfg.TRAINER.PLOTPP.PRETRAIN_DIR)
        
        device_count = torch.cuda.device_count()
        if cfg.DATASET.NAME == 'ImageNet':
            self.device = torch.device("cuda:0")
            device1 = torch.device("cuda")
            self.model.to(self.device)
            self.model.text_encoder.to(device1)
            self.model.text_encoder=nn.DataParallel(self.model.text_encoder)
        elif device_count > 1:
            print(f"Multiple GPUs detected (n_gpus={device_count}), use all of them!")
            self.device = torch.device("cuda")
            self.model.to(self.device)
            self.model = nn.DataParallel(self.model)
        else: